darwin: pipeline listenIn's sends so write batches can form

On darwin, listenIn encrypts into one SendBatch while a sender goroutine
writes the other (batch.SendPipeline). The sender takes whatever built up
during its last write, so writes overlap tun reads and packets queued
behind a write leave in one WriteBatch, which a connected socket sends as
one sendmsg_x. Two batches bound memory; a full producer waits for the
swap rather than dropping. One producer and one sender keep per-peer order
equal to message-counter order.

Everywhere else the branch is a false constant and listenIn is unchanged.
-tags send_pipeline builds the pipeline on any platform, so its tests and
the e2e suite can run through it on linux.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

overlay/batch/send_pipeline.go

@@ -0,0 +1,156 @@
+//go:build darwin || send_pipeline
+
+package batch
+
+import (
+	"sync"
+)
+
+// SendPipeline splits encrypt from send. The producer (a tun reader) encrypts packets into one SendBatch while a
+// sender goroutine writes the other, and they swap whenever the sender is free. The sender takes everything the
+// producer committed since its last take, so packets that arrive during a write leave together in the next
+// WriteBatch, and the reader returns to the tun while the kernel is still sending.
+//
+// The producer holds the batch for one packet at a time (Begin, reserve and commit, End) so the sender can take the
+// batch between any two packets, never while a reserved slot is still being written.
+//
+// Ownership: the producer touches its batch only between Begin and End, and the sender swaps batches only under
+// the same lock, so a slot the producer reserved is never sent half written. A batch goes back to the producer
+// only after its WriteBatch returned and it was emptied, so no slot is reused before its send completes.
+//
+// Memory is bounded at two batches. When the producer fills its batch while the sender is still writing the other,
+// Begin waits for the swap. That slows the tun reader instead of dropping, as the unpipelined path does.
+//
+// Order: a single sender writes batches in the order they were filled, each in commit order, so every peer's
+// packets leave in the order their message counters were taken.
+type SendPipeline struct {
+	out      batchWriter
+	cap      int
+	maxWrite int
+	onWrite  func(queued, written int, err error)
+
+	mu   sync.Mutex
+	cond sync.Cond
+	cur  *SendBatch // the producer's batch, guarded by mu
+	gen  uint64     // swaps so far, guarded by mu
+
+	ready chan struct{} // cap 1: cur may have packets for the sender
+	stop  chan struct{}
+	done  chan struct{}
+	once  sync.Once
+}
+
+// NewSendPipeline starts a pipeline over two batches of batchCap slots and arenaSize arena bytes each. It writes
+// to out at most maxWrite packets per WriteBatch. If onWrite is not nil, it is called on the sender goroutine
+// after every WriteBatch with how many packets that call was given, how many went out, and its error. Close stops
+// the pipeline.
+func NewSendPipeline(out batchWriter, batchCap, arenaSize, maxWrite int, onWrite func(queued, written int, err error)) *SendPipeline {
+	batchCap = max(batchCap, 1)
+	p := &SendPipeline{
+		out:      out,
+		cap:      batchCap,
+		maxWrite: max(maxWrite, 1),
+		onWrite:  onWrite,
+		cur:      NewSendBatch(out, batchCap, arenaSize),
+		ready:    make(chan struct{}, 1),
+		stop:     make(chan struct{}),
+		done:     make(chan struct{}),
+	}
+	p.cond.L = &p.mu
+	go p.run(NewSendBatch(out, batchCap, arenaSize))
+	return p
+}
+
+// Begin returns a batch with room for at least one more packet. The batch stays the producer's until End. If the
+// batch is full, Begin first waits for the sender to swap it out. Producer only.
+func (p *SendPipeline) Begin() *SendBatch {
+	p.mu.Lock()
+	if p.cur.Len() >= p.cap {
+		g := p.gen
+		p.wake()
+		for p.gen == g {
+			p.cond.Wait()
+		}
+	}
+	return p.cur
+}
+
+// End gives up the batch Begin returned. With wake set it also tells the sender the batch has packets for it, which
+// the producer does once a run of packets is in (the end of a tun read) rather than per packet, so the sender takes
+// the run whole when it is idle. A sender still writing takes the batch when it finishes, woken or not.
+func (p *SendPipeline) End(wake bool) {
+	n := p.cur.Len()
+	p.mu.Unlock()
+	if wake && n > 0 {
+		p.wake()
+	}
+}
+
+func (p *SendPipeline) wake() {
+	select {
+	case p.ready <- struct{}{}:
+	default:
+	}
+}
+
+// Close writes whatever the producer left in its batch, stops the sender, and waits for it. The producer calls it,
+// outside Begin/End, once it is done. Safe to call more than once.
+func (p *SendPipeline) Close() {
+	p.once.Do(func() { close(p.stop) })
+	<-p.done
+}
+
+func (p *SendPipeline) run(mine *SendBatch) {
+	defer close(p.done)
+	for {
+		stopping := false
+		select {
+		case <-p.ready:
+		case <-p.stop:
+			stopping = true
+		}
+		// Keep taking whatever the producer committed during the last write, without waiting to be woken, so a
+		// busy sender goes straight from one write to the next and each write carries everything that built up.
+		// Once stopping, the producer is done and this writes what it left.
+		for p.take(&mine) {
+			p.write(mine)
+		}
+		if stopping {
+			return
+		}
+	}
+}
+
+// take swaps the producer's batch for the empty *mine if the producer's has packets, and reports whether it did.
+func (p *SendPipeline) take(mine **SendBatch) bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.cur.Len() == 0 {
+		return false
+	}
+	*mine, p.cur = p.cur, *mine
+	p.gen++
+	p.cond.Broadcast()
+	return true
+}
+
+// write sends b's packets, maxWrite at a time, then empties b.
+func (p *SendPipeline) write(b *SendBatch) {
+	for off := 0; off < len(b.bufs); {
+		n := min(p.maxWrite, len(b.bufs)-off)
+		written, err := p.out.WriteBatch(b.bufs[off:off+n], b.dsts[off:off+n])
+		if p.onWrite != nil {
+			p.onWrite(n, written, err)
+		}
+		off += n
+	}
+	b.recycle()
+}
+
+// recycle empties b without writing it. SendPipeline uses it because it writes a batch's packets itself.
+func (b *SendBatch) recycle() {
+	clear(b.bufs)
+	b.bufs = b.bufs[:0]
+	b.dsts = b.dsts[:0]
+	b.arena.Reset()
+}

