the definitive tun offloads branch (#1704)

interface.go

@@ -4,9 +4,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"log/slog"
 	"net/netip"
+	"runtime"
 	"slices"
 	"sync"
 	"sync/atomic"
@@ -14,12 +14,15 @@ import (
 
 	"github.com/gaissmai/bart"
 	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/util"
 
 	"github.com/slackhq/nebula/cert"
 	"github.com/slackhq/nebula/config"
 	"github.com/slackhq/nebula/firewall"
 	"github.com/slackhq/nebula/header"
 	"github.com/slackhq/nebula/overlay"
+	"github.com/slackhq/nebula/overlay/batch"
+	"github.com/slackhq/nebula/overlay/tio"
 	"github.com/slackhq/nebula/udp"
 )
 
@@ -49,7 +52,19 @@ type InterfaceConfig struct {
 	reQueryWait     time.Duration
 
 	ConntrackCacheTimeout time.Duration
-	l                     *slog.Logger
+
+	// 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.
+	CpuAffinity []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
+
+	l *slog.Logger
 }
 
 type Interface struct {
@@ -73,7 +88,16 @@ type Interface struct {
 	routines              int
 	disconnectInvalid     atomic.Bool
 	closed                atomic.Bool
-	relayManager          *relayManager
+	// cpuAffinity, when non-empty, names the CPUs each TUN reader goroutine
+	// 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
+	// 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
+	relayManager *relayManager
 
 	tryPromoteEvery atomic.Uint32
 	reQueryEvery    atomic.Uint32
@@ -90,8 +114,14 @@ type Interface struct {
 
 	ctx     context.Context
 	writers []udp.Conn
-	readers []io.ReadWriteCloser
-	wg      sync.WaitGroup
+	queues  []tio.Queue
+	// batchers is one per tun queue, wrapping queues[i]. readOutsidePackets
+	// commits plaintext into the batcher; the plaintext is decrypted
+	// in place inside the UDP receive buffers, so listenOut must call Flush
+	// at the end of each UDP recvmmsg batch, before those buffers are
+	// reused (every udp.Conn ListenOut guarantees that ordering).
+	batchers []*batch.MultiCoalescer
+	wg       sync.WaitGroup
 
 	// fatalErr holds the first unexpected reader error that caused shutdown.
 	// nil means "no fatal error" (yet)
@@ -102,18 +132,13 @@ type Interface struct {
 	metricHandshakes    metrics.Histogram
 	messageMetrics      *MessageMetrics
 	cachedPacketMetrics *cachedPacketMetrics
+	metricTxDropped     metrics.Counter
 
 	l *slog.Logger
 }
 
 type EncWriter interface {
-	SendVia(via *HostInfo,
-		relay *Relay,
-		ad,
-		nb,
-		out []byte,
-		nocopy bool,
-	)
+	SendVia(via *HostInfo, relay *Relay, ad, nb, out []byte, nocopy bool, q int)
 	SendMessageToVpnAddr(t header.MessageType, st header.MessageSubType, vpnAddr netip.Addr, p, nb, out []byte)
 	SendMessageToHostInfo(t header.MessageType, st header.MessageSubType, hostinfo *HostInfo, p, nb, out []byte)
 	Handshake(vpnAddr netip.Addr)
@@ -172,6 +197,10 @@ func NewInterface(ctx context.Context, c *InterfaceConfig) (*Interface, error) {
 		return nil, errors.New("no connection manager")
 	}
 
+	if c.routines <= 1 {
+		c.PinThreads = false //pinning is not useful unless there's more than one tun reader
+	}
+
 	cs := c.pki.getCertState()
 	ifce := &Interface{
 		ctx:                   ctx,
@@ -189,7 +218,7 @@ func NewInterface(ctx context.Context, c *InterfaceConfig) (*Interface, error) {
 		routines:              c.routines,
 		version:               c.version,
 		writers:               make([]udp.Conn, c.routines),
-		readers:               make([]io.ReadWriteCloser, c.routines),
+		batchers:              make([]*batch.MultiCoalescer, c.routines),
 		myVpnNetworks:         cs.myVpnNetworks,
 		myVpnNetworksTable:    cs.myVpnNetworksTable,
 		myVpnAddrs:            cs.myVpnAddrs,
@@ -198,8 +227,11 @@ func NewInterface(ctx context.Context, c *InterfaceConfig) (*Interface, error) {
 		relayManager:          c.relayManager,
 		connectionManager:     c.connectionManager,
 		conntrackCacheTimeout: c.ConntrackCacheTimeout,
+		cpuAffinity:           c.CpuAffinity,
+		pinThreads:            c.PinThreads,
 
 		metricHandshakes: metrics.GetOrRegisterHistogram("handshakes", nil, metrics.NewExpDecaySample(1028, 0.015)),
+		metricTxDropped:  metrics.GetOrRegisterCounter("udp.tx.dropped", nil),
 		messageMetrics:   c.MessageMetrics,
 		cachedPacketMetrics: &cachedPacketMetrics{
 			sent:    metrics.GetOrRegisterCounter("hostinfo.cached_packets.sent", nil),
@@ -240,25 +272,36 @@ func (f *Interface) activate() error {
 		"boringcrypto", boringEnabled(),
 	)
 
-	if f.routines > 1 {
-		if !f.inside.SupportsMultiqueue() || !f.outside.SupportsMultipleReaders() {
-			f.routines = 1
-			f.l.Warn("routines is not supported on this platform, falling back to a single routine")
-		}
+	if f.routines > 1 && !f.outside.SupportsMultipleReaders() {
+		f.routines = 1
+		f.l.Warn("multiple udp readers are not supported on this platform, falling back to a single routine")
 	}
 
+	// Prepare the tun queues. A device that can't open that many hands back
+	// fewer (a single queue on platforms without multiqueue support) and we
+	// size the reader routines to what we actually got.
+	queues, err := f.inside.Queues(f.routines)
+	if err != nil {
+		return err
+	}
+	if len(queues) < f.routines {
+		// TODO: this clamp is only safe because it is unreachable when the
+		// udp side has multiple readers (linux Queues opens exactly n or
+		// errors; every other platform already clamped routines to 1 above).
+		// If a platform ever returns fewer queues than routines with
+		// SO_REUSEPORT sockets already bound, the surplus sockets get no
+		// listenOut and the kernel blackholes every flow it hashes to them —
+		// fail loudly or close the extra sockets instead.
+		f.l.Warn("tun multiqueue is not supported on this platform, falling back to fewer routines",
+			"requested", f.routines, "opened", len(queues))
+		f.routines = len(queues)
+	}
+	f.queues = queues
+
 	metrics.GetOrRegisterGauge("routines", nil).Update(int64(f.routines))
 
-	// Prepare n tun queues
-	var reader io.ReadWriteCloser = f.inside
-	for i := 0; i < f.routines; i++ {
-		if i > 0 {
-			reader, err = f.inside.NewMultiQueueReader()
-			if err != nil {
-				return err
-			}
-		}
-		f.readers[i] = reader
+	for i := range f.queues {
+		f.batchers[i] = batch.NewMultiCoalescer(f.queues[i], f.l)
 	}
 
 	// On error the caller owns the cleanup, Control.Start cancels the service context
@@ -281,7 +324,7 @@ func (f *Interface) run() {
 	// Launch n queues to read packets from tun dev
 	for i := 0; i < f.routines; i++ {
 		f.wg.Go(func() {
-			f.listenIn(f.readers[i], i)
+			f.listenIn(f.queues[i], i)
 		})
 	}
 
@@ -306,6 +349,31 @@ func (f *Interface) onFatal(err error) {
 	}
 }
 
+type rxContext struct {
+	q       int
+	scratch []byte
+	// nb is a re-usable nonce buffer for decrypt calls to use
+	nb           []byte
+	h            *header.H
+	fwPacket     *firewall.ParsedPacket
+	hostmapCache map[uint32]*HostInfo
+	lhh          *LightHouseHandler
+	ctCache      *firewall.ConntrackCacheTicker
+}
+
+func newRxContext(f *Interface, q int) *rxContext {
+	return &rxContext{
+		q:            q,
+		scratch:      make([]byte, mtu),
+		nb:           make([]byte, 12, 12),
+		h:            &header.H{},
+		fwPacket:     &firewall.ParsedPacket{},
+		hostmapCache: map[uint32]*HostInfo{},
+		lhh:          f.lightHouse.NewRequestHandler(),
+		ctCache:      firewall.NewConntrackCacheTicker(f.ctx, f.l, f.conntrackCacheTimeout),
+	}
+}
+
 func (f *Interface) listenOut(i int) {
 	var li udp.Conn
 	if i > 0 {
@@ -314,16 +382,20 @@ func (f *Interface) listenOut(i int) {
 		li = f.outside
 	}
 
-	ctCache := firewall.NewConntrackCacheTicker(f.ctx, f.l, f.conntrackCacheTimeout)
-	lhh := f.lightHouse.NewRequestHandler()
-	plaintext := make([]byte, udp.MTU)
-	h := &header.H{}
-	fwPacket := &firewall.Packet{}
-	nb := make([]byte, 12, 12)
+	rxc := newRxContext(f, i)
 
-	err := li.ListenOut(func(fromUdpAddr netip.AddrPort, payload []byte) {
-		f.readOutsidePackets(ViaSender{UdpAddr: fromUdpAddr}, plaintext[:0], payload, h, fwPacket, lhh, nb, i, ctCache.Get())
-	})
+	listener := func(fromUdpAddr netip.AddrPort, payload []byte) {
+		f.readOutsidePackets(ViaSender{UdpAddr: fromUdpAddr}, payload, rxc)
+	}
+
+	flusher := func() {
+		if err := f.batchers[i].Flush(); err != nil {
+			f.l.Error("Failed to flush tun coalescer", "error", err)
+		}
+		clear(rxc.hostmapCache)
+	}
+
+	err := li.ListenOut(listener, flusher)
 
 	// An error after teardown began is shutdown noise, the closed flag covers resources
 	// Close releases itself and the cancelled ctx covers ones torn down by their owners
@@ -336,16 +408,42 @@ func (f *Interface) listenOut(i int) {
 	f.l.Debug("underlay reader is done", "reader", i)
 }
 
-func (f *Interface) listenIn(reader io.ReadWriteCloser, i int) {
-	packet := make([]byte, mtu)
-	out := make([]byte, mtu)
-	fwPacket := &firewall.Packet{}
+func (f *Interface) pinThisThread(i int) {
+	var cpu 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]
+	} 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()
+	}
+	if err := util.PinThreadToCPU(cpu); err != nil {
+		f.l.Warn("failed to pin tun reader 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)
+	fwPacket := &firewall.ParsedPacket{}
 	nb := make([]byte, 12, 12)
 
 	conntrackCache := firewall.NewConntrackCacheTicker(f.ctx, f.l, f.conntrackCacheTimeout)
 
 	for {
-		n, err := reader.Read(packet)
+		pkts, err := queue.Read()
 		if err != nil {
 			// Same shutdown noise handling as listenOut
 			if !f.closed.Load() && f.ctx.Err() == nil {
@@ -355,12 +453,35 @@ func (f *Interface) listenIn(reader io.ReadWriteCloser, i int) {
 			break
 		}
 
-		f.consumeInsidePacket(packet[:n], fwPacket, nb, out, i, conntrackCache.Get())
+		for _, pkt := range pkts {
+			f.consumeInsidePacket(pkt, fwPacket, nb, sb, rejectBuf, i, conntrackCache.Get())
+			// Flush incrementally once a full sendmmsg batch has
+			// accumulated so the first packets of a deep read drain
+			// hit the wire while the rest are still being encrypted.
+			if sb.Len() >= batch.SendBatchCap {
+				f.flushSendBatch(sb, i)
+			}
+		}
+		f.flushSendBatch(sb, i)
 	}
 
 	f.l.Debug("overlay reader is done", "reader", i)
 }
 
+// flushSendBatch drains sb to the underlay and accounts for anything it could not deliver. A shortfall means
+// specific destinations were undeliverable (a stale remote, a reject rule), which the backend logs per peer at
+// debug; here it is only a counter, so one unreachable peer cannot spam a log line per batch.
+func (f *Interface) flushSendBatch(sb *batch.SendBatch, q int) {
+	queued := sb.Len()
+	written, err := sb.Flush()
+	if err != nil {
+		f.l.Error("Failed to write outgoing batch", "error", err, "writer", q)
+	}
+	if dropped := queued - written; dropped > 0 {
+		f.metricTxDropped.Inc(int64(dropped))
+	}
+}
+
 func (f *Interface) RegisterConfigChangeCallbacks(c *config.C) {
 	c.RegisterReloadCallback(f.reloadFirewall)
 	c.RegisterReloadCallback(f.reloadSendRecvError)