WireGuard inbound: Reduce memory held by idle UDP flows (#6856)

https://github.com/XTLS/Xray-core/pull/6856#issuecomment-5870639238

proxy/tun/stack_gvisor.go

@@ -101,7 +101,7 @@ func (t *stackGVisor) Start() error {
 	// Use custom UDP packet handler, instead of strict gVisor forwarder, for FullCone NAT support
 	udpForwarder := newUdpConnectionHandler(t.handler.HandleConnection, t.writeRawUDPPacket)
 	ipStack.SetTransportProtocolHandler(udp.ProtocolNumber, func(id stack.TransportEndpointID, pkt *stack.PacketBuffer) bool {
-		data := pkt.Clone().Data().AsRange().ToSlice()
+		data := pkt.Data().AsRange().ToSlice()
 		// if len(data) == 0 {
 		// 	return false
 		// }

proxy/wireguard/queue_bench_test.go

@@ -0,0 +1,80 @@
+package wireguard
+
+import (
+	"runtime"
+	"testing"
+)
+
+const benchBatch = 64
+
+// Raw cost of queueing and draining a small burst, as one flow's reader does.
+func BenchmarkQueueBurstChan(b *testing.B) {
+	ch := make(chan *packet, udpQueueLimit)
+	p := &packet{}
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		for j := 0; j < benchBatch; j++ {
+			ch <- p
+		}
+		for j := 0; j < benchBatch; j++ {
+			<-ch
+		}
+	}
+}
+
+func BenchmarkQueueBurstPacketQueue(b *testing.B) {
+	q := newPacketQueue(udpQueueLimit)
+	p := &packet{}
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		for j := 0; j < benchBatch; j++ {
+			q.push(p)
+		}
+		for j := 0; j < benchBatch; j++ {
+			q.pop()
+		}
+	}
+}
+
+// Producer and consumer on different goroutines; the producer yields when the
+// queue is full instead of spinning, like a blocking channel send would.
+func BenchmarkQueueStreamChan(b *testing.B) {
+	ch := make(chan *packet, udpQueueLimit)
+	p := &packet{}
+	done := make(chan struct{})
+	go func() {
+		for range ch {
+		}
+		close(done)
+	}()
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ch <- p
+	}
+	close(ch)
+	<-done
+}
+
+func BenchmarkQueueStreamPacketQueue(b *testing.B) {
+	q := newPacketQueue(udpQueueLimit)
+	p := &packet{}
+	done := make(chan struct{})
+	go func() {
+		for {
+			if _, ok := q.pop(); !ok {
+				break
+			}
+		}
+		close(done)
+	}()
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for !q.push(p) {
+			runtime.Gosched()
+		}
+	}
+	q.close()
+	<-done
+}

proxy/wireguard/tun.go