overlay/batch/send_pipeline_test.go

@@ -0,0 +1,400 @@
+//go:build darwin || send_pipeline
+
+package batch
+
+import (
+	"encoding/binary"
+	"errors"
+	"net/netip"
+	"sync"
+	"testing"
+	"time"
+	"unsafe"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
+)
+
+// pipeWriter records every WriteBatch a SendPipeline makes. It checks each packet's bytes against the sequence
+// number they carry while the call is in progress, and can hold a call until the test releases it.
+type pipeWriter struct {
+	t *testing.T
+
+	mu     sync.Mutex
+	seqs   []uint32
+	dsts   []netip.AddrPort
+	calls  []int
+	inCall [][]byte // the bufs of the call in progress, for overlap checks
+
+	gate    chan struct{} // if not nil, each call waits for a receive from it
+	entered chan int      // if not nil, each call sends its size here on entry
+	short   int           // each call reports this many fewer written
+	err     error
+}
+
+func (w *pipeWriter) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) {
+	w.mu.Lock()
+	w.inCall = bufs
+	w.mu.Unlock()
+	if w.entered != nil {
+		w.entered <- len(bufs)
+	}
+	if w.gate != nil {
+		<-w.gate
+	}
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	for i, b := range bufs {
+		seq := binary.BigEndian.Uint32(b)
+		if !checkPkt(b, seq) {
+			w.t.Errorf("packet %d changed while being sent", seq)
+		}
+		w.seqs = append(w.seqs, seq)
+		w.dsts = append(w.dsts, addrs[i])
+	}
+	w.calls = append(w.calls, len(bufs))
+	w.inCall = nil
+	return max(len(bufs)-w.short, 0), w.err
+}
+
+func (w *pipeWriter) sent() ([]uint32, []netip.AddrPort, []int) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return append([]uint32(nil), w.seqs...), append([]netip.AddrPort(nil), w.dsts...), append([]int(nil), w.calls...)
+}
+
+// fillPkt writes seq and a pattern derived from it into b, as the encrypt step writes a slot.
+func fillPkt(b []byte, seq uint32) {
+	binary.BigEndian.PutUint32(b, seq)
+	for i := 4; i < len(b); i++ {
+		b[i] = byte(seq) ^ byte(i)
+	}
+}
+
+func checkPkt(b []byte, seq uint32) bool {
+	if len(b) < 4 || binary.BigEndian.Uint32(b) != seq {
+		return false
+	}
+	for i := 4; i < len(b); i++ {
+		if b[i] != byte(seq)^byte(i) {
+			return false
+		}
+	}
+	return true
+}
+
+func pktSize(seq uint32) int { return 16 + int(seq%97) }
+
+var testDsts = []netip.AddrPort{
+	netip.MustParseAddrPort("10.0.0.1:4242"),
+	netip.MustParseAddrPort("10.0.0.2:4242"),
+	netip.MustParseAddrPort("[fd00::3]:4242"),
+}
+
+// commit adds packet seq to p the way listenIn does: Begin, reserve, encrypt, commit, End.
+func commit(p *SendPipeline, seq uint32, last bool) {
+	sb := p.Begin()
+	slot := sb.Reserve(pktSize(seq))
+	fillPkt(slot, seq)
+	sb.Commit(slot, testDsts[seq%uint32(len(testDsts))])
+	p.End(last)
+}
+
+func TestSendPipelineOrder(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	w := &pipeWriter{t: t}
+	// A small arena makes the batches grow theirs, and a small maxWrite splits takes across several writes.
+	p := NewSendPipeline(w, 16, 64, 5, nil)
+
+	const n = 20000
+	seq := uint32(0)
+	for seq < n {
+		// Reads of 1 to 40 packets, with an occasional pause so the sender goes idle and is woken.
+		read := 1 + int(seq*7%40)
+		for j := 0; j < read && seq < n; j++ {
+			commit(p, seq, j == read-1 || seq == n-1)
+			seq++
+		}
+		if seq%1000 < 40 {
+			time.Sleep(50 * time.Microsecond)
+		}
+	}
+	p.Close()
+
+	seqs, dsts, calls := w.sent()
+	require.Len(t, seqs, n)
+	for i, s := range seqs {
+		require.Equal(t, uint32(i), s, "packet %d out of order", i)
+		require.Equal(t, testDsts[s%uint32(len(testDsts))], dsts[i])
+	}
+	for _, c := range calls {
+		assert.LessOrEqual(t, c, 5, "WriteBatch got more than maxWrite")
+		assert.Positive(t, c)
+	}
+}
+
+func TestSendPipelineConcurrentPeersKeepOrder(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	// A slow writer makes the sender take many packets per write, interleaving the three peers in each batch.
+	w := &pipeWriter{t: t}
+	slow := &slowWriter{pipeWriter: w, d: 20 * time.Microsecond}
+	p := NewSendPipeline(slow, 128, 128*64, 128, nil)
+	const n = 5000
+	for i := uint32(0); i < n; i++ {
+		commit(p, i, i%8 == 7 || i == n-1)
+	}
+	p.Close()
+
+	seqs, dsts, _ := w.sent()
+	require.Len(t, seqs, n)
+	last := map[netip.AddrPort]int64{}
+	for i, s := range seqs {
+		prev, ok := last[dsts[i]]
+		if ok {
+			require.Greater(t, int64(s), prev, "peer %v went backwards", dsts[i])
+		}
+		last[dsts[i]] = int64(s)
+	}
+}
+
+type slowWriter struct {
+	*pipeWriter
+	d time.Duration
+}
+
+func (s *slowWriter) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) {
+	time.Sleep(s.d)
+	return s.pipeWriter.WriteBatch(bufs, addrs)
+}
+
+// TestSendPipelineBackpressureAndReuse holds the sender inside a write and checks that the producer can fill
+// exactly one more batch, that none of its slots overlap the packets being written, and that the next Begin waits
+// until the write returns.
+func TestSendPipelineBackpressureAndReuse(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	const batchCap = 8
+	w := &pipeWriter{t: t, gate: make(chan struct{}), entered: make(chan int, 16)}
+	p := NewSendPipeline(w, batchCap, 64, 128, nil)
+
+	seq := uint32(0)
+	for ; seq < batchCap; seq++ {
+		commit(p, seq, seq == batchCap-1)
+	}
+	require.Equal(t, batchCap, <-w.entered, "sender should take the first batch whole")
+
+	w.mu.Lock()
+	inFlight := append([][]byte(nil), w.inCall...)
+	w.mu.Unlock()
+
+	// The second batch fills while the first is being written, in slots that don't touch the first's bytes.
+	for ; seq < 2*batchCap; seq++ {
+		sb := p.Begin()
+		slot := sb.Reserve(pktSize(seq))
+		for _, b := range inFlight {
+			assert.False(t, overlaps(slot, b), "slot for packet %d overlaps packet %d, still being sent", seq, binary.BigEndian.Uint32(b))
+		}
+		fillPkt(slot, seq)
+		sb.Commit(slot, testDsts[0])
+		p.End(true)
+	}
+
+	// Both batches are now spoken for, so the next Begin must wait for the write to return.
+	begun := make(chan struct{})
+	go func() {
+		commit(p, seq, true)
+		close(begun)
+	}()
+	select {
+	case <-begun:
+		t.Fatal("Begin returned with both batches in use: the pipeline is not bounded")
+	case <-time.After(50 * time.Millisecond):
+	}
+	for i, b := range inFlight {
+		assert.True(t, checkPkt(b, uint32(i)), "packet %d was overwritten before its send completed", i)
+	}
+
+	w.gate <- struct{}{} // first write returns; the sender swaps and starts on the second batch
+	assert.Equal(t, batchCap, <-w.entered)
+	select {
+	case <-begun:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Begin did not return after the sender freed a batch")
+	}
+	w.gate <- struct{}{}
+	go func() {
+		for range w.entered {
+			w.gate <- struct{}{}
+		}
+	}()
+	p.Close()
+	close(w.entered)
+
+	seqs, _, calls := w.sent()
+	require.Len(t, seqs, 2*batchCap+1)
+	for i, s := range seqs {
+		assert.Equal(t, uint32(i), s)
+	}
+	assert.Equal(t, []int{batchCap, batchCap, 1}, calls)
+}
+
+func overlaps(a, b []byte) bool {
+	if len(a) == 0 || len(b) == 0 {
+		return false
+	}
+	a0, b0 := uintptr(unsafe.Pointer(&a[0])), uintptr(unsafe.Pointer(&b[0]))
+	return a0 < b0+uintptr(len(b)) && b0 < a0+uintptr(len(a))
+}
+
+// TestSendPipelineBatchesBuildUp checks the point of the pipeline: while one write is in the kernel, everything the
+// producer commits leaves together in the next write.
+func TestSendPipelineBatchesBuildUp(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	w := &pipeWriter{t: t, gate: make(chan struct{}), entered: make(chan int, 16)}
+	p := NewSendPipeline(w, 128, 128*64, 128, nil)
+
+	commit(p, 0, true)
+	require.Equal(t, 1, <-w.entered)
+	// Ten one-packet reads while the first write is held.
+	for i := uint32(1); i <= 10; i++ {
+		commit(p, i, true)
+	}
+	w.gate <- struct{}{}
+	require.Equal(t, 10, <-w.entered, "packets queued during a write should go out in one WriteBatch")
+	w.gate <- struct{}{}
+	p.Close()
+
+	_, _, calls := w.sent()
+	assert.Equal(t, []int{1, 10}, calls)
+}
+
+func TestSendPipelineCloseFlushes(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	w := &pipeWriter{t: t}
+	p := NewSendPipeline(w, 32, 32*64, 128, nil)
+	// End(false) never wakes the sender, so only Close's final take sends these.
+	for i := uint32(0); i < 5; i++ {
+		commit(p, i, false)
+	}
+	p.Close()
+	p.Close() // idempotent
+	seqs, _, _ := w.sent()
+	assert.Equal(t, []uint32{0, 1, 2, 3, 4}, seqs)
+}
+
+func TestSendPipelineCloseEmpty(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	w := &pipeWriter{t: t}
+	p := NewSendPipeline(w, 4, 64, 128, nil)
+	sb := p.Begin()
+	p.End(true) // nothing committed: no write
+	assert.Zero(t, sb.Len())
+	p.Close()
+	_, _, calls := w.sent()
+	assert.Empty(t, calls)
+}
+
+// TestSendPipelineCloseWaitsForWrite checks that Close returns only after a write in progress has returned, so the
+// caller can't free anything the sender still uses.
+func TestSendPipelineCloseWaitsForWrite(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	w := &pipeWriter{t: t, gate: make(chan struct{}), entered: make(chan int, 4)}
+	p := NewSendPipeline(w, 4, 64, 128, nil)
+	commit(p, 0, true)
+	<-w.entered
+	commit(p, 1, false)
+
+	closed := make(chan struct{})
+	go func() {
+		p.Close()
+		close(closed)
+	}()
+	select {
+	case <-closed:
+		t.Fatal("Close returned during a write")
+	case <-time.After(50 * time.Millisecond):
+	}
+	w.gate <- struct{}{}
+	<-w.entered // Close's final take writes packet 1
+	w.gate <- struct{}{}
+	select {
+	case <-closed:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Close did not return")
+	}
+	seqs, _, _ := w.sent()
+	assert.Equal(t, []uint32{0, 1}, seqs)
+}
+
+func TestSendPipelineOnWriteAndDrops(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	werr := errors.New("boom")
+	w := &pipeWriter{t: t, short: 1, err: werr}
+	type call struct {
+		queued, written int
+		err             error
+	}
+	var got []call
+	p := NewSendPipeline(w, 16, 16*64, 4, func(q, n int, err error) {
+		got = append(got, call{q, n, err}) // sender goroutine only
+	})
+	for i := uint32(0); i < 10; i++ {
+		commit(p, i, false)
+	}
+	p.Close()
+
+	assert.Equal(t, []call{{4, 3, werr}, {4, 3, werr}, {2, 1, werr}}, got)
+}
+
+// TestSendPipelineStressRace runs a producer against a sender whose writes take varying time, for the race
+// detector, and checks nothing is lost, reordered or overwritten.
+func TestSendPipelineStressRace(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	w := &pipeWriter{t: t}
+	jitter := &jitterWriter{pipeWriter: w}
+	p := NewSendPipeline(jitter, 32, 256, 16, nil)
+	const n = 50000
+	for i := uint32(0); i < n; i++ {
+		commit(p, i, i%3 == 0)
+	}
+	p.Close()
+	seqs, _, _ := w.sent()
+	require.Len(t, seqs, n)
+	for i, s := range seqs {
+		if uint32(i) != s {
+			t.Fatalf("packet %d arrived at %d", s, i)
+		}
+	}
+}
+
+type jitterWriter struct {
+	*pipeWriter
+	i int
+}
+
+func (j *jitterWriter) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, error) {
+	j.i++
+	if j.i%7 == 0 {
+		time.Sleep(time.Duration(j.i%5) * 10 * time.Microsecond)
+	}
+	return j.pipeWriter.WriteBatch(bufs, addrs)
+}
+
+type discardWriter struct{}
+
+func (discardWriter) WriteBatch(bufs [][]byte, _ []netip.AddrPort) (int, error) {
+	return len(bufs), nil
+}
+
+// BenchmarkSendPipeline measures a packet through Begin, reserve, commit and End with a writer that costs nothing;
+// its allocs/op is the pipeline's steady-state heap cost per packet.
+func BenchmarkSendPipeline(b *testing.B) {
+	p := NewSendPipeline(discardWriter{}, SendBatchCap, SendBatchCap*1500, 128, nil)
+	defer p.Close()
+	b.ReportAllocs()
+	for i := range b.N {
+		sb := p.Begin()
+		sb.Commit(sb.Reserve(1400), testDsts[0])
+		p.End(i%8 == 7)
+	}
+}