tun: Size staged slots from the workers and listen.batch

The TX and RX slot counts were fixed at eight and six. TX now uses one slot
per worker plus the one being read and the one being sent, at least eight.
RX keeps up to 384 recvmmsg entries in flight in four to six slots, so a
larger listen.batch takes fewer slots and RX memory never exceeds today's.
tun.staged.tx_slots and rx_slots override either, down to two. The defaults
at listen.batch 64 with up to six workers are unchanged: eight TX slots of
about 1.1 MiB and six RX slots of 4 MiB with GRO.

examples/config.yml

@@ -286,8 +286,14 @@ tun:
   # also delays the other tunnels of the same read, and readers are not guaranteed fairness.
   # Idle batches below 256 TX segments or RX packets stay on the reader. With multiple routines, those sends still
   # take tunnel reservation credit.
+  # tx_slots and rx_slots bound the batches each routine has in flight per direction, minimum 2. A TX slot holds one
+  # tun read and its ciphertext, about 1.1 MiB; the default is 8, or workers + 2 if more. An RX slot holds one
+  # recvmmsg batch, listen.batch * 64 KiB with udp_offloads or * 9 KiB without; the default keeps up to 384 receive
+  # entries in flight in 4 to 6 slots: 6 slots or 24 MiB at the default listen.batch of 64. Lower them to save memory.
   #staged:
   #  workers: 0
+  #  tx_slots: 8
+  #  rx_slots: 6
 
   # Linux only. pin_threads pins each tun reader/encrypt OS thread to a single CPU. This keeps every goroutine's
   # batched sends flowing through one XPS-selected NIC TX ring, so packets within a flow stay ordered on the wire

flow.go

@@ -16,8 +16,11 @@ import (
 // behind which an ordered stage sends or writes every batch in read order. It reuses the inline packet handlers:
 // staging changes ownership and scheduling, not policy. Each routine has its own workers and ordered stages, so order
 // is per routine. Full queues block the reader. A read error drains the batches in flight, while shutdown waits for
-// the crypto and then drops packets whose destination has closed. TX uses eight slots and RX six; at listen.batch=64
-// with GRO their buffers take about 9 MiB and 24 MiB per routine. Small-MSS TX arenas can grow.
+// the crypto and then drops packets whose destination has closed. A TX slot holds one tun read and its ciphertext,
+// about 1.1 MiB; there are eight, or two more than the workers. An RX slot holds one recvmmsg batch, listen.batch
+// times 64 KiB with GRO, and the count keeps up to 384 entries in flight in four to six slots: six and 24 MiB at
+// listen.batch=64.
+// tun.staged.tx_slots and rx_slots override the counts. Small-MSS TX arenas can grow.
 //
 // Metrics are registered only for active pipelines. flow.{tx,rx}.batches counts reads handed to the ordered stage and
 // flow.{tx,rx}.inline those an idle reader handled itself (TX inline also counts oversized reads); their rates show
@@ -33,6 +36,36 @@ type stagedConfig struct {
 	rx      bool
 	// inlineBelow is the TX segment or RX packet count below which an idle reader skips the worker handoff
 	inlineBelow int
+	// txSlots and rxSlots are how many batches each routine can have in flight per direction
+	txSlots, rxSlots int
+}
+
+const (
+	// stagedTxSlots is the default TX slot count, unless the workers need more
+	stagedTxSlots = 8
+	// stagedRxEntries is how many recvmmsg entries the RX slots hold in flight by default
+	stagedRxEntries = 6 * 64
+	// stagedMinSlots lets the reader fill one batch while another is in flight
+	stagedMinSlots = 2
+	// stagedMinRxSlots keeps the default RX pipeline deep enough for large listen.batch values
+	stagedMinRxSlots = 4
+	// stagedMaxRxSlots is the measured default, a smaller listen.batch does not take more slots
+	stagedMaxRxSlots = 6
+)
+
+// stagedSlots sizes the TX and RX slots. TX reads are bounded by the tun queue, so one slot per worker plus the one
+// being read and the one being sent, and at least the eight measured. RX slots keep stagedRxEntries in flight, four
+// to six, so a larger listen.batch means fewer slots and RX memory never exceeds the six measured. A positive override replaces either, down to stagedMinSlots.
+func stagedSlots(workers, batch, txOverride, rxOverride int) (tx, rx int) {
+	tx = max(stagedTxSlots, workers+2)
+	rx = min(max(stagedMinRxSlots, (stagedRxEntries+batch-1)/max(batch, 1)), stagedMaxRxSlots)
+	if txOverride > 0 {
+		tx = max(txOverride, stagedMinSlots)
+	}
+	if rxOverride > 0 {
+		rx = max(rxOverride, stagedMinSlots)
+	}
+	return tx, rx
 }
 
 // parseStagedConfig enables staging on Linux only; queues and conns that cannot hand off batches also stay inline.
@@ -52,7 +85,10 @@ func parseStagedConfig(c *config.C, l *slog.Logger) stagedConfig {
 		l.Warn("tun.staged cannot encrypt in parallel with boringcrypto or fips140; transmit stays inline")
 		sc.tx = false
 	}
-	l.Info("Staged packet processing enabled", "workers", sc.workers, "tx", sc.tx, "rx", sc.rx)
+	sc.txSlots, sc.rxSlots = stagedSlots(sc.workers, c.GetInt("listen.batch", 64),
+		c.GetInt("tun.staged.tx_slots", 0), c.GetInt("tun.staged.rx_slots", 0))
+	l.Info("Staged packet processing enabled", "workers", sc.workers, "tx", sc.tx, "rx", sc.rx,
+		"txSlots", sc.txSlots, "rxSlots", sc.rxSlots)
 	return sc
 }
 

flow_rx.go

@@ -16,7 +16,6 @@ type handoffConn interface {
 	ListenOutHandoff(udp.EncReader, func([][]byte) [][]byte) error
 }
 
-const stagedRxSlots = 6
 const rxChunkMin = 16
 
 // rxItem is one received packet. A data packet for an established tunnel has ci set; out and err are written by
@@ -108,8 +107,8 @@ func newStagedRx(f *Interface, q int, wrxc *rxContext) *stagedRx {
 		h:             &header.H{},
 		hostmapCache:  map[uint32]*HostInfo{},
 		wrxc:          wrxc,
-		ordered:       make(chan *rxBatch, stagedRxSlots),
-		returned:      make(chan *rxBatch, stagedRxSlots),
+		ordered:       make(chan *rxBatch, f.staged.rxSlots),
+		returned:      make(chan *rxBatch, f.staged.rxSlots),
 		writerDone:    make(chan struct{}),
 		readBusy:      stageBusy("rx_read", q),
 		decryptBusy:   stageBusy("rx_decrypt", q),
@@ -118,8 +117,8 @@ func newStagedRx(f *Interface, q int, wrxc *rxContext) *stagedRx {
 		metricInline:  metrics.GetOrRegisterCounter("flow.rx.inline", nil),
 	}
 	workers := f.staged.workers
-	p.pool = newFlowPool(workers, stagedRxSlots*2*workers, p.decryptBusy)
-	for range stagedRxSlots {
+	p.pool = newFlowPool(workers, f.staged.rxSlots*2*workers, p.decryptBusy)
+	for range f.staged.rxSlots {
 		// Two chunks per worker, so the writer can start committing while the workers are still on the rest
 		b := &rxBatch{items: make([]rxItem, 0, 256), chunks: make([]rxChunk, 2*workers)}
 		for i := range b.chunks {

flow_test.go

@@ -355,12 +355,12 @@ func delayCrypto(peer *flowTestPeer, before func(uint64)) {
 
 // stagedOn enables staging with n workers.
 func stagedOn(n int) stagedConfig {
-	return stagedConfig{workers: n, tx: true, rx: true}
+	return stagedConfig{workers: n, tx: true, rx: true, txSlots: stagedTxSlots, rxSlots: 6}
 }
 
 // adaptive is staging with the inline threshold parseStagedConfig sets.
 func adaptive(n int) stagedConfig {
-	return stagedConfig{workers: n, tx: true, rx: true, inlineBelow: 256}
+	return stagedConfig{workers: n, tx: true, rx: true, inlineBelow: 256, txSlots: stagedTxSlots, rxSlots: 6}
 }
 
 func stagedName(sc stagedConfig) string {

flow_tx.go

@@ -234,8 +234,6 @@ type handoffQueue interface {
 	SwapRxBuf([]byte) []byte
 }
 
-const stagedTxSlots = 8
-
 // txSlot is one staged batch. The reader fills items, a worker encrypts them into sb and signals done, and the
 // ordered sender flushes sb.
 type txSlot struct {
@@ -292,8 +290,8 @@ func newStagedTx(f *Interface, queue handoffQueue, q int) *stagedTx {
 		queue:         queue,
 		sb:            batch.NewSendBatch(f.writers[q], batch.SendBatchCap, batch.SendBatchCap*(udp.MTU+32)),
 		nb:            make([]byte, 12),
-		ordered:       make(chan *txSlot, stagedTxSlots),
-		returned:      make(chan *txSlot, stagedTxSlots),
+		ordered:       make(chan *txSlot, f.staged.txSlots),
+		returned:      make(chan *txSlot, f.staged.txSlots),
 		senderDone:    make(chan struct{}),
 		readBusy:      stageBusy("tx_read", q),
 		encryptBusy:   stageBusy("tx_encrypt", q),
@@ -309,12 +307,12 @@ func newStagedTx(f *Interface, queue handoffQueue, q int) *stagedTx {
 			p.senderCPU = f.cpuAffinity[(q+1)%n]
 		}
 	}
-	p.pool = newFlowPool(f.staged.workers, stagedTxSlots, p.encryptBusy)
+	p.pool = newFlowPool(f.staged.workers, f.staged.txSlots, p.encryptBusy)
 
 	// Ciphertext adds per-segment headers and tags; the arena can grow for small segments.
 	buf := queue.SwapRxBuf(nil)
 	bufSize := len(buf)
-	for i := range stagedTxSlots {
+	for i := range f.staged.txSlots {
 		if i > 0 {
 			buf = make([]byte, bufSize)
 		}

main_test.go

@@ -142,6 +142,26 @@ func TestParseStagedConfig(t *testing.T) {
 		require.Equal(t, stagedConfig{}, staged(2))
 		return
 	}
-	want := stagedConfig{workers: 2, tx: !noiseutil.EncryptLockNeeded, rx: true, inlineBelow: 256}
+	want := stagedConfig{workers: 2, tx: !noiseutil.EncryptLockNeeded, rx: true, inlineBelow: 256, txSlots: 8, rxSlots: 6}
 	require.Equal(t, want, staged(2))
+
+	c := config.NewC(l)
+	c.Settings["listen"] = map[string]any{"batch": 128}
+	c.Settings["tun"] = map[string]any{"staged": map[string]any{"workers": 2, "tx_slots": 3, "rx_slots": 1}}
+	got := parseStagedConfig(c, l)
+	require.Equal(t, []int{3, 2}, []int{got.txSlots, got.rxSlots}, "overrides apply, down to the minimum")
+}
+
+func TestStagedSlots(t *testing.T) {
+	for _, tt := range []struct{ workers, batch, tx, rx int }{
+		{workers: 2, batch: 64, tx: 8, rx: 6},
+		{workers: 12, batch: 64, tx: 14, rx: 6},
+		{workers: 2, batch: 96, tx: 8, rx: 4},
+		{workers: 2, batch: 512, tx: 8, rx: 4},
+		{workers: 2, batch: 32, tx: 8, rx: 6},
+		{workers: 2, batch: 1, tx: 8, rx: 6},
+	} {
+		tx, rx := stagedSlots(tt.workers, tt.batch, 0, 0)
+		require.Equal(t, []int{tt.tx, tt.rx}, []int{tx, rx}, "workers %d batch %d", tt.workers, tt.batch)
+	}
 }