tunmux,wireguardfilter: preserve owned packet batches

Transfer pooled packet ownership through ChanTUN and Mux so forwarding avoids
redundant clones while retaining exact release, queue, and close semantics.
Preserve opportunistic batches across host, profile, local-service, filter, and
TSNet paths without latency timers or protocol changes.

Compact accepted filter buffers in place and cover ownership, ordering,
partial-admission, error, close, delegated-size, TSNet, and benchmark paths.

Author: maisem

profile/tsnetengine_test.go

@@ -21,6 +21,12 @@ func TestTSNetEngineAcceptsProvidedTunBeforeStart(t *testing.T) {
 	if engine == nil {
 		t.Fatal("engine is nil")
 	}
+	if got := tun.BatchSize(); got != 128 {
+		t.Fatalf("provided TUN batch size = %d, want 128", got)
+	}
+	if got := engine.cfg.Tun.BatchSize(); got != tun.BatchSize() {
+		t.Fatalf("engine TUN batch size = %d, want %d", got, tun.BatchSize())
+	}
 }
 
 func TestExitNodePrefsUsesBackendResolutionPath(t *testing.T) {

profile/wireguardengine_benchmark_test.go

@@ -14,13 +14,14 @@ import (
 	"github.com/maisem/tailmix/packetmap"
 	"github.com/maisem/tailmix/tunmux"
 	"github.com/maisem/tailmix/wireguardcfg"
+	"github.com/tailscale/wireguard-go/conn"
 	"tailscale.com/net/packet"
 )
 
 var wireGuardBenchmarkPacketSizes = [...]int{64, 512, 1280}
 
 const (
-	wireGuardBenchmarkWindow       = 256
+	wireGuardBenchmarkWindow       = 2 * conn.IdealBatchSize
 	wireGuardBenchmarkReceiveLimit = 5 * time.Second
 	wireGuardBenchmarkDrainQuiet   = 20 * time.Millisecond
 	wireGuardBenchmarkDrainLimit   = 2 * time.Second
@@ -112,12 +113,15 @@ func benchmarkWireGuardSteadyState(b *testing.B, fixture *wireGuardBenchmarkFixt
 	for completed := 0; completed < b.N; {
 		batchSize := min(wireGuardBenchmarkWindow, b.N-completed)
 		for i := range batchSize {
-			fixture.source.Outbound <- packets[i]
+			if err := fixture.source.InjectOutbound(packets[i]); err != nil {
+				b.Fatal(err)
+			}
 		}
 		resetWireGuardBenchmarkTimer(timer, wireGuardBenchmarkReceiveLimit)
 		for range batchSize {
 			select {
-			case <-fixture.destination.Inbound:
+			case packet := <-fixture.destination.Inbound():
+				packet.Release()
 			case err := <-fixture.muxError():
 				b.Fatalf("WireGuard benchmark mux stopped: %v", err)
 			case <-timer.C:
@@ -147,7 +151,8 @@ func benchmarkWireGuardOfferedLoad(b *testing.B, fixture *wireGuardBenchmarkFixt
 			select {
 			case <-stopDrain:
 				return
-			case <-fixture.destination.Inbound:
+			case packet := <-fixture.destination.Inbound():
+				packet.Release()
 				delivered.Add(1)
 			}
 		}
@@ -158,10 +163,9 @@ func benchmarkWireGuardOfferedLoad(b *testing.B, fixture *wireGuardBenchmarkFixt
 	b.ReportAllocs()
 	b.ResetTimer()
 	for i := range b.N {
-		select {
-		case fixture.source.Outbound <- packets[i%len(packets)]:
+		if fixture.source.TryInjectOutbound(packets[i%len(packets)]) {
 			accepted++
-		default:
+		} else {
 			backpressure++
 		}
 	}
@@ -369,19 +373,17 @@ func (f *wireGuardBenchmarkFixture) registerCleanup(tb testing.TB, hostTUNs []*t
 func (f *wireGuardBenchmarkFixture) establishSession(tb testing.TB, packetSize int) {
 	tb.Helper()
 	packet := wireGuardBenchmarkPacket(f.sourceIP, f.destinationIP, packetSize, 0)
-	select {
-	case f.source.Outbound <- packet:
-	case err := <-f.muxError():
-		tb.Fatalf("WireGuard benchmark mux stopped during warmup: %v", err)
-	case <-time.After(wireGuardBenchmarkReceiveLimit):
-		tb.Fatalf("timed out after %v injecting WireGuard warmup packet", wireGuardBenchmarkReceiveLimit)
+	if err := f.source.InjectOutbound(packet); err != nil {
+		tb.Fatalf("inject WireGuard warmup packet: %v", err)
 	}
 	select {
-	case got := <-f.destination.Inbound:
+	case got := <-f.destination.Inbound():
 		want := wireGuardBenchmarkPacket(f.expectedSourceIP, f.expectedDestination, packetSize, 0)
-		if !bytes.Equal(got, want) {
-			tb.Fatalf("WireGuard benchmark warmup packet changed in transit\n got: %x\nwant: %x", got, want)
+		if !bytes.Equal(got.Bytes(), want) {
+			got.Release()
+			tb.Fatalf("WireGuard benchmark warmup packet changed in transit")
 		}
+		got.Release()
 	case err := <-f.muxError():
 		tb.Fatalf("WireGuard benchmark mux stopped during warmup: %v", err)
 	case <-time.After(wireGuardBenchmarkReceiveLimit):
@@ -397,21 +399,20 @@ func (f *wireGuardBenchmarkFixture) verifyWindow(tb testing.TB, packetSize int)
 	if err != nil {
 		tb.Fatal(err)
 	}
-	deadline := time.NewTimer(wireGuardBenchmarkReceiveLimit)
-	defer deadline.Stop()
 	for _, packet := range packets {
-		select {
-		case f.source.Outbound <- packet:
-		case err := <-f.muxError():
-			tb.Fatalf("WireGuard benchmark mux stopped during correctness preflight: %v", err)
-		case <-deadline.C:
-			tb.Fatalf("timed out after %v injecting correctness packet", wireGuardBenchmarkReceiveLimit)
+		if err := f.source.InjectOutbound(packet); err != nil {
+			tb.Fatalf("inject WireGuard correctness packet: %v", err)
 		}
 	}
+
+	deadline := time.NewTimer(wireGuardBenchmarkReceiveLimit)
+	defer deadline.Stop()
 	for !validator.complete() {
 		select {
-		case packet := <-f.destination.Inbound:
-			if err := validator.observe(packet); err != nil {
+		case packet := <-f.destination.Inbound():
+			err := validator.observe(packet.Bytes())
+			packet.Release()
+			if err != nil {
 				tb.Fatal(err)
 			}
 		case err := <-f.muxError():
@@ -424,8 +425,9 @@ func (f *wireGuardBenchmarkFixture) verifyWindow(tb testing.TB, packetSize int)
 	quiet := time.NewTimer(wireGuardBenchmarkDrainQuiet)
 	defer quiet.Stop()
 	select {
-	case packet := <-f.destination.Inbound:
-		err := validator.observe(packet)
+	case packet := <-f.destination.Inbound():
+		err := validator.observe(packet.Bytes())
+		packet.Release()
 		if err == nil {
 			tb.Fatal("WireGuard benchmark accepted an excess correctness delivery")
 		}

profile/wireguardengine_test.go

@@ -450,10 +450,12 @@ func assertFilteredWrite(t *testing.T, e *WireGuardEngine, raw *tunmux.ChanTUN,
 		t.Fatalf("filtered Write() = (%d, %v)", n, err)
 	}
 	select {
-	case got := <-raw.Inbound:
+	case got := <-raw.Inbound():
 		if !want {
-			t.Fatalf("packet unexpectedly accepted: %v", got)
+			got.Release()
+			t.Fatal("packet unexpectedly accepted")
 		}
+		got.Release()
 	default:
 		if want {
 			t.Fatal("packet unexpectedly dropped")

tunmux/chantun.go

@@ -4,28 +4,42 @@ import (
 	"errors"
 	"io"
 	"os"
-	"slices"
+	"sync"
 
+	"github.com/tailscale/wireguard-go/conn"
+	"github.com/tailscale/wireguard-go/device"
 	"github.com/tailscale/wireguard-go/tun"
 )
 
+const (
+	chanTUNQueueSize = 1024
+	chanTUNMTU       = 1280
+)
+
 type ChanTUN struct {
-	Inbound  chan []byte
-	Outbound chan []byte
+	inbound  chan Packet
+	outbound chan Packet
 	events   chan tun.Event
 	closed   chan struct{}
+	pool     *packetPool
 	name     string
 	mtu      int
+
+	closeOnce sync.Once
+	sendMu    sync.Mutex
+	senders   sync.WaitGroup
+	isClosed  bool
 }
 
 func NewChanTUN(name string) *ChanTUN {
 	t := &ChanTUN{
-		Inbound:  make(chan []byte, 1024),
-		Outbound: make(chan []byte, 1024),
+		inbound:  make(chan Packet, chanTUNQueueSize),
+		outbound: make(chan Packet, chanTUNQueueSize),
 		events:   make(chan tun.Event, 1),
 		closed:   make(chan struct{}),
+		pool:     newPacketPool(device.MessageTransportHeaderSize, chanTUNMTU),
 		name:     name,
-		mtu:      1280,
+		mtu:      chanTUNMTU,
 	}
 	t.events <- tun.EventUp
 	return t
@@ -34,52 +48,153 @@ func NewChanTUN(name string) *ChanTUN {
 func (t *ChanTUN) File() *os.File { return nil }
 
 func (t *ChanTUN) Close() error {
-	select {
-	case <-t.closed:
-	default:
+	t.closeOnce.Do(func() {
+		t.sendMu.Lock()
+		t.isClosed = true
 		close(t.closed)
-	}
+		t.sendMu.Unlock()
+		t.senders.Wait()
+		close(t.inbound)
+		close(t.outbound)
+		for packet := range t.inbound {
+			packet.Release()
+		}
+		for packet := range t.outbound {
+			packet.Release()
+		}
+	})
 	return nil
 }
 
 func (t *ChanTUN) Read(bufs [][]byte, sizes []int, offset int) (int, error) {
-	select {
-	case <-t.closed:
+	batchSize := min(len(bufs), len(sizes), conn.IdealBatchSize)
+	if batchSize == 0 {
+		return 0, nil
+	}
+	packet, ok := <-t.outbound
+	if !ok {
 		return 0, io.EOF
-	case pkt := <-t.Outbound:
-		sizes[0] = copy(bufs[0][offset:], pkt)
-		return 1, nil
 	}
+	copyPacket := func(i int, packet Packet) {
+		sizes[i] = copy(bufs[i][offset:], packet.Bytes())
+		packet.Release()
+	}
+	copyPacket(0, packet)
+	n := 1
+	for n < batchSize {
+		select {
+		case packet, ok = <-t.outbound:
+			if !ok {
+				return n, nil
+			}
+			copyPacket(n, packet)
+			n++
+		default:
+			return n, nil
+		}
+	}
+	return n, nil
 }
 
 func (t *ChanTUN) Write(bufs [][]byte, offset int) (int, error) {
+	if !t.beginSend() {
+		return 0, errors.New("tun closed")
+	}
+	defer t.senders.Done()
 	for _, buf := range bufs {
-		pkt := buf[offset:]
-		if len(pkt) == 0 {
+		data := buf[offset:]
+		if len(data) == 0 {
 			continue
 		}
+		packet := t.pool.copy(data)
 		select {
 		case <-t.closed:
+			packet.Release()
 			return 0, errors.New("tun closed")
-		case t.Inbound <- slices.Clone(pkt):
+		case t.inbound <- packet:
 		default:
-			// A real TUN drops under sustained backpressure. Do the same
-			// instead of blocking a profile engine's receive path forever.
+			packet.Release()
 		}
 	}
 	return len(bufs), nil
 }
 
+// InjectOutbound copies borrowed packet bytes into owned storage and blocks
+// until ownership transfers to the outbound queue or the TUN closes.
 func (t *ChanTUN) InjectOutbound(pkt []byte) error {
+	packet := t.pool.copy(pkt)
+	if err := t.injectOutboundPacket(packet); err != nil {
+		packet.Release()
+		return err
+	}
+	return nil
+}
+
+// TryInjectOutbound copies borrowed packet bytes into owned storage and tries
+// to transfer ownership without blocking. It reports whether the transfer
+// succeeded.
+func (t *ChanTUN) TryInjectOutbound(pkt []byte) bool {
+	packet := t.pool.copy(pkt)
+	if !t.beginSend() {
+		packet.Release()
+		return false
+	}
+	defer t.senders.Done()
+	select {
+	case <-t.closed:
+		packet.Release()
+		return false
+	case t.outbound <- packet:
+		return true
+	default:
+		packet.Release()
+		return false
+	}
+}
+
+// Inbound returns packets written by the profile engine. Receiving transfers
+// ownership to the caller, which must release or forward each Packet.
+func (t *ChanTUN) Inbound() <-chan Packet { return t.inbound }
+
+func (t *ChanTUN) injectOutboundPacket(packet Packet) error {
+	if !t.beginSend() {
+		return errors.New("tun closed")
+	}
+	defer t.senders.Done()
 	select {
 	case <-t.closed:
 		return errors.New("tun closed")
-	case t.Outbound <- slices.Clone(pkt):
+	case t.outbound <- packet:
 		return nil
 	}
 }
 
+func (t *ChanTUN) tryInjectInboundPacket(packet Packet) (bool, error) {
+	if !t.beginSend() {
+		return false, errors.New("tun closed")
+	}
+	defer t.senders.Done()
+	select {
+	case <-t.closed:
+		return false, errors.New("tun closed")
+	case t.inbound <- packet:
+		return true, nil
+	default:
+		return false, nil
+	}
+}
+
+func (t *ChanTUN) beginSend() bool {
+	t.sendMu.Lock()
+	defer t.sendMu.Unlock()
+	if t.isClosed {
+		return false
+	}
+	t.senders.Add(1)
+	return true
+}
+
 func (t *ChanTUN) MTU() (int, error)        { return t.mtu, nil }
 func (t *ChanTUN) Name() (string, error)    { return t.name, nil }
 func (t *ChanTUN) Events() <-chan tun.Event { return t.events }
-func (t *ChanTUN) BatchSize() int           { return 1 }
+func (t *ChanTUN) BatchSize() int           { return conn.IdealBatchSize }