Merge pull request #62 from LaurenceJJones/split/upstream-dev-relay-session-index

perf(relay): index WireGuard sessions by receiver index

.github/workflows/test.yml

@@ -19,7 +19,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@4b73464bb391d4059bd26b0524d20df3927bd417 # v6.3.0
         with:
-          go-version: 1.25
+          go-version: 1.26
 
       - name: Build go
         run: go build

relay/relay.go

@@ -161,6 +161,9 @@ type UDPProxyServer struct {
 	// Session tracking for WireGuard peers
 	// Key format: "senderIndex:receiverIndex"
 	wgSessions sync.Map
+	// Session index for O(1) lookup by receiver index
+	// Key: receiverIndex (uint32), Value: *WireGuardSession
+	sessionsByReceiverIndex sync.Map
 	// Communication pattern tracking for rebuilding sessions
 	// Key format: "clientIP:clientPort-destIP:destPort"
 	commPatterns sync.Map
@@ -586,6 +589,8 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 		} else {
 			metrics.RecordSession(relayIfname, 1)
 		}
+		// Also index by sender index for O(1) lookup in transport data path
+		s.sessionsByReceiverIndex.Store(senderIndex, session)
 
 		// Forward the response to the original sender
 		for _, dest := range proxyMapping.Destinations {
@@ -615,21 +620,15 @@ func (s *UDPProxyServer) handleWireGuardPacket(packet []byte, remoteAddr *net.UD
 		// Data packet: forward only to the established session peer
 		// logger.Debug("Received transport data with receiver index %d from %s", receiverIndex, remoteAddr)
 
-		// Look up the session based on the receiver index
+		// Look up the session based on the receiver index - O(1) lookup instead of O(n) Range
 		var destAddr *net.UDPAddr
 
-		// First check for existing sessions to see if we know where to send this packet
-		s.wgSessions.Range(func(k, v interface{}) bool {
-			session := v.(*WireGuardSession)
-			// Check if session matches (read lock for check)
-			if session.GetSenderIndex() == receiverIndex {
-				// Found matching session - get dest addr and update last seen
-				destAddr = session.GetDestAddr()
-				session.UpdateLastSeen()
-				return false // stop iteration
-			}
-			return true // continue iteration
-		})
+		// Fast path: direct index lookup by receiver index
+		if sessionObj, ok := s.sessionsByReceiverIndex.Load(receiverIndex); ok {
+			session := sessionObj.(*WireGuardSession)
+			destAddr = session.GetDestAddr()
+			session.UpdateLastSeen()
+		}
 
 		if destAddr != nil {
 			// We found a specific peer to forward to
@@ -767,6 +766,8 @@ func (s *UDPProxyServer) handleResponses(conn *net.UDPConn, destAddr *net.UDPAdd
 				} else {
 					metrics.RecordSession(relayIfname, 1)
 				}
+				// Also index by sender index for O(1) lookup
+				s.sessionsByReceiverIndex.Store(senderIndex, session)
 				logger.Debug("Stored session mapping: %s -> %s", sessionKey, destAddr.String())
 			} else if ok && buffer[0] == WireGuardMessageTypeTransportData {
 				// Track communication pattern for session rebuilding (reverse direction)