@@ -85,7 +85,7 @@ func CreateForwarder(gstack *stack.Stack, handler func(conn net.Conn, dest net.D
 	}
 
 	gstack.SetTransportProtocolHandler(udp.ProtocolNumber, func(id stack.TransportEndpointID, pkt *stack.PacketBuffer) bool {
-		data := pkt.Clone().Data().AsRange().ToSlice()
+		data := pkt.Data().AsRange().ToSlice()
 		// if len(data) == 0 {
 		// 	return false
 		// }
@@ -112,12 +112,7 @@ func (m *udpManager) feed(src net.Destination, dst net.Destination, data []byte)
 	m.mutex.RLock()
 	uc, ok := m.m[src.NetAddr()]
 	if ok {
-		select {
-		case uc.queue <- &packet{
-			p:    data,
-			dest: &dst,
-		}:
-		default:
+		if !uc.queue.push(&packet{p: data, dest: &dst}) {
 			errors.LogDebug(context.Background(), "drop udp with size ", len(data), " to ", dst.NetAddr(), " original ", uc.dst.NetAddr(), " > queue full")
 		}
 		m.mutex.RUnlock()
@@ -131,7 +126,7 @@ func (m *udpManager) feed(src net.Destination, dst net.Destination, data []byte)
 	uc, ok = m.m[src.NetAddr()]
 	if !ok {
 		uc = &udpConn{
-			queue: make(chan *packet, 1024),
+			queue: newPacketQueue(udpQueueLimit),
 			src:   src,
 			dst:   dst,
 		}
@@ -145,12 +140,7 @@ func (m *udpManager) feed(src net.Destination, dst net.Destination, data []byte)
 		go m.handler(uc, dst)
 	}
 
-	select {
-	case uc.queue <- &packet{
-		p:    data,
-		dest: &dst,
-	}:
-	default:
+	if !uc.queue.push(&packet{p: data, dest: &dst}) {
 		errors.LogDebug(context.Background(), "drop udp with size ", len(data), " to ", dst.NetAddr(), " original ", uc.dst.NetAddr(), " > queue full 2")
 	}
 }
@@ -158,7 +148,7 @@ func (m *udpManager) feed(src net.Destination, dst net.Destination, data []byte)
 func (m *udpManager) close(uc *udpConn) {
 	if !uc.closed {
 		uc.closed = true
-		close(uc.queue)
+		uc.queue.close()
 		delete(m.m, uc.src.NetAddr())
 	}
 }
@@ -232,7 +222,7 @@ type packet struct {
 }
 
 type udpConn struct {
-	queue     chan *packet
+	queue     *packetQueue
 	src       net.Destination
 	dst       net.Destination
 	writeFunc func(payload []byte, src net.Destination, dst net.Destination) error
@@ -242,7 +232,7 @@ type udpConn struct {
 
 func (c *udpConn) ReadMultiBuffer() (buf.MultiBuffer, error) {
 	for {
-		q, ok := <-c.queue
+		q, ok := c.queue.pop()
 		if !ok {
 			return nil, io.EOF
 		}
@@ -261,7 +251,7 @@ func (c *udpConn) ReadMultiBuffer() (buf.MultiBuffer, error) {
 }
 
 func (c *udpConn) Read(p []byte) (int, error) {
-	q, ok := <-c.queue
+	q, ok := c.queue.pop()
 	if !ok {
 		return 0, io.EOF
 	}
@@ -324,3 +314,68 @@ func (c *udpConn) SetReadDeadline(t time.Time) error {
 func (c *udpConn) SetWriteDeadline(t time.Time) error {
 	return nil
 }
+
+// udpQueueLimit bounds the packets waiting for one UDP flow; more are dropped.
+const udpQueueLimit = 1024
+
+// packetQueue holds the packets waiting for one UDP flow. Unlike a buffered
+// channel of the same bound it only allocates for packets actually queued, so
+// the many idle flows kept until the idle timeout cost next to nothing.
+type packetQueue struct {
+	mu     sync.Mutex
+	items  []*packet
+	limit  int
+	notify chan struct{}
+	closed bool
+}
+
+func newPacketQueue(limit int) *packetQueue {
+	return &packetQueue{limit: limit, notify: make(chan struct{}, 1)}
+}
+
+// push queues p and reports whether it was accepted.
+func (q *packetQueue) push(p *packet) bool {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	if q.closed || len(q.items) >= q.limit {
+		return false
+	}
+	q.items = append(q.items, p)
+	select {
+	case q.notify <- struct{}{}:
+	default:
+	}
+	return true
+}
+
+// pop blocks until a packet is queued or the queue is closed and drained.
+func (q *packetQueue) pop() (*packet, bool) {
+	for {
+		q.mu.Lock()
+		if len(q.items) > 0 {
+			p := q.items[0]
+			q.items[0] = nil
+			q.items = q.items[1:]
+			if len(q.items) == 0 {
+				q.items = nil
+			}
+			q.mu.Unlock()
+			return p, true
+		}
+		if q.closed {
+			q.mu.Unlock()
+			return nil, false
+		}
+		q.mu.Unlock()
+		<-q.notify
+	}
+}
+
+func (q *packetQueue) close() {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+	if !q.closed {
+		q.closed = true
+		close(q.notify)
+	}
+}

proxy/wireguard/tun_test.go

@@ -0,0 +1,91 @@
+package wireguard
+
+import (
+	"testing"
+	"time"
+
+	"github.com/xtls/xray-core/common/net"
+)
+
+// BenchmarkUDPManagerNewSession measures what one new UDP flow costs the
+// inbound while it stays open: QUIC and DNS open many short flows, and each
+// one lives until the connection idle timeout.
+func BenchmarkUDPManagerNewSession(b *testing.B) {
+	m := &udpManager{
+		handler: func(conn net.Conn, dest net.Destination) {},
+		m:       make(map[string]*udpConn),
+	}
+	dst := net.UDPDestination(net.ParseAddress("1.1.1.1"), 443)
+	payload := make([]byte, 1200)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		src := net.UDPDestination(net.IPAddress([]byte{10, byte(i >> 16), byte(i >> 8), byte(i)}), net.Port(1024+i%60000))
+		m.feed(src, dst, payload)
+	}
+}
+
+func TestPacketQueueOrderAndClose(t *testing.T) {
+	q := newPacketQueue(udpQueueLimit)
+	for i := 0; i < 3; i++ {
+		if !q.push(&packet{p: []byte{byte(i)}}) {
+			t.Fatalf("push %d rejected", i)
+		}
+	}
+	for i := 0; i < 3; i++ {
+		p, ok := q.pop()
+		if !ok || p.p[0] != byte(i) {
+			t.Fatalf("pop %d: got %v, %v", i, p, ok)
+		}
+	}
+	q.close()
+	if _, ok := q.pop(); ok {
+		t.Fatal("pop after close returned a packet")
+	}
+	if q.push(&packet{}) {
+		t.Fatal("push after close accepted")
+	}
+}
+
+func TestPacketQueueLimit(t *testing.T) {
+	q := newPacketQueue(udpQueueLimit)
+	for i := 0; i < udpQueueLimit; i++ {
+		if !q.push(&packet{}) {
+			t.Fatalf("push %d rejected below the limit", i)
+		}
+	}
+	if q.push(&packet{}) {
+		t.Fatal("push above the limit accepted")
+	}
+}
+
+func TestPacketQueueCloseUnblocksReader(t *testing.T) {
+	q := newPacketQueue(udpQueueLimit)
+	done := make(chan bool)
+	go func() {
+		_, ok := q.pop()
+		done <- ok
+	}()
+	q.close()
+	select {
+	case ok := <-done:
+		if ok {
+			t.Fatal("blocked pop returned a packet after close")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("close did not wake the reader")
+	}
+}
+
+func TestPacketQueueDropsDrainedStorage(t *testing.T) {
+	q := newPacketQueue(udpQueueLimit)
+	for i := 0; i < 100; i++ {
+		q.push(&packet{})
+	}
+	for i := 0; i < 100; i++ {
+		q.pop()
+	}
+	if q.items != nil {
+		t.Fatalf("drained queue still holds %d slots", cap(q.items))
+	}
+}