From 08775afd6543ea9c979b9484f2fcd6ce132526d8 Mon Sep 17 00:00:00 2001 From: Denis <72341893+kroticw@users.noreply.github.com> Date: Mon, 5 Oct 2026 04:56:36 +0300 Subject: [PATCH] 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 | 2 +- proxy/wireguard/queue_bench_test.go | 80 +++++++++++++++++++++++++ proxy/wireguard/tun.go | 91 +++++++++++++++++++++++------ proxy/wireguard/tun_test.go | 91 +++++++++++++++++++++++++++++ 4 files changed, 245 insertions(+), 19 deletions(-) create mode 100644 proxy/wireguard/queue_bench_test.go create mode 100644 proxy/wireguard/tun_test.go diff --git a/proxy/tun/stack_gvisor.go b/proxy/tun/stack_gvisor.go index 8584616e9..de1735117 100644 --- a/proxy/tun/stack_gvisor.go +++ b/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 // } diff --git a/proxy/wireguard/queue_bench_test.go b/proxy/wireguard/queue_bench_test.go new file mode 100644 index 000000000..3380c84a8 --- /dev/null +++ b/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 +} diff --git a/proxy/wireguard/tun.go b/proxy/wireguard/tun.go index b4bcdc371..2d3a4a970 100644 --- a/proxy/wireguard/tun.go +++ b/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) + } +} diff --git a/proxy/wireguard/tun_test.go b/proxy/wireguard/tun_test.go new file mode 100644 index 000000000..88a45336f --- /dev/null +++ b/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)) + } +}