perf(amneziawg): return gVisor's pooled buffers on the embedded data path

Every packet crossing the embedded AmneziaWG interface allocated instead of
reusing gVisor's pools, in both directions. stackTun.Write injected each
decrypted packet and never called DecRef, so the packet buffer and its chunk
were never returned; stackTun.Read copied each view out and never released
it. gVisor's own link endpoints settle the ownership question -- loopback.go
and sharedmem.go both DecRef immediately after DeliverNetworkPacket, because
the injector owns the buffer.

AttachUDPHandler compounded it by cloning a packet buffer it then dropped on
the floor, on top of a Data().AsRange().ToSlice() that already returns an
owned copy, so the clone bought nothing and stranded a pooled buffer plus a
cloned view per datagram.

Measured with the benchmarks added here:

  stackTunWrite (upload)     794ns -> 107ns   4 -> 0 allocs
  stackTunRead  (download)   707ns -> 129ns   3 -> 0 allocs
  UDP datagram, end to end  2.69us -> 1.58us  8 -> 2 allocs

The remaining UDP allocation is the ToSlice copy itself. Through a real
handshaked tunnel -- both devices in one process over loopback, so
ChaCha20-Poly1305 and the UDP syscalls dominate -- it is worth -48% bytes/op
and -33% allocs/op, and about +4.8% throughput in each direction (n=18,
p<=0.01). On a small VPS, where the allocation pressure is not spread over
24 idle cores, the throughput share should be larger; that part is reasoning,
not something measured here.

The three regression tests assert allocations per packet rather than timing,
since the defect is the pool miss, not the nanoseconds. Thresholds leave room
for the extra allocation -race adds.

internal/amneziawgnet/bench_test.go

