Reduce VPN control packet latency

agents/rap-node-agent/internal/vpnruntime/fabric_transport.go

@@ -1794,17 +1794,22 @@ func (t *FabricPacketTransport) ReceiveGatewayPacketBatch(ctx context.Context, t
 type FabricPacketInbox struct {
 	capacity int
 	mu       sync.Mutex
-	queues   map[string]chan mesh.VPNPacketBatchPayload
+	queues   map[string]*fabricPacketInboxQueue
 	dropped  uint64
 }
 
+type fabricPacketInboxQueue struct {
+	normal   chan mesh.VPNPacketBatchPayload
+	priority chan mesh.VPNPacketBatchPayload
+}
+
 func NewFabricPacketInbox(capacity int) *FabricPacketInbox {
 	if capacity <= 0 {
 		capacity = 4096
 	}
 	return &FabricPacketInbox{
 		capacity: capacity,
-		queues:   map[string]chan mesh.VPNPacketBatchPayload{},
+		queues:   map[string]*fabricPacketInboxQueue{},
 	}
 }
 
@@ -1860,12 +1865,46 @@ func (i *FabricPacketInbox) Receive(ctx context.Context, vpnConnectionID, direct
 	defer timer.Stop()
 	queue := i.queue(vpnConnectionID, direction)
 	for {
+		select {
+		case payload := <-queue.priority:
+			packets := cleanPacketBatch(payload.Packets)
+			if len(packets) == 0 {
+				continue
+			}
+			return packets, nil
+		default:
+		}
+		if len(queue.normal) > 0 {
+			priorityTimer := time.NewTimer(2 * time.Millisecond)
+			select {
+			case <-ctx.Done():
+				priorityTimer.Stop()
+				return nil, ctx.Err()
+			case <-timer.C:
+				priorityTimer.Stop()
+				return nil, nil
+			case payload := <-queue.priority:
+				priorityTimer.Stop()
+				packets := cleanPacketBatch(payload.Packets)
+				if len(packets) == 0 {
+					continue
+				}
+				return packets, nil
+			case <-priorityTimer.C:
+			}
+		}
 		select {
 		case <-ctx.Done():
 			return nil, ctx.Err()
 		case <-timer.C:
 			return nil, nil
-		case payload := <-queue:
+		case payload := <-queue.priority:
+			packets := cleanPacketBatch(payload.Packets)
+			if len(packets) == 0 {
+				continue
+			}
+			return packets, nil
+		case payload := <-queue.normal:
 			packets := cleanPacketBatch(payload.Packets)
 			if len(packets) == 0 {
 				continue
@@ -1877,8 +1916,12 @@ func (i *FabricPacketInbox) Receive(ctx context.Context, vpnConnectionID, direct
 
 func (i *FabricPacketInbox) enqueue(payload mesh.VPNPacketBatchPayload) error {
 	queue := i.queue(payload.VPNConnectionID, payload.Direction)
+	target := queue.normal
+	if payload.Direction == FabricDirectionGatewayToClient && batchHasTCPControlPacket(payload.Packets) {
+		target = queue.priority
+	}
 	select {
-	case queue <- payload:
+	case target <- payload:
 	default:
 		i.mu.Lock()
 		i.dropped++
@@ -1887,21 +1930,41 @@ func (i *FabricPacketInbox) enqueue(payload mesh.VPNPacketBatchPayload) error {
 	return nil
 }
 
-func (i *FabricPacketInbox) queue(vpnConnectionID, direction string) chan mesh.VPNPacketBatchPayload {
+func (i *FabricPacketInbox) queue(vpnConnectionID, direction string) *fabricPacketInboxQueue {
 	key := vpnConnectionID + "\x00" + direction
 	i.mu.Lock()
 	defer i.mu.Unlock()
 	if i.queues == nil {
-		i.queues = map[string]chan mesh.VPNPacketBatchPayload{}
+		i.queues = map[string]*fabricPacketInboxQueue{}
 	}
 	queue, ok := i.queues[key]
 	if !ok {
-		queue = make(chan mesh.VPNPacketBatchPayload, i.capacity)
+		priorityCapacity := maxInt(1, i.capacity/4)
+		queue = &fabricPacketInboxQueue{
+			normal:   make(chan mesh.VPNPacketBatchPayload, i.capacity),
+			priority: make(chan mesh.VPNPacketBatchPayload, priorityCapacity),
+		}
 		i.queues[key] = queue
 	}
 	return queue
 }
 
+func batchHasTCPControlPacket(packets [][]byte) bool {
+	for _, packet := range packets {
+		if isTCPControlPacket(packet) {
+			return true
+		}
+	}
+	return false
+}
+
+func maxInt(a, b int) int {
+	if a > b {
+		return a
+	}
+	return b
+}
+
 func (i *FabricPacketInbox) Dropped() uint64 {
 	if i == nil {
 		return 0
@@ -1933,7 +1996,7 @@ func (i *FabricPacketInbox) Snapshot() FabricPacketInboxSnapshot {
 	snapshot.Dropped = i.dropped
 	snapshot.QueueCount = len(i.queues)
 	for key, queue := range i.queues {
-		snapshot.QueueDepths[strings.ReplaceAll(key, "\x00", ":")] = len(queue)
+		snapshot.QueueDepths[strings.ReplaceAll(key, "\x00", ":")] = len(queue.normal) + len(queue.priority)
 	}
 	return snapshot
 }