udp: Add opt-in ordered staged receive processing

flow.go

@@ -0,0 +1,80 @@
+package nebula
+
+import (
+	"fmt"
+	"log/slog"
+	"runtime"
+	"sync"
+	"time"
+
+	"github.com/rcrowley/go-metrics"
+	"github.com/slackhq/nebula/config"
+)
+
+// Staged receive (tun.staged.workers) moves decryption off each udp listener onto worker goroutines, behind which an
+// ordered writer commits every batch in arrival order. It reuses the inline packet handlers: staging changes ownership
+// and scheduling, not policy. Each routine has its own workers and writer, so order is per routine. Full queues block
+// the listener. A read error drains the batches in flight, while shutdown waits for the workers and then drops what
+// is left. Each routine keeps six batches; at listen.batch=64 with GRO their buffers take about 24 MiB per routine.
+//
+// Metrics are registered only for active pipelines. flow.rx.batches counts batches handed to the writer and
+// flow.rx.inline those an idle listener handled itself; their rates show whether staging is used.
+// flow.stage.<stage>.q<queue>.busy_ns locates the limiting stage: rx_read, rx_decrypt or rx_write. 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. Small batches sent to the writer include decryption in
+// writer time.
+
+// stagedConfig is tun.staged. workers is per routine; zero keeps receive inline.
+type stagedConfig struct {
+	workers int
+	// inlineBelow is the packet count below which an idle listener skips the worker handoff
+	inlineBelow int
+}
+
+// parseStagedConfig enables staging on Linux only; 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{}
+	}
+	// 256 is the measured crossover
+	sc.inlineBelow = 256
+	l.Info("Staged receive processing enabled", "workers", sc.workers)
+	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 listener's crypto jobs. The listener 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)
+}