@@ -0,0 +1,244 @@
+package amneziawgnet
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/netip"
+	"testing"
+	"time"
+
+	awgconn "github.com/amnezia-vpn/amneziawg-go/v3/conn"
+	"github.com/amnezia-vpn/amneziawg-go/v3/device"
+	"github.com/amnezia-vpn/amneziawg-go/v3/tun/netstack"
+	"gvisor.dev/gvisor/pkg/buffer"
+	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
+	"gvisor.dev/gvisor/pkg/tcpip/link/channel"
+
+	"github.com/mhsanaei/3x-ui/v3/internal/amneziawg"
+	"github.com/mhsanaei/3x-ui/v3/internal/util/wireguard"
+)
+
+// BenchmarkStackTunWrite measures the upload path's per-packet cost: one
+// decrypted packet handed from amneziawg-go into the gVisor stack.
+func BenchmarkStackTunWrite(b *testing.B) {
+	tun := &stackTun{ep: channel.New(tunQueueDepth, 1420, ""), mtu: 1420}
+	defer tun.ep.Close()
+
+	packet := make([]byte, 1400)
+	packet[0] = 0x45
+	bufs := [][]byte{packet}
+
+	b.SetBytes(int64(len(packet)))
+	b.ReportAllocs()
+	b.ResetTimer()
+	for range b.N {
+		if _, err := tun.Write(bufs, 0); err != nil {
+			b.Fatalf("Write: %v", err)
+		}
+	}
+}
+
+// BenchmarkStackTunRead measures the download path's per-packet cost: one
+// packet drained out of the stack for amneziawg-go to encrypt.
+func BenchmarkStackTunRead(b *testing.B) {
+	tun := &stackTun{incomingPacket: make(chan *buffer.View, tunQueueDepth)}
+	packet := make([]byte, 1400)
+	buf := [][]byte{make([]byte, 2048)}
+	sizes := make([]int, 1)
+
+	b.SetBytes(int64(len(packet)))
+	b.ReportAllocs()
+	b.ResetTimer()
+	for range b.N {
+		tun.incomingPacket <- buffer.NewViewWithData(packet)
+		if _, err := tun.Read(buf, sizes, 0); err != nil {
+			b.Fatalf("Read: %v", err)
+		}
+	}
+}
+
+// BenchmarkUDPDatagramDelivery measures one datagram travelling the whole
+// inbound path: stack injection, routing, and the UDP transport handler.
+func BenchmarkUDPDatagramDelivery(b *testing.B) {
+	tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("10.78.0.1")}, 1420)
+	if err != nil {
+		b.Fatalf("createNetTUNWithStack: %v", err)
+	}
+	defer tun.Close()
+
+	src := netip.MustParseAddrPort("10.78.0.2:40000")
+	dst := netip.MustParseAddrPort("10.78.9.9:5353")
+	payload := make([]byte, 1024)
+	AttachUDPHandler(gstack, func(netip.AddrPort, netip.AddrPort, []byte) {})
+
+	bufs := [][]byte{udpDatagram(src, dst, payload)}
+	st := tun.(*stackTun)
+
+	b.SetBytes(int64(len(payload)))
+	b.ReportAllocs()
+	b.ResetTimer()
+	for range b.N {
+		if _, err := st.Write(bufs, 0); err != nil {
+			b.Fatalf("Write: %v", err)
+		}
+	}
+}
+
+// Two destinations the benchmark forwarder tells apart: one drains what the
+// client sends, the other streams at the client. Neither is routed anywhere.
+const (
+	benchDiscardPort = 9001
+	benchSourcePort  = 9002
+)
+
+// benchTunnel is a live AmneziaWG pair -- this package's server Device and a
+// stock amneziawg-go client -- talking real encrypted UDP over loopback.
+type benchTunnel struct {
+	clientNet *netstack.Net
+	closeFn   func()
+}
+
+// newBenchTunnel brings up both devices and blocks until the handshake has
+// actually completed, so no setup cost lands inside the measured loop.
+func newBenchTunnel(b *testing.B, listenPort int, serverAddr, clientAddr string) *benchTunnel {
+	b.Helper()
+	serverPriv, serverPub, err := wireguard.GenerateWireguardKeypair()
+	if err != nil {
+		b.Fatalf("server keypair: %v", err)
+	}
+	clientPriv, clientPub, err := wireguard.GenerateWireguardKeypair()
+	if err != nil {
+		b.Fatalf("client keypair: %v", err)
+	}
+
+	inst := amneziawg.Instance{
+		Id:            90,
+		InterfaceName: "awgbench",
+		ListenPort:    listenPort,
+		PrivateKey:    serverPriv,
+		PublicKey:     serverPub,
+		Address:       []string{serverAddr + "/24"},
+		MTU:           1420,
+		Obfuscation: amneziawg.Obfuscation31{
+			Jc: 4, Jmin: 40, Jmax: 70,
+			S1: 20, S2: 30, S3: 20, S4: 20,
+		},
+		Peers: []amneziawg.Peer{{
+			Email:      "[email protected]",
+			PublicKey:  clientPub,
+			AllowedIPs: []string{clientAddr + "/32"},
+		}},
+	}
+
+	dev, err := newUnconfiguredDevice(inst, DeviceOptions{})
+	if err != nil {
+		b.Fatalf("newUnconfiguredDevice: %v", err)
+	}
+	AttachTCPForwarder(dev.Stack, func(conn *gonet.TCPConn, dest netip.AddrPort) {
+		defer conn.Close()
+		switch dest.Port() {
+		case benchDiscardPort:
+			_, _ = io.Copy(io.Discard, conn)
+		case benchSourcePort:
+			chunk := make([]byte, 64<<10)
+			for {
+				if _, err := conn.Write(chunk); err != nil {
+					return
+				}
+			}
+		}
+	})
+	if err := dev.Configure(inst, DeviceOptions{}); err != nil {
+		b.Fatalf("Configure: %v", err)
+	}
+
+	clientTun, clientNet, err := netstack.CreateNetTUN(
+		[]netip.Addr{netip.MustParseAddr(clientAddr)},
+		[]netip.Addr{netip.MustParseAddr("1.1.1.1")}, 1420)
+	if err != nil {
+		b.Fatalf("client CreateNetTUN: %v", err)
+	}
+	clientDev := device.NewDevice(clientTun, awgconn.NewDefaultBind(), device.NewLogger(device.LogLevelSilent, ""))
+
+	clientPrivHex, err := wireguard.KeyToHex(clientPriv)
+	if err != nil {
+		b.Fatalf("client key to hex: %v", err)
+	}
+	serverPubHex, err := wireguard.KeyToHex(serverPub)
+	if err != nil {
+		b.Fatalf("server key to hex: %v", err)
+	}
+	conf := fmt.Sprintf(
+		"private_key=%s\njc=4\njmin=40\njmax=70\ns1=20\ns2=30\ns3=20\ns4=20\npublic_key=%s\nendpoint=127.0.0.1:%d\nallowed_ip=0.0.0.0/0\n",
+		clientPrivHex, serverPubHex, listenPort)
+	if err := clientDev.IpcSet(conf); err != nil {
+		b.Fatalf("client IpcSet: %v", err)
+	}
+	if err := clientDev.Up(); err != nil {
+		b.Fatalf("client Up: %v", err)
+	}
+
+	t := &benchTunnel{clientNet: clientNet, closeFn: func() {
+		clientDev.Close()
+		dev.Close()
+	}}
+	// Prove the handshake really completed before anything is timed.
+	probe := t.dial(b, netip.MustParseAddrPort(fmt.Sprintf("%s:%d", serverAddr, benchDiscardPort)))
+	probe.Close()
+	return t
+}
+
+// dial opens one tunnelled connection, retrying while the handshake settles.
+func (t *benchTunnel) dial(b *testing.B, dest netip.AddrPort) *gonet.TCPConn {
+	b.Helper()
+	deadline := time.Now().Add(15 * time.Second)
+	for {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		conn, err := t.clientNet.DialContextTCPAddrPort(ctx, dest)
+		cancel()
+		if err == nil {
+			return conn
+		}
+		if time.Now().After(deadline) {
+			b.Fatalf("dial %v through tunnel: %v", dest, err)
+		}
+	}
+}
+
+// BenchmarkTunnelThroughput is the end-to-end number: real bytes through a
+// real handshaked AmneziaWG tunnel, in both directions.
+func BenchmarkTunnelThroughput(b *testing.B) {
+	const chunkSize = 64 << 10
+	const serverAddr = "10.203.0.1"
+	tun := newBenchTunnel(b, 58714, serverAddr, "10.203.0.2")
+	defer tun.closeFn()
+
+	b.Run("upload", func(b *testing.B) {
+		conn := tun.dial(b, netip.MustParseAddrPort(fmt.Sprintf("%s:%d", serverAddr, benchDiscardPort)))
+		defer conn.Close()
+		chunk := make([]byte, chunkSize)
+		b.SetBytes(chunkSize)
+		b.ReportAllocs()
+		b.ResetTimer()
+		for range b.N {
+			if _, err := conn.Write(chunk); err != nil {
+				b.Fatalf("upload write: %v", err)
+			}
+		}
+	})
+
+	b.Run("download", func(b *testing.B) {
+		conn := tun.dial(b, netip.MustParseAddrPort(fmt.Sprintf("%s:%d", serverAddr, benchSourcePort)))
+		defer conn.Close()
+		chunk := make([]byte, chunkSize)
+		b.SetBytes(chunkSize)
+		b.ReportAllocs()
+		b.ResetTimer()
+		for range b.N {
+			if _, err := io.ReadFull(conn, chunk); err != nil {
+				b.Fatalf("download read: %v", err)
+			}
+		}
+	})
+}

internal/amneziawgnet/netstack.go

@@ -105,7 +105,8 @@ func (t *stackTun) Events() <-chan awgtun.Event { return t.events }
 func (t *stackTun) MTU() (int, error)           { return t.mtu, nil }
 func (t *stackTun) BatchSize() int              { return 1 }
 
-// Read drains incomingPacket into buf, supporting batched reads.
+// Read drains incomingPacket into buf, supporting batched reads. Each view is
+// released once copied out, so the download path reuses gVisor's pooled chunks.
 func (t *stackTun) Read(buf [][]byte, sizes []int, offset int) (int, error) {
 	var view *buffer.View
 	select {
@@ -114,6 +115,7 @@ func (t *stackTun) Read(buf [][]byte, sizes []int, offset int) (int, error) {
 	case view = <-t.incomingPacket:
 	}
 	n, err := view.Read(buf[0][offset:])
+	view.Release()
 	if err != nil {
 		return 0, err
 	}
@@ -123,6 +125,7 @@ func (t *stackTun) Read(buf [][]byte, sizes []int, offset int) (int, error) {
 		select {
 		case view = <-t.incomingPacket:
 			n, err := view.Read(buf[count][offset:])
+			view.Release()
 			if err != nil {
 				return count, nil
 			}
@@ -135,6 +138,8 @@ func (t *stackTun) Read(buf [][]byte, sizes []int, offset int) (int, error) {
 	return count, nil
 }
 
+// Write injects each packet into the stack. The injector owns the packet
+// buffer -- DecRef returns it and its chunk to gVisor's pools (see loopback.go).
 func (t *stackTun) Write(buf [][]byte, offset int) (int, error) {
 	for _, b := range buf {
 		packet := b[offset:]
@@ -148,8 +153,10 @@ func (t *stackTun) Write(buf [][]byte, offset int) (int, error) {
 		case 6:
 			t.ep.InjectInbound(header.IPv6ProtocolNumber, pkb)
 		default:
+			pkb.DecRef()
 			return 0, syscall.EAFNOSUPPORT
 		}
+		pkb.DecRef()
 	}
 	return len(buf), nil
 }
@@ -176,6 +183,7 @@ func (t *stackTun) WriteNotify() {
 	select {
 	case t.incomingPacket <- view:
 	case <-t.done:
+		view.Release()
 	}
 }
 

internal/amneziawgnet/netstack_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"gvisor.dev/gvisor/pkg/buffer"
+	"gvisor.dev/gvisor/pkg/tcpip/link/channel"
 )
 
 // TestStackTunReadDrainsBufferedBatch is a regression test for a real
@@ -86,3 +87,44 @@ func TestStackTunReadStopsAtBufCapacity(t *testing.T) {
 		t.Errorf("leftover packet = %v, want [3]", got)
 	}
 }
+
+// TestStackTunWriteReturnsPacketBuffersToPool locks in gVisor's ownership rule
+// for the upload path: whoever calls InjectInbound must DecRef the packet.
+func TestStackTunWriteReturnsPacketBuffersToPool(t *testing.T) {
+	tun := &stackTun{ep: channel.New(tunQueueDepth, 1420, ""), mtu: 1420}
+	defer tun.ep.Close()
+
+	packet := make([]byte, 1400)
+	packet[0] = 0x45 // IPv4, version nibble is all Write inspects
+	bufs := [][]byte{packet}
+
+	allocs := testing.AllocsPerRun(1000, func() {
+		if _, err := tun.Write(bufs, 0); err != nil {
+			t.Fatalf("Write: %v", err)
+		}
+	})
+	// 0 once pooled, 4 when every packet buffer is stranded; -race adds ~1.
+	if allocs > 1 {
+		t.Fatalf("Write allocates %v times per packet, want <=1: injected packet buffers are not being returned to gVisor's pools", allocs)
+	}
+}
+
+// TestStackTunReadReturnsViewsToPool is the download-path counterpart: a view
+// that is copied out but never released strands its pooled chunk.
+func TestStackTunReadReturnsViewsToPool(t *testing.T) {
+	tun := &stackTun{incomingPacket: make(chan *buffer.View, tunQueueDepth)}
+	packet := make([]byte, 1400)
+	buf := [][]byte{make([]byte, 2048)}
+	sizes := make([]int, 1)
+
+	allocs := testing.AllocsPerRun(1000, func() {
+		tun.incomingPacket <- buffer.NewViewWithData(packet)
+		if _, err := tun.Read(buf, sizes, 0); err != nil {
+			t.Fatalf("Read: %v", err)
+		}
+	})
+	// 0 once the drained view goes back to viewPool, 3 when it does not.
+	if allocs > 1 {
+		t.Fatalf("Read allocates %v times per packet, want <=1: drained views are not being released", allocs)
+	}
+}

internal/amneziawgnet/udp.go

@@ -32,7 +32,9 @@ func AttachUDPHandler(gstack *stack.Stack, handler UDPHandler) {
 	enablePromiscuousRouting(gstack)
 
 	gstack.SetTransportProtocolHandler(udp.ProtocolNumber, func(id stack.TransportEndpointID, pkt *stack.PacketBuffer) bool {
-		data := pkt.Clone().Data().AsRange().ToSlice()
+		// ToSlice already returns an owned copy, so cloning pkt here would only
+		// strand a pooled packet buffer and its chunks on every datagram.
+		data := pkt.Data().AsRange().ToSlice()
 		src := netip.AddrPortFrom(addrFromTcpip(id.RemoteAddress), id.RemotePort)
 		dst := netip.AddrPortFrom(addrFromTcpip(id.LocalAddress), id.LocalPort)
 		handler(src, dst, data)

internal/amneziawgnet/udp_test.go

@@ -9,6 +9,8 @@ import (
 	awgconn "github.com/amnezia-vpn/amneziawg-go/v3/conn"
 	"github.com/amnezia-vpn/amneziawg-go/v3/device"
 	"github.com/amnezia-vpn/amneziawg-go/v3/tun/netstack"
+	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/header"
 
 	"github.com/mhsanaei/3x-ui/v3/internal/amneziawg"
 	"github.com/mhsanaei/3x-ui/v3/internal/util/wireguard"
@@ -154,3 +156,64 @@ func TestNewDeviceUDPHandlerAndReply(t *testing.T) {
 		return
 	}
 }
+
+// udpDatagram builds a complete IPv4/UDP packet, the shape stackTun.Write
+// expects from amneziawg-go after decryption.
+func udpDatagram(src, dst netip.AddrPort, payload []byte) []byte {
+	total := header.IPv4MinimumSize + header.UDPMinimumSize + len(payload)
+	p := make([]byte, total)
+	ip := header.IPv4(p)
+	ip.Encode(&header.IPv4Fields{
+		TotalLength: uint16(total),
+		TTL:         64,
+		Protocol:    uint8(header.UDPProtocolNumber),
+		SrcAddr:     tcpip.AddrFromSlice(src.Addr().AsSlice()),
+		DstAddr:     tcpip.AddrFromSlice(dst.Addr().AsSlice()),
+	})
+	ip.SetChecksum(^ip.CalculateChecksum())
+	u := header.UDP(p[header.IPv4MinimumSize:])
+	u.Encode(&header.UDPFields{
+		SrcPort: src.Port(),
+		DstPort: dst.Port(),
+		Length:  uint16(header.UDPMinimumSize + len(payload)),
+	})
+	copy(p[header.IPv4MinimumSize+header.UDPMinimumSize:], payload)
+	return p
+}
+
+// TestAttachUDPHandlerDoesNotStrandPacketBuffers drives a real datagram all the
+// way through the stack: Range.ToSlice already copies, so cloning pkt only leaks.
+func TestAttachUDPHandlerDoesNotStrandPacketBuffers(t *testing.T) {
+	tun, gstack, err := createNetTUNWithStack([]netip.Addr{netip.MustParseAddr("10.77.0.1")}, 1420)
+	if err != nil {
+		t.Fatalf("createNetTUNWithStack: %v", err)
+	}
+	defer tun.Close()
+
+	src := netip.MustParseAddrPort("10.77.0.2:40000")
+	dst := netip.MustParseAddrPort("10.77.9.9:5353")
+	payload := make([]byte, 512)
+	var delivered int
+	AttachUDPHandler(gstack, func(gotSrc, gotDst netip.AddrPort, got []byte) {
+		if gotSrc != src || gotDst != dst || len(got) != len(payload) {
+			t.Errorf("handler got (%v -> %v, %d bytes), want (%v -> %v, %d bytes)", gotSrc, gotDst, len(got), src, dst, len(payload))
+		}
+		delivered++
+	})
+
+	bufs := [][]byte{udpDatagram(src, dst, payload)}
+	st := tun.(*stackTun)
+	allocs := testing.AllocsPerRun(500, func() {
+		if _, err := st.Write(bufs, 0); err != nil {
+			t.Fatalf("Write: %v", err)
+		}
+	})
+	if delivered == 0 {
+		t.Fatal("handler never ran: the datagram never reached the UDP transport handler")
+	}
+	// 2 once nothing is stranded (1 is ToSlice itself), 8 with the leaked
+	// clone plus the un-released packet buffer; -race adds ~1.
+	if allocs > 4 {
+		t.Fatalf("UDP delivery allocates %v times per datagram, want <=4: pooled packet buffers are being stranded", allocs)
+	}
+}