Jelajahi Sumber

tun: Add opt-in staged transmit processing

Set tun.staged.workers to a positive count on Linux, with tun.use_offloads
enabled, to move encryption off each tun reader. The reader filters and
reserves counters for a whole read, hands the batch to that many encrypt
workers, and an ordered sender flushes the batches in read order. Eight slots
per routine bound the batches in flight; a full pipeline blocks the reader.
Zero workers, the default, keeps today's inline path, as do tun queues that
cannot hand off their read buffer and builds whose AES-GCM needs nonce-ordered
encryption (BoringCrypto, FIPS).

IMPORTANT: staged senders need peers with the 8192-entry replay window of
master. Released v1.11.x peers use 1024 and will drop reordered batches. The
window is not negotiated.

Readers keep today's pinning. The sender is only pinned with an explicit
tun.cpu_affinity: sender i uses entry i+1, unless that is its reader's CPU,
so a one-entry list leaves the sender unpinned rather than sharing it: two
busy transmit threads on one CPU serialise (eight TCP streams over 100 GbE
fell from 25.9 to 22.8 Gbit/s). Without a list the sender is not pinned. Placing it by queue index on a
hybrid CPU put it on an efficiency core, and one TCP stream over 100 GbE
reached 18.6-19.6 Gbit/s pinned against 20.2-20.6 unpinned.

New API, each needed across a package boundary:
- tio.Offload.SwapRxBuf hands the buffer of the last Read to the reader and
  reads into the one given. It is a method on the concrete queue, found through
  an unexported interface in nebula, so tio.Queue and its other
  implementations are unchanged.
- batch.SendBatch.Reset drops queued packets without sending, for a sender
  whose sockets are already closed; Flush now uses it.

Metrics, registered only for active pipelines: flow.tx.batches counts reads
handed to the sender, flow.stage.{tx_read,tx_encrypt,tx_send}.q<N>.busy_ns the
elapsed busy time of each stage, to find the limiting one.

Tests cover ordering and counters under backpressure with both ciphers (AES
skipped where nonce order is required), read errors with batches in flight,
shutdown, exhaustion, relays, fallback, firewall parity, and the actual reader
and send thread affinity for every placement, with pinning on and off: the
sender is pinned only to an explicit entry other than its reader's CPU.
Ryan Huber 5 hari lalu
induk
melakukan
f3f98e1690
11 mengubah file dengan 832 tambahan dan 24 penghapusan
  1. 12 1
      examples/config.yml
  2. 81 0
      flow.go
  3. 112 1
      flow_test.go
  4. 226 0
      flow_tx.go
  5. 288 1
      flow_tx_test.go
  6. 43 16
      interface.go
  7. 5 1
      main.go
  8. 22 0
      main_test.go
  9. 9 4
      overlay/batch/tx_batch.go
  10. 15 0
      overlay/tio/tio_gso_linux.go
  11. 19 0
      overlay/tio/tun_linux_offload_test.go

+ 12 - 1
examples/config.yml

@@ -275,9 +275,18 @@ tun:
   # the use_offloads setting controls if Nebula will attempt to enable GSO and GRO for the tun device. Linux only, not reloadable.
   #use_offloads: false
 
+  # Linux only. staged.workers enables staged packet processing: each routine hands its batches to this many crypto
+  # workers per direction, then sends or writes them in order. 0, the default, keeps processing inline. Not reloadable.
+  # Transmit needs use_offloads; BoringCrypto and FIPS builds keep it inline because AES-GCM there needs ordered nonces.
+  # IMPORTANT: staged senders require peers with at least the 8192-entry replay window of this release. Released
+  # v1.11.x peers use 1024 and are incompatible. The window is not negotiated.
+  #staged:
+  #  workers: 0
+
   # 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
   # instead of being sprayed across multiple TX rings and reordered. Not reloadable. Coerced to false if routines <= 1.
+  # Staged senders are only pinned with an explicit cpu_affinity, see below; otherwise the scheduler places them.
   #pin_threads: true
 
   # pin_threads_key helps the CPU-auto-selector shuffle which CPUs are chosen for pinning.
@@ -287,7 +296,9 @@ tun:
 
   # Linux only. cpu_affinity overrides which CPUs the tun reader threads pin to: a list of CPU IDs, one per routine
   # (see the top-level `routines` setting). Lists shorter than `routines` are modulo-cycled across the queues; extra
-  # entries are ignored. IDs must be within the process's allowed CPU set, so this respects taskset / cgroup cpusets;
+  # entries are ignored. Staged sender i pins to entry i+1, modulo-cycled the same way, unless that is its reader's
+  # CPU, as with a one-entry list: a sender never shares its reader's CPU and stays unpinned instead.
+  # IDs must be within the process's allowed CPU set, so this respects taskset / cgroup cpusets;
   # a non-integer or not-allowed entry disables the override, leaving the default pin selection described below.
   # Only meaningful while pin_threads is true. Not reloadable.
   # When unset (or rejected), the default spread prefers performance cores on heterogeneous CPUs (ARM big.LITTLE,

+ 81 - 0
flow.go

@@ -0,0 +1,81 @@
+package nebula
+
+import (
+	"fmt"
+	"log/slog"
+	"runtime"
+	"sync"
+	"time"
+
+	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/config"
+	"github.com/slackhq/nebula/noiseutil"
+)
+
+// Staged transmit (tun.staged.workers) moves encryption off each tun reader onto worker goroutines, behind which an
+// ordered sender flushes every batch in read order. It reuses the inline filter and encryption: staging changes
+// ownership and scheduling, not policy. Each routine has its own workers and sender, 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 sockets have closed. Eight slots take about 9 MiB per routine; small-MSS arenas can grow.
+//
+// Metrics are registered only for active pipelines. flow.tx.batches counts reads handed to the sender.
+// flow.stage.<stage>.q<queue>.busy_ns locates the limiting stage: tx_read, tx_encrypt or tx_send. Take rates; 1e9
+// ns/s fills a serial stage. These are elapsed times that exclude read syscalls and downstream waits; descheduling
+// inflates them and worker times add up across the pool.
+
+// stagedConfig is tun.staged. workers is per routine and direction; zero keeps processing inline.
+type stagedConfig struct {
+	workers int
+	tx      bool
+}
+
+// parseStagedConfig enables staging on Linux only; queues and conns that cannot hand off batches also stay inline.
+func parseStagedConfig(c *config.C, l *slog.Logger) stagedConfig {
+	sc := stagedConfig{workers: c.GetInt("tun.staged.workers", 0)}
+	if sc.workers <= 0 {
+		return stagedConfig{}
+	}
+	if runtime.GOOS != "linux" {
+		l.Warn("tun.staged is only supported on Linux; packet processing stays inline")
+		return stagedConfig{}
+	}
+	sc.tx = true
+	if noiseutil.EncryptLockNeeded {
+		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)
+	return sc
+}
+
+// flowJob is one batch of crypto work. It owns its packet buffers until run returns; jobs and queues are reused.
+type flowJob interface{ run(nb []byte) }
+
+// flowPool runs one reader's crypto jobs. The reader is the only producer and closes the pool once its ordered stage
+// has drained.
+type flowPool struct {
+	jobs chan flowJob
+	wg   sync.WaitGroup
+}
+
+func newFlowPool(workers, queueLen int, busy metrics.Counter) *flowPool {
+	p := &flowPool{jobs: make(chan flowJob, queueLen)}
+	for range workers {
+		p.wg.Go(func() {
+			nb := make([]byte, 12)
+			for job := range p.jobs {
+				t := time.Now()
+				job.run(nb)
+				busy.Inc(int64(time.Since(t)))
+			}
+		})
+	}
+	return p
+}
+
+func (p *flowPool) close() { close(p.jobs); p.wg.Wait() }
+
+// stageBusy is the busy time counter of one stage of routine q.
+func stageBusy(stage string, q int) metrics.Counter {
+	return metrics.GetOrRegisterCounter(fmt.Sprintf("flow.stage.%s.q%d.busy_ns", stage, q), nil)
+}

+ 112 - 1
flow_test.go

@@ -7,7 +7,9 @@ import (
 	"io"
 	"math"
 	"net/netip"
+	"os"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -18,11 +20,13 @@ import (
 	"github.com/slackhq/nebula/firewall"
 	"github.com/slackhq/nebula/header"
 	"github.com/slackhq/nebula/iputil"
+	"github.com/slackhq/nebula/noiseutil"
 	"github.com/slackhq/nebula/overlay/batch"
 	"github.com/slackhq/nebula/overlay/tio"
 	"github.com/slackhq/nebula/test"
 	"github.com/slackhq/nebula/udp"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 )
 
 // flowTestMSS is the segment payload size of the synthetic TSO superpackets.
@@ -106,7 +110,7 @@ type flowTestPeer struct {
 
 // newFlowTestPair builds two interfaces joined by one real tunnel (keys from an actual handshake). sender sends from
 // flowTestLocal to flowTestRemote, receiver is the far side.
-func newFlowTestPair(t testing.TB, cipher noise.CipherFunc) (sender, receiver *flowTestPeer) {
+func newFlowTestPair(t testing.TB, cipher noise.CipherFunc, sc stagedConfig) (sender, receiver *flowTestPeer) {
 	t.Helper()
 	initR, respR := runTestHandshakeWith(t, cipher)
 	sci, err := newConnectionStateFromResult(initR)
@@ -115,6 +119,7 @@ func newFlowTestPair(t testing.TB, cipher noise.CipherFunc) (sender, receiver *f
 	require.NoError(t, err)
 	sender = newFlowTestPeer(t, flowTestLocal, flowTestRemote, sci, 100, 200)
 	receiver = newFlowTestPeer(t, flowTestRemote, flowTestLocal, rci, 200, 100)
+	sender.f.staged, receiver.f.staged = sc, sc
 	return sender, receiver
 }
 
@@ -254,3 +259,109 @@ func decryptWire(t testing.TB, receiver *flowTestPeer, wire [][]byte) []decodedS
 	}
 	return out
 }
+
+// flowBatchQueue is a handoffQueue that hands out scripted Read batches, then fails with io.EOF. Every batch is copied
+// into the current receive buffer, so a test sees the buffer ownership the Offload queue has.
+type flowBatchQueue struct {
+	batches [][]tio.Packet
+	rxBuf   []byte
+	out     []tio.Packet
+	closed  chan struct{}
+	// reads counts the batches handed out
+	reads atomic.Int64
+}
+
+func newFlowBatchQueue(batches [][]tio.Packet) *flowBatchQueue {
+	return &flowBatchQueue{batches: batches, rxBuf: make([]byte, 8*65536), closed: make(chan struct{})}
+}
+
+func (q *flowBatchQueue) Read() ([]tio.Packet, error) {
+	select {
+	case <-q.closed:
+		return nil, os.ErrClosed
+	default:
+	}
+	n := int(q.reads.Load())
+	if n == len(q.batches) {
+		return nil, io.EOF
+	}
+	q.out = q.out[:0]
+	off := 0
+	for _, p := range q.batches[n] {
+		l := copy(q.rxBuf[off:], p.Bytes)
+		q.out = append(q.out, tio.Packet{Bytes: q.rxBuf[off : off+l : off+l], GSO: p.GSO})
+		off += l
+	}
+	q.reads.Add(1)
+	return q.out, nil
+}
+
+func (q *flowBatchQueue) SwapRxBuf(buf []byte) []byte {
+	if buf == nil {
+		buf = make([]byte, len(q.rxBuf))
+	}
+	old := q.rxBuf
+	q.rxBuf = buf
+	// Scribble over the buffer we are handed, a staged batch that still pointed into it would be corrupted
+	for i := range buf {
+		buf[i] = 0xa5
+	}
+	return old
+}
+
+func (q *flowBatchQueue) Write(b []byte) (int, error) { return len(b), nil }
+
+func (q *flowBatchQueue) Close() error { close(q.closed); return nil }
+
+// delayedCipher changes scheduling through the real cipher interface, only in tests.
+type delayedCipher struct {
+	noiseutil.CipherState
+	before func(uint64)
+}
+
+func (c delayedCipher) EncryptDanger(out, ad, plain []byte, n uint64, nb []byte) ([]byte, error) {
+	c.before(n)
+	return c.CipherState.EncryptDanger(out, ad, plain, n, nb)
+}
+
+func (c delayedCipher) DecryptDanger(out, ad, ciphertext []byte, n uint64, nb []byte) ([]byte, error) {
+	c.before(n)
+	return c.CipherState.DecryptDanger(out, ad, ciphertext, n, nb)
+}
+
+func delayCrypto(peer *flowTestPeer, before func(uint64)) {
+	ci := peer.hostinfo.ConnectionState
+	ci.eKey = delayedCipher{ci.eKey, before}
+	ci.dKey = delayedCipher{ci.dKey, before}
+}
+
+// stagedOn enables staging with n workers.
+func stagedOn(n int) stagedConfig {
+	return stagedConfig{workers: n, tx: true}
+}
+
+func stagedName(sc stagedConfig) string {
+	switch {
+	case sc.workers == 0:
+		return "inline"
+	default:
+		return "staged"
+	}
+}
+
+// TestFlowStagedFallback: a tun queue that cannot hand off its buffer, or a conn that cannot hand off its batches,
+// keeps that direction on the inline path.
+func TestFlowStagedFallback(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(3))
+	batches, want := flowTestBatches(10, 2, 1, 4)
+	tx := metrics.GetOrRegisterCounter("flow.tx.batches", nil)
+	tx0 := tx.Count()
+	sender.f.listenIn(plainQueue{newFlowBatchQueue(batches)}, 0)
+	wire := sender.conn.captured()
+	require.Equal(t, want, decryptWire(t, receiver, wire))
+	require.Equal(t, tx0, tx.Count())
+}
+
+// plainQueue and plainConn hide SwapRxBuf and ListenOutHandoff, so they cannot hand off their batches.
+type plainQueue struct{ tio.Queue }

+ 226 - 0
flow_tx.go

@@ -3,7 +3,9 @@ package nebula
 import (
 	"errors"
 	"net/netip"
+	"time"
 
+	"github.com/rcrowley/go-metrics"
 	"github.com/slackhq/nebula/firewall"
 	"github.com/slackhq/nebula/header"
 	"github.com/slackhq/nebula/noiseutil"
@@ -231,3 +233,227 @@ func (r *txReserver) reserve(items []txItem) []txItem {
 	clear(items[n:])
 	return items[:n]
 }
+
+// handoffQueue is a tun queue that can give up the buffer its last Read filled. SwapRxBuf(nil) allocates a
+// correctly sized replacement; later swaps recycle the buffers handed back.
+type handoffQueue interface {
+	tio.Queue
+	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 {
+	f     *Interface
+	items []txItem
+	// buf owns the queue receive buffer the items point into, swapped out at dispatch
+	buf  []byte
+	sb   *batch.SendBatch
+	done chan struct{}
+	// reserved is how many counters the batch reserved, for the reader's maxReserved budget
+	reserved uint64
+}
+
+func (s *txSlot) run(nb []byte) {
+	for i := range s.items {
+		s.f.encryptTxItem(&s.items[i], nb, s.sb)
+	}
+	s.done <- struct{}{}
+}
+
+// stagedTx is the staged transmit pipeline of one tun reader. Everything but the workers and the sender runs on the
+// reader goroutine.
+type stagedTx struct {
+	txReserver
+	q         int
+	cpu       int // reader CPU
+	senderCPU int // -1 leaves the sender within its inherited affinity
+	queue     handoffQueue
+	pool      *flowPool
+
+	// free holds the reader's idle slots. A dispatched slot goes to the workers and, in the same order, to the
+	// sender on ordered; the sender gives it back on returned once it is on the wire.
+	free       []*txSlot
+	ordered    chan *txSlot
+	returned   chan *txSlot
+	inFlight   int
+	senderDone chan struct{}
+	// reserved is the counters held by in-flight slots, maxReserved this reader's half-window share
+	reserved    uint64
+	maxReserved uint64
+
+	readBusy, encryptBusy, sendBusy metrics.Counter
+	metricBatches                   metrics.Counter
+}
+
+// newStagedTx needs a queue that hands over whole read batches, the offload (TSO/USO) queue on linux. On any other
+// queue a read is one packet, and staging it would cost a handoff per packet.
+func newStagedTx(f *Interface, queue handoffQueue, q int) *stagedTx {
+	p := &stagedTx{
+		txReserver:    txReserver{f: f, tunnels: make([]txTunnel, 0, 4)},
+		q:             q,
+		senderCPU:     -1,
+		queue:         queue,
+		ordered:       make(chan *txSlot, stagedTxSlots),
+		returned:      make(chan *txSlot, stagedTxSlots),
+		senderDone:    make(chan struct{}),
+		maxReserved:   (ReplayWindow / 2) / uint64(max(f.routines, 1)),
+		readBusy:      stageBusy("tx_read", q),
+		encryptBusy:   stageBusy("tx_encrypt", q),
+		sendBusy:      stageBusy("tx_send", q),
+		metricBatches: metrics.GetOrRegisterCounter("flow.tx.batches", nil),
+	}
+	if f.pinThreads {
+		p.cpu = f.threadCPU(q)
+		// Only an explicit list pins the sender, and never onto its reader's CPU, where the two would serialise.
+		// Automatic placement by queue index can put a busy sender on a slow core; unpinned, the scheduler moves it.
+		if n := len(f.cpuAffinity); n > 0 && f.cpuAffinity[(q+1)%n] != p.cpu {
+			p.senderCPU = f.cpuAffinity[(q+1)%n]
+		}
+	}
+	p.pool = newFlowPool(f.staged.workers, stagedTxSlots, 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 {
+		if i > 0 {
+			buf = make([]byte, bufSize)
+		}
+		p.free = append(p.free, &txSlot{
+			f:    f,
+			buf:  buf,
+			sb:   batch.NewSendBatch(f.writers[q], 4*batch.SendBatchCap, bufSize+bufSize/4),
+			done: make(chan struct{}, 1),
+		})
+	}
+	go p.sender()
+	return p
+}
+
+// run is listenIn's loop with the staged pipeline: read, then hand each batch on, until the queue fails.
+func (p *stagedTx) run(fwPacket *firewall.ParsedPacket, rejectBuf []byte, conntrackCache *firewall.ConntrackCacheTicker) error {
+	defer p.close()
+	if p.f.pinThreads {
+		p.f.pinThreadToCPU(p.q, p.cpu)
+	}
+	for {
+		pkts, err := p.queue.Read()
+		if err != nil {
+			return err
+		}
+		t := time.Now()
+		stall := p.handle(pkts, fwPacket, rejectBuf, conntrackCache.Get())
+		p.readBusy.Inc(int64(time.Since(t) - stall))
+	}
+}
+
+// handle takes one tun read through the pipeline. It returns once the batch is queued, or sent when the reader sent
+// it itself, with how long it waited for the sender.
+func (p *stagedTx) handle(pkts []tio.Packet, fwPacket *firewall.ParsedPacket, rejectBuf []byte, localCache firewall.ConntrackCache) time.Duration {
+	p.reap()
+
+	slot, stall := p.acquire()
+	items := p.collect(slot.items[:0], pkts, fwPacket, rejectBuf, p.q, localCache)
+	if len(items) == 0 {
+		slot.items = items
+		p.free = append(p.free, slot)
+		return stall
+	}
+
+	var total uint64
+	for _, t := range p.tunnels {
+		total += t.n
+	}
+	// Over this reader's share of the window, wait for its oldest batches to go out first. A lone batch always goes.
+	if p.inFlight > 0 && p.reserved+total > p.maxReserved {
+		t := time.Now()
+		for p.inFlight > 0 && p.reserved+total > p.maxReserved {
+			p.release(<-p.returned)
+		}
+		stall += time.Since(t)
+	}
+	items = p.reserve(items)
+	slot.items = items
+	if len(items) == 0 {
+		p.free = append(p.free, slot)
+		return stall
+	}
+
+	// The packets now live in the slot, and the queue reads into the slot's old buffer
+	slot.buf = p.queue.SwapRxBuf(slot.buf)
+
+	slot.reserved = total
+	p.reserved += total
+	p.inFlight++
+	p.metricBatches.Inc(1)
+	// Counters were reserved above on this goroutine and the slot takes its place in the sender's order right here,
+	// before any other batch can, so the sender puts them on the wire in counter order.
+	p.ordered <- slot
+	p.pool.jobs <- slot
+	return stall
+}
+
+// acquire returns an idle slot, waiting for the sender to give one back when all of them are in flight, and how long
+// it waited.
+func (p *stagedTx) acquire() (*txSlot, time.Duration) {
+	var stall time.Duration
+	if len(p.free) == 0 {
+		t := time.Now()
+		p.release(<-p.returned)
+		stall = time.Since(t)
+	}
+	s := p.free[len(p.free)-1]
+	p.free = p.free[:len(p.free)-1]
+	return s, stall
+}
+
+// reap takes back every slot the sender has finished with, without waiting.
+func (p *stagedTx) reap() {
+	for {
+		select {
+		case s := <-p.returned:
+			p.release(s)
+		default:
+			return
+		}
+	}
+}
+
+func (p *stagedTx) release(s *txSlot) {
+	p.reserved -= s.reserved
+	s.reserved = 0
+	p.inFlight--
+	p.free = append(p.free, s)
+}
+
+// sender is the ordered send stage: it puts each batch on the wire once its worker is done, in dispatch order.
+func (p *stagedTx) sender() {
+	defer close(p.senderDone)
+	if p.senderCPU >= 0 {
+		p.f.pinThreadToCPU(p.q, p.senderCPU)
+	}
+	for s := range p.ordered {
+		<-s.done
+		t := time.Now()
+		if p.f.closed.Load() {
+			// The sockets are gone, a flush would only log errors
+			s.sb.Reset()
+		} else {
+			p.f.flushSendBatch(s.sb, p.q)
+		}
+		p.sendBusy.Inc(int64(time.Since(t)))
+		clear(s.items)
+		s.items = s.items[:0]
+		p.returned <- s
+	}
+}
+
+// close sends everything still in flight, in order, and stops the sender and the workers.
+func (p *stagedTx) close() {
+	close(p.ordered)
+	<-p.senderDone
+	p.pool.close()
+}

+ 288 - 1
flow_tx_test.go

@@ -3,19 +3,32 @@
 package nebula
 
 import (
+	"fmt"
+	"net/netip"
+	"os"
+	"os/exec"
 	"testing"
+	"time"
 
 	"github.com/flynn/noise"
+	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/cpupick"
 	"github.com/slackhq/nebula/firewall"
+	"github.com/slackhq/nebula/header"
+	"github.com/slackhq/nebula/noiseutil"
 	"github.com/slackhq/nebula/overlay/batch"
 	"github.com/slackhq/nebula/overlay/tio"
+	"github.com/slackhq/nebula/udp"
+	"github.com/slackhq/nebula/util"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
+	"golang.org/x/sys/unix"
 )
 
 // TestReservedEncryptionMatchesInline encrypts a batch with counters reserved up front: the receiver accepts every
 // segment, in order.
 func TestReservedEncryptionMatchesInline(t *testing.T) {
-	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly)
+	sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedConfig{})
 	packets, want := flowTestBatches(1, 3, 3, 9)
 	r := txReserver{f: sender.f}
 	items := r.collect(nil, packets[0], &firewall.ParsedPacket{}, make([]byte, mtu), 0, nil)
@@ -42,3 +55,277 @@ func TestSegmentCount(t *testing.T) {
 		}
 	}
 }
+
+// TestFlowStagedTXOrderAndCounters runs the staged transmit pipeline with shuffled workers and a slow socket, so the
+// sender lags and the reader runs into backpressure, then fails the tun read with batches in flight: the wire must
+// carry every segment once, in tun read order, under strictly increasing counters that were each reserved exactly
+// once, and the read error is still reported.
+func TestFlowStagedTXOrderAndCounters(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	for _, cipher := range []noise.CipherFunc{noiseutil.CipherAESGCM, noise.CipherChaChaPoly} {
+		t.Run(cipher.CipherName(), func(t *testing.T) {
+			if cipher.CipherName() == noiseutil.CipherAESGCM.CipherName() && noiseutil.EncryptLockNeeded {
+				t.Skip("AES-GCM requires encryption in nonce order")
+			}
+			sender, receiver := newFlowTestPair(t, cipher, stagedOn(4))
+			delayCrypto(sender, func(n uint64) {
+				if n%54 == 0 {
+					time.Sleep(time.Millisecond)
+				}
+			})
+			// Far slower than the reader, even under the race detector
+			sender.conn.delay = 8 * time.Millisecond
+			batches, want := flowTestBatches(60, 6, 3, 9)
+			start := sender.hostinfo.ConnectionState.messageCounter.Load()
+			tx := metrics.GetOrRegisterCounter("flow.tx.batches", nil)
+			tx0 := tx.Count()
+			sender.f.listenIn(newFlowBatchQueue(batches), 0)
+			require.Equal(t, int64(len(batches)), tx.Count()-tx0, "every batch uses the ordered sender")
+			require.Equal(t, want, decryptWire(t, receiver, sender.conn.captured()), "wire order must match tun read order")
+			require.Equal(t, start+uint64(len(want)), sender.hostinfo.ConnectionState.messageCounter.Load())
+			require.NotNil(t, sender.f.fatalErr.Load())
+		})
+	}
+}
+
+// TestFlowStagedTXCloseInFlight closes the interface and tun while batches are in flight, 20 times over.
+func TestFlowStagedTXCloseInFlight(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	for range 20 {
+		sender, _ := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(3))
+		sender.conn.delay = time.Millisecond
+		batches, _ := flowTestBatches(400, 4, 2, 8)
+		q := newFlowBatchQueue(batches)
+		done := make(chan struct{})
+		go func() { sender.f.listenIn(q, 0); close(done) }()
+		require.Eventually(t, func() bool { return q.reads.Load() > 5 }, 5*time.Second, 100*time.Microsecond)
+		sender.f.closed.Store(true)
+		require.NoError(t, q.Close())
+		<-done
+	}
+}
+
+// TestFlowStagedTXExhausted drops every batch of a tunnel whose counter runs out, like inline sends do.
+func TestFlowStagedTXExhausted(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	sender, _ := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(2))
+	ci := sender.hostinfo.ConnectionState
+	ci.messageCounter.Store(RejectAfterMessages - 3)
+	sender.f.listenIn(newFlowBatchQueue([][]tio.Packet{{flowTCPSuperpacket(0, 0, 10)}, {flowTCPSuperpacket(0, 1, 2)}}), 0)
+	require.Empty(t, sender.conn.captured())
+	require.Equal(t, RejectAfterMessages, ci.messageCounter.Load())
+	require.Equal(t, int64(12), sender.f.messageMetrics.txExhausted.Count())
+}
+
+// TestFlowStagedTXRelay sends staged batches, and an empty read, through a relay: outer and inner layers both verify,
+// in order, and go to the relay.
+func TestFlowStagedTXRelay(t *testing.T) {
+	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
+	for _, window := range []uint64{ReplayWindow} {
+		t.Run(fmt.Sprintf("window=%d", window), func(t *testing.T) {
+			sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, stagedOn(2))
+			sender.hostinfo.ConnectionState.window = NewBits(window)
+			receiver.hostinfo.ConnectionState.window = NewBits(window)
+			relayRemote, _, relayPeer := addTestRelay(t, sender, window)
+			batches, want := flowTestBatches(2, 3, 1, 6)
+			var payload []byte
+			sender.f.listenIn(newFlowBatchQueue(append(batches, nil)), 0)
+			wire := sender.conn.captured()
+			require.Len(t, wire, len(want)+len(payload))
+			var got []decodedSeg
+			var tiny []byte
+			var lastOuter, lastInner uint64
+			for i, pkt := range wire {
+				require.Equal(t, relayRemote, sender.conn.addrs[i])
+				h := &header.H{}
+				require.NoError(t, h.Parse(pkt))
+				require.Greater(t, h.MessageCounter, lastOuter)
+				lastOuter = h.MessageCounter
+				require.NoError(t, relayPeer.VerifyRelay(nil, h.MessageCounter, pkt, make([]byte, 12)))
+				pkt = pkt[header.Len : len(pkt)-16]
+				require.NoError(t, h.Parse(pkt))
+				require.Greater(t, h.MessageCounter, lastInner)
+				lastInner = h.MessageCounter
+				plain, err := receiver.hostinfo.ConnectionState.Decrypt(receiver.f.l, h.MessageCounter, pkt, make([]byte, 12))
+				require.NoError(t, err)
+				if len(plain) == 41 {
+					tiny = append(tiny, plain[40])
+				} else {
+					got = append(got, decodeSeg(plain[40:]))
+				}
+			}
+			require.Equal(t, want, got)
+			require.Equal(t, payload, tiny)
+		})
+	}
+}
+
+// addTestRelay makes sender reach its peer only through a relay tunnel with replay windows of window. It returns the
+// relay's underlay address, its hostinfo and the relay's side of the tunnel.
+func addTestRelay(t testing.TB, sender *flowTestPeer, window uint64) (netip.AddrPort, *HostInfo, *ConnectionState) {
+	relayInit, relayResp := runTestHandshake(t)
+	relayCI, err := newConnectionStateFromResult(relayInit)
+	require.NoError(t, err)
+	relayPeer, err := newConnectionStateFromResult(relayResp)
+	require.NoError(t, err)
+	relayCI.window, relayPeer.window = NewBits(window), NewBits(window)
+
+	relayAddr := netip.MustParseAddr("10.0.0.3")
+	relayRemote := netip.MustParseAddrPort("192.0.2.3:4242")
+	relay := &Relay{Type: TerminalType, State: Established, LocalIndex: 5, RemoteIndex: 6, PeerAddr: flowTestRemote}
+	relayHI := &HostInfo{
+		vpnAddrs:        []netip.Addr{relayAddr},
+		ConnectionState: relayCI,
+		relayState: RelayState{
+			relayForByAddr: map[netip.Addr]*Relay{flowTestRemote: relay},
+			relayForByIdx:  map[uint32]*Relay{relay.LocalIndex: relay},
+		},
+	}
+	relayHI.remote.Store(&relayRemote)
+	sender.f.hostMap.Hosts[relayAddr] = relayHI
+	sender.hostinfo.remote.Store(&netip.AddrPort{})
+	sender.hostinfo.relayState.relays = []netip.Addr{relayAddr}
+	return relayRemote, relayHI, relayPeer
+}
+
+// affinityQueue makes every read start with an idle pipeline: before each read it waits until conn has sent what was
+// read before, recording how many sends that took, then samples the reader's actual thread affinity.
+type affinityQueue struct {
+	*flowBatchQueue
+	conn       *affinityConn
+	sent       int
+	boundaries []int
+	masks      []unix.CPUSet
+}
+
+func (q *affinityQueue) Read() ([]tio.Packet, error) {
+	if n := q.reads.Load(); n > 0 {
+		for _, p := range q.batches[n-1] {
+			q.sent += int(segmentCount(p))
+		}
+	}
+	q.boundaries = append(q.boundaries, q.conn.waitFor(q.sent))
+	q.masks = append(q.masks, threadMask())
+	return q.flowBatchQueue.Read()
+}
+
+// affinityConn captures every send, sampling the sending thread's actual affinity in each WriteBatch.
+type affinityConn struct {
+	flowCaptureConn
+	masks []unix.CPUSet
+}
+
+func (c *affinityConn) WriteBatch(p [][]byte, addrs []netip.AddrPort) (int, error) {
+	mask := threadMask()
+	c.mu.Lock()
+	c.masks = append(c.masks, mask)
+	c.mu.Unlock()
+	return c.flowCaptureConn.WriteBatch(p, addrs)
+}
+
+// waitFor waits until n packets were sent and returns how many WriteBatch calls that took.
+func (c *affinityConn) waitFor(n int) int {
+	for len(c.captured()) < n {
+		time.Sleep(50 * time.Microsecond)
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return len(c.masks)
+}
+
+func threadMask() unix.CPUSet {
+	var mask unix.CPUSet
+	if err := unix.SchedGetaffinity(0, &mask); err != nil {
+		panic(err)
+	}
+	return mask
+}
+
+// requireFlowCPU checks every mask is just cpu, or with cpu -1, the whole allowed set.
+func requireFlowCPU(t *testing.T, masks []unix.CPUSet, cpu int, allowed []int) {
+	t.Helper()
+	require.NotEmpty(t, masks)
+	for _, mask := range masks {
+		if cpu >= 0 {
+			require.Equal(t, 1, mask.Count())
+			require.True(t, mask.IsSet(cpu), "expected CPU %d", cpu)
+			continue
+		}
+		require.Equal(t, len(allowed), mask.Count())
+		for _, c := range allowed {
+			require.True(t, mask.IsSet(c))
+		}
+	}
+}
+
+// TestFlowStagedTXAffinity sends large, small and large reads, each from an idle pipeline, and checks the order on the
+// wire and where the reader and each send actually ran, for every CPU placement, with one routine and with several. A
+// staged sender is only pinned to an explicit entry other than its reader's CPU, and inline sends stay on the reader.
+// The reader only reads again once its earlier reads are on the wire, so this also checks the sender never waits for
+// the reader. Pinning threads changes the process, so the test runs in a child.
+func TestFlowStagedTXAffinity(t *testing.T) {
+	if os.Getenv("NEBULA_FLOW_AFFINITY_TEST") != "1" {
+		cmd := exec.Command(os.Args[0], "-test.run=^TestFlowStagedTXAffinity$")
+		cmd.Env = append(os.Environ(), "NEBULA_FLOW_AFFINITY_TEST=1")
+		out, err := cmd.CombinedOutput()
+		require.NoError(t, err, "%s", out)
+		return
+	}
+	allowed, err := util.AllowedCPUs()
+	require.NoError(t, err)
+	for _, placement := range []string{"fallback", "automatic", "explicit", "explicit-single"} {
+		for _, pin := range []bool{false, true} {
+			for _, sc := range sendPaths() {
+				for _, routines := range []int{1, 2} {
+					t.Run(fmt.Sprintf("%s/pin=%t/%s/routines=%d", placement, pin, stagedName(sc), routines), func(t *testing.T) {
+						if pin && len(allowed) < 2 {
+							t.Skip("requires at least two allowed CPUs")
+						}
+						sender, receiver := newFlowTestPair(t, noise.CipherChaChaPoly, sc)
+						f := sender.f
+						f.pinThreads, f.routines = pin, routines
+						switch placement {
+						case "automatic":
+							f.defaultCpuAffinity = cpupick.Default(routines, 4242, nil)
+						case "explicit":
+							f.cpuAffinity = []int{allowed[len(allowed)-1], allowed[0]}
+						case "explicit-single":
+							f.cpuAffinity = allowed[len(allowed)-1:]
+						}
+						readerCPU, sendCPU := -1, -1
+						if pin {
+							readerCPU = f.threadCPU(0)
+							sendCPU = readerCPU
+							if sc.tx {
+								// Never the reader's CPU, and no automatic choice: unpinned unless an explicit entry fits
+								sendCPU = -1
+								if placement == "explicit" {
+									sendCPU = f.cpuAffinity[1]
+								}
+							}
+						}
+						smallCPU := sendCPU
+						conn := &affinityConn{}
+						f.writers = []udp.Conn{conn}
+						batches, want := flowTestBatches(3, 8, 1, 45)
+						batches[1] = batches[1][:1]
+						want = append(want[:405], want[720:]...)
+						q := &affinityQueue{flowBatchQueue: newFlowBatchQueue(batches), conn: conn}
+						f.listenIn(q, 0)
+						require.Equal(t, want, decryptWire(t, receiver, conn.captured()))
+						b := q.boundaries
+						requireFlowCPU(t, q.masks, readerCPU, allowed)
+						requireFlowCPU(t, conn.masks[:b[1]], sendCPU, allowed)
+						requireFlowCPU(t, conn.masks[b[1]:b[2]], smallCPU, allowed)
+						requireFlowCPU(t, conn.masks[b[2]:], sendCPU, allowed)
+					})
+				}
+			}
+		}
+	}
+}
+
+// sendPaths is inline, staged and, with the threshold, adaptive transmit.
+func sendPaths() []stagedConfig {
+	return []stagedConfig{{}, stagedOn(2)}
+}

+ 43 - 16
interface.go

@@ -57,14 +57,19 @@ type InterfaceConfig struct {
 
 	// CpuAffinity, when non-empty, names the CPUs each TUN reader goroutine
 	// should pin to. Queue i pins to CpuAffinity[i % len(CpuAffinity)] —
-	// shorter lists than `routines` cycle. Empty list keeps the default
-	// pin-to-(i % NumCPU) behavior. Only consulted when PinThreads is true.
+	// shorter lists than `routines` cycle. Staged sender i uses entry i+1,
+	// unless that is its reader's CPU. Empty uses the automatic order. Only
+	// consulted when PinThreads is true.
 	CpuAffinity []int
+	// defaultCpuAffinity is the automatic pin order for readers, kept apart
+	// from CpuAffinity because staged senders are only pinned to an explicit list.
+	defaultCpuAffinity []int
 	// PinThreads controls whether each TUN reader OS thread is pinned to a
 	// single CPU (via tun.pin_threads, default true). Pinning keeps each
 	// goroutine's sendmmsg on one XPS-selected NIC TX ring so per-flow
 	// packets stay ordered on the wire.
 	PinThreads bool
+	staged     stagedConfig
 
 	l *slog.Logger
 }
@@ -94,11 +99,13 @@ type Interface struct {
 	// should pin to. Queue i pins to cpuAffinity[i % len(cpuAffinity)].
 	// Empty falls back to the default pin-to-(allowed CPU) behavior.
 	// Only consulted when pinThreads is true.
-	cpuAffinity []int
+	cpuAffinity        []int
+	defaultCpuAffinity []int
 	// pinThreads controls whether listenIn pins each TUN reader OS thread to
 	// a CPU at all (tun.pin_threads, default true). When false, threads are
 	// left free to migrate as on stock nebula.
 	pinThreads   bool
+	staged       stagedConfig
 	relayManager *relayManager
 
 	tryPromoteEvery atomic.Uint32
@@ -247,7 +254,9 @@ func NewInterface(ctx context.Context, c *InterfaceConfig) (*Interface, error) {
 		connectionManager:     c.connectionManager,
 		conntrackCacheTimeout: c.ConntrackCacheTimeout,
 		cpuAffinity:           c.CpuAffinity,
+		defaultCpuAffinity:    c.defaultCpuAffinity,
 		pinThreads:            c.PinThreads,
+		staged:                c.staged,
 		udpTxErrLimiter:       udpTxErrLimiter,
 		udpRelayTxErrLimiter:  udpRelayTxErrLimiter,
 		txOversizeLimiter:     txOversizeLimiter,
@@ -433,32 +442,30 @@ func (f *Interface) listenOut(i int) {
 	f.l.Debug("underlay reader is done", "reader", i)
 }
 
-func (f *Interface) pinThisThread(i int) {
-	var cpu int
+// threadCPU is the CPU tun reader i pins to.
+func (f *Interface) threadCPU(i int) int {
 	if n := len(f.cpuAffinity); n > 0 {
 		// Explicit tun.cpu_affinity list wins; parseCpuAffinity already
 		// validated the entries against the allowed CPU set.
-		cpu = f.cpuAffinity[i%n]
+		return f.cpuAffinity[i%n]
+	} else if n := len(f.defaultCpuAffinity); n > 0 {
+		return f.defaultCpuAffinity[i%n]
 	} else if allowed, err := util.AllowedCPUs(); err == nil && len(allowed) > 0 {
 		// Default: spread queues across the CPUs we're actually allowed to
 		// run on. Under a cpuset/taskset mask these aren't 0..NumCPU-1, so
 		// i % NumCPU would pick unrunnable IDs and every pin would fail.
-		cpu = allowed[i%len(allowed)]
-	} else {
-		cpu = i % runtime.NumCPU()
+		return allowed[i%len(allowed)]
 	}
+	return i % runtime.NumCPU()
+}
+
+func (f *Interface) pinThreadToCPU(i, cpu int) {
 	if err := util.PinThreadToCPU(cpu); err != nil {
-		f.l.Warn("failed to pin tun reader to CPU", "queue", i, "cpu", cpu, "err", err)
+		f.l.Warn("failed to pin tun processing thread to CPU", "queue", i, "cpu", cpu, "err", err)
 	}
 }
 
 func (f *Interface) listenIn(queue tio.Queue, i int) {
-	// Pinning this thread (and goroutine) to a single CPU keeps every sendmmsg from this goroutine going through the
-	// same TX ring on the nic, so the wire sees per-flow order. Skip entirely when tun.pin_threads is false.
-	if f.pinThreads {
-		f.pinThisThread(i)
-	}
-
 	rejectBuf := make([]byte, mtu)
 	arenaSize := batch.SendBatchCap * (udp.MTU + 32)
 	sb := batch.NewSendBatch(f.writers[i], batch.SendBatchCap, arenaSize)
@@ -467,6 +474,26 @@ func (f *Interface) listenIn(queue tio.Queue, i int) {
 
 	conntrackCache := firewall.NewConntrackCacheTicker(f.ctx, f.l, f.conntrackCacheTimeout)
 
+	if sq, ok := queue.(handoffQueue); ok && f.staged.tx {
+		// Every stage on its own goroutine, see flow_tx.go
+		err := newStagedTx(f, sq, i).run(fwPacket, rejectBuf, conntrackCache)
+		if !f.closed.Load() && f.ctx.Err() == nil {
+			f.l.Error("Error while reading outbound packet, closing", "error", err, "reader", i)
+			f.onFatal(err)
+		}
+		f.l.Debug("overlay reader is done", "reader", i)
+		return
+	}
+	if f.staged.tx {
+		f.l.Info("tun queue cannot hand off read batches, transmit stays inline", "reader", i)
+	}
+
+	// Pinning this thread (and goroutine) to a single CPU keeps every sendmmsg from this goroutine going through the
+	// same TX ring on the nic, so the wire sees per-flow order. Skip entirely when tun.pin_threads is false.
+	if f.pinThreads {
+		f.pinThreadToCPU(i, f.threadCPU(i))
+	}
+
 	for {
 		pkts, err := queue.Read()
 		if err != nil {

+ 5 - 1
main.go

@@ -244,6 +244,8 @@ func Main(c *config.C, configTest bool, buildVersion string, l *slog.Logger, dev
 
 	pinThreads := c.GetBool("tun.pin_threads", true)
 	cpuAffinity := parseCpuAffinity(c, l, routines)
+	var defaultCpuAffinity []int
+	staged := parseStagedConfig(c, l)
 	if pinThreads && routines > 1 && len(cpuAffinity) == 0 && !configTest {
 		// The operator didn't choose pin CPUs, so pick a default set that
 		// prefers performance cores and doesn't stack co-located instances
@@ -270,7 +272,7 @@ func Main(c *config.C, configTest bool, buildVersion string, l *slog.Logger, dev
 			l.Warn("tun.pin_threads_key is invalid, using PID")
 		}
 
-		cpuAffinity = cpupick.Default(routines, key, l)
+		defaultCpuAffinity = cpupick.Default(routines, key, l)
 	}
 
 	ifConfig := &InterfaceConfig{
@@ -295,7 +297,9 @@ func Main(c *config.C, configTest bool, buildVersion string, l *slog.Logger, dev
 		punchy:                punchy,
 		ConntrackCacheTimeout: conntrackCacheTimeout,
 		CpuAffinity:           cpuAffinity,
+		defaultCpuAffinity:    defaultCpuAffinity,
 		PinThreads:            pinThreads,
+		staged:                staged,
 		l:                     l,
 	}
 

+ 22 - 0
main_test.go

@@ -5,12 +5,14 @@ import (
 	"net/netip"
 	"os"
 	"path/filepath"
+	"runtime"
 	"testing"
 	"time"
 
 	"github.com/slackhq/nebula/cert"
 	cert_test "github.com/slackhq/nebula/cert_test"
 	"github.com/slackhq/nebula/config"
+	"github.com/slackhq/nebula/noiseutil"
 	"github.com/slackhq/nebula/test"
 	"github.com/slackhq/nebula/util"
 	"github.com/stretchr/testify/assert"
@@ -123,3 +125,23 @@ firewall:
 	require.NoError(t, err)
 	require.Nil(t, ctrl, "a config test hands back nothing to stop, so it must stop itself")
 }
+
+func TestParseStagedConfig(t *testing.T) {
+	l := test.NewLogger()
+	staged := func(workers any) stagedConfig {
+		c := config.NewC(l)
+		if workers != nil {
+			c.Settings["tun"] = map[string]any{"staged": map[string]any{"workers": workers}}
+		}
+		return parseStagedConfig(c, l)
+	}
+	for _, workers := range []any{nil, -1, 0} {
+		require.Equal(t, stagedConfig{}, staged(workers))
+	}
+	if runtime.GOOS != "linux" {
+		require.Equal(t, stagedConfig{}, staged(2))
+		return
+	}
+	want := stagedConfig{workers: 2, tx: !noiseutil.EncryptLockNeeded}
+	require.Equal(t, want, staged(2))
+}

+ 9 - 4
overlay/batch/tx_batch.go

@@ -43,6 +43,14 @@ func (b *SendBatch) Commit(pkt []byte, dst netip.AddrPort) {
 	b.dsts = append(b.dsts, dst)
 }
 
+// Reset drops every queued packet without sending it.
+func (b *SendBatch) Reset() {
+	clear(b.bufs)
+	b.bufs = b.bufs[:0]
+	b.dsts = b.dsts[:0]
+	b.arena.Reset()
+}
+
 // Flush writes every queued packet and reports how many actually went out. A short count means some destinations
 // were undeliverable; the batch is drained either way.
 func (b *SendBatch) Flush() (int, error) {
@@ -51,9 +59,6 @@ func (b *SendBatch) Flush() (int, error) {
 	if len(b.bufs) > 0 {
 		written, err = b.out.WriteBatch(b.bufs, b.dsts)
 	}
-	clear(b.bufs)
-	b.bufs = b.bufs[:0]
-	b.dsts = b.dsts[:0]
-	b.arena.Reset()
+	b.Reset()
 	return written, err
 }

+ 15 - 0
overlay/tio/tio_gso_linux.go

@@ -219,6 +219,21 @@ func (r *Offload) Read() ([]Packet, error) {
 	return r.pending, nil
 }
 
+// SwapRxBuf installs buf as the receive buffer for the following Reads and returns the one backing the packets
+// of the most recent Read, so the caller can keep them past the next Read. A nil buf allocates a replacement of
+// the required size. Single-reader only, like Read.
+func (r *Offload) SwapRxBuf(buf []byte) []byte {
+	if buf == nil {
+		buf = make([]byte, tunRxBufCap)
+	}
+	if len(buf) < tunRxBufCap {
+		panic(fmt.Sprintf("tio: SwapRxBuf buffer is %d bytes, need %d", len(buf), tunRxBufCap))
+	}
+	old := r.rxBuf
+	r.rxBuf = buf
+	return old
+}
+
 // logDroppedRead reports a tun packet dropped for a bad/unsupported virtio
 // header. Debug-gated so the happy path never pays for attribute assembly.
 func (r *Offload) logDroppedRead(err error) {

+ 19 - 0
overlay/tio/tun_linux_offload_test.go

@@ -911,6 +911,25 @@ func TestOffloadWriteZeroLength(t *testing.T) {
 	}
 }
 
+// TestOffloadSwapRxBuf: SwapRxBuf hands over the buffer the last Read filled and reads into the one it is given.
+func TestOffloadSwapRxBuf(t *testing.T) {
+	o := &Offload{rxBuf: make([]byte, tunRxBufCap)}
+	read := o.rxBuf
+	owned := o.SwapRxBuf(nil)
+	if &owned[0] != &read[0] || len(o.rxBuf) != tunRxBufCap || &o.rxBuf[0] == &read[0] {
+		t.Fatal("SwapRxBuf(nil) must hand over the last read buffer and allocate a new one")
+	}
+	if back := o.SwapRxBuf(owned); &back[0] == &owned[0] || &o.rxBuf[0] != &owned[0] {
+		t.Fatal("SwapRxBuf must install the buffer it is given")
+	}
+	defer func() {
+		if recover() == nil {
+			t.Fatal("SwapRxBuf must reject a short buffer")
+		}
+	}()
+	o.SwapRxBuf(make([]byte, 10))
+}
+
 // TestWriteGSOSuperpacketGeometry decodes the vnet header the kernel would see for a multi-segment write:
 // the GSO type must match the proto and IP version
 // gso_size must be the per-segment size (the kernel rejects a superpacket with gso_size == 0),