Simplify staged receive workers and buffer handoff

flow.go

@@ -3,7 +3,6 @@ package nebula
 import (
 	"log/slog"
 	"runtime"
-	"sync"
 
 	"github.com/slackhq/nebula/config"
 )
@@ -39,28 +38,3 @@ func parseStagedConfig(c *config.C, l *slog.Logger) stagedConfig {
 	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) *flowPool {
-	p := &flowPool{jobs: make(chan flowJob, queueLen)}
-	for range workers {
-		p.wg.Go(func() {
-			nb := make([]byte, 12)
-			for job := range p.jobs {
-				job.run(nb)
-			}
-		})
-	}
-	return p
-}
-
-func (p *flowPool) close() { close(p.jobs); p.wg.Wait() }

flow_rx.go

@@ -3,16 +3,18 @@ package nebula
 import (
 	"log/slog"
 	"net/netip"
+	"sync"
 
 	"github.com/rcrowley/go-metrics"
 	"github.com/slackhq/nebula/header"
 	"github.com/slackhq/nebula/udp"
 )
 
-// handoffConn is a udp conn that can hand each received batch over. The callback returns replacement buffers with
-// the same count and at least the same lengths.
+// handoffConn is a udp conn that can hand each received batch over. The callback returns replacement buffers in the
+// conn's HandoffShape.
 type handoffConn interface {
 	ListenOutHandoff(udp.EncReader, func([][]byte) [][]byte) error
+	HandoffShape() (n, size int)
 }
 
 const stagedRxSlots = 6
@@ -31,16 +33,10 @@ type rxItem struct {
 
 // rxChunk is a contiguous run of a batch's items for one worker to decrypt.
 type rxChunk struct {
-	l     *slog.Logger
 	items []rxItem
 	done  chan struct{}
 }
 
-func (c *rxChunk) run(nb []byte) {
-	decryptRxItems(c.l, c.items, nb)
-	c.done <- struct{}{}
-}
-
 // decryptRxItems checks the data packets among items against the replay window and decrypts the ones it accepts in
 // place, as Decrypt does. Recording the counters is left to readOutsidePackets on the writer, which also handles
 // the packets without a tunnel.
@@ -81,9 +77,11 @@ func (b *rxBatch) reset() {
 
 // stagedRx is the staged receive pipeline of one udp listener. deliver and handoff run on the listener goroutine.
 type stagedRx struct {
-	f    *Interface
-	q    int
-	pool *flowPool
+	f *Interface
+	q int
+	// jobs feeds the decrypt workers. The listener is the only producer and closes it once the writer has drained.
+	jobs    chan *rxChunk
+	workers sync.WaitGroup
 	// h and hostmapCache are the listener's own. wrxc is the commit context: the writer's, or the listener's while
 	// nothing is in flight.
 	h            *header.H
@@ -114,13 +112,15 @@ func newStagedRx(f *Interface, q int, wrxc *rxContext) *stagedRx {
 		metricBatches: metrics.GetOrRegisterCounter("flow.rx.batches", nil),
 		metricInline:  metrics.GetOrRegisterCounter("flow.rx.inline", nil),
 	}
-	workers := f.staged.workers
-	p.pool = newFlowPool(workers, stagedRxSlots*2*workers)
+	n := f.staged.workers
+	p.jobs = make(chan *rxChunk, stagedRxSlots*2*n)
+	for range n {
+		p.workers.Go(p.decrypt)
+	}
 	for range stagedRxSlots {
 		// 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)}
+		b := &rxBatch{items: make([]rxItem, 0, 256), chunks: make([]rxChunk, 2*n)}
 		for i := range b.chunks {
-			b.chunks[i].l = f.l
 			b.chunks[i].done = make(chan struct{}, 1)
 		}
 		p.free = append(p.free, b)
@@ -131,9 +131,26 @@ func newStagedRx(f *Interface, q int, wrxc *rxContext) *stagedRx {
 	return p
 }
 
+// decrypt is a worker: it decrypts the chunks it is handed and signals each one done.
+func (p *stagedRx) decrypt() {
+	nb := make([]byte, 12)
+	for c := range p.jobs {
+		decryptRxItems(p.f.l, c.items, nb)
+		c.done <- struct{}{}
+	}
+}
+
 // listen runs the conn's receive loop through the pipeline until it fails, then drains the pipeline. The conn hands
 // its buffers over, so a batch moves between stages without a copy.
 func (p *stagedRx) listen(li handoffConn) error {
+	// The first batch reads into the conn's own buffers; every other one brings its own
+	n, size := li.HandoffShape()
+	for _, b := range p.free {
+		b.bufs = make([][]byte, n)
+		for i := range b.bufs {
+			b.bufs[i] = make([]byte, size)
+		}
+	}
 	defer p.close()
 	return li.ListenOutHandoff(p.deliver, p.handoff)
 }
@@ -175,14 +192,7 @@ func (p *stagedRx) classify(it *rxItem) {
 func (p *stagedRx) handoff(bufs [][]byte) [][]byte {
 	p.cur.bufs = bufs
 	p.endBatch()
-	b := p.cur
-	if len(b.bufs) != len(bufs) || len(b.bufs[0]) < len(bufs[0]) {
-		b.bufs = make([][]byte, len(bufs))
-		for i := range b.bufs {
-			b.bufs[i] = make([]byte, len(bufs[0]))
-		}
-	}
-	return b.bufs
+	return p.cur.bufs
 }
 
 // endBatch hands the batch to the workers and the ordered writer, then takes a free batch.
@@ -216,7 +226,7 @@ func (p *stagedRx) endBatch() {
 	}
 	// The jobs go out before the batch does: once the writer has it, it may finish and reset it at any moment
 	for i := range b.nChunks {
-		p.pool.jobs <- &b.chunks[i]
+		p.jobs <- &b.chunks[i]
 	}
 	p.inFlight++
 	p.ordered <- b
@@ -301,5 +311,6 @@ func (p *stagedRx) close() {
 	}
 	close(p.ordered)
 	<-p.writerDone
-	p.pool.close()
+	close(p.jobs)
+	p.workers.Wait()
 }

flow_test.go

@@ -403,11 +403,16 @@ func (c *rxConn) ListenOut(r udp.EncReader, flush func()) error {
 	})
 }
 
-func (c *rxConn) ListenOutHandoff(r udp.EncReader, handoff func(bufs [][]byte) [][]byte) error {
-	n := c.batchLen
-	for _, size := range c.sizes {
-		n = max(n, size)
+func (c *rxConn) HandoffShape() (n, size int) {
+	n = c.batchLen
+	for _, s := range c.sizes {
+		n = max(n, s)
 	}
+	return n, udp.MTU
+}
+
+func (c *rxConn) ListenOutHandoff(r udp.EncReader, handoff func(bufs [][]byte) [][]byte) error {
+	n, _ := c.HandoffShape()
 	bufs := make([][]byte, n)
 	for i := range bufs {
 		bufs[i] = make([]byte, udp.MTU)

udp/udp_linux.go

@@ -246,22 +246,27 @@ func getFrom(names [][]byte, i int, isV4 bool) netip.AddrPort {
 }
 
 func (u *StdConn) ListenOut(r EncReader, flush func()) error {
-	return u.listenOut(r, flush, nil)
+	return u.ListenOutHandoff(r, func(bufs [][]byte) [][]byte {
+		flush()
+		return bufs
+	})
 }
 
-// ListenOutHandoff is ListenOut for a reader that keeps the packets: after each batch, handoff takes the batch's
-// buffers and returns replacements, of the same count and at least the same size, to receive the next one into.
+// HandoffShape is the count and size of the buffers ListenOutHandoff hands over after each batch.
+func (u *StdConn) HandoffShape() (n, size int) {
+	if u.groSupported {
+		return u.batch, udpGROBufferSize
+	}
+	return u.batch, MTU
+}
+
+// ListenOutHandoff is the receive loop. After each batch, handoff takes the batch's buffers and returns the ones to
+// receive the next batch into, in HandoffShape: the same ones for a reader that is done with the packets, or
+// replacements for one that keeps them.
 func (u *StdConn) ListenOutHandoff(r EncReader, handoff func(bufs [][]byte) [][]byte) error {
-	return u.listenOut(r, nil, handoff)
-}
-
-// listenOut is the receive loop. After each batch it calls flush, or with handoff set, hands the batch's buffers
-// over and points the recvmmsg iovecs at the ones it gets back.
-func (u *StdConn) listenOut(r EncReader, flush func(), handoff func(bufs [][]byte) [][]byte) error {
-	bufSize := MTU
+	_, bufSize := u.HandoffShape()
 	cmsgSpace := 0
 	if u.groSupported {
-		bufSize = udpGROBufferSize
 		cmsgSpace = unix.CmsgSpace(udpGROCmsgPayload)
 	}
 	msgs, buffers, names, _ := prepareRawMessages(u.batch, bufSize, cmsgSpace)
@@ -295,21 +300,20 @@ func (u *StdConn) listenOut(r EncReader, flush func(), handoff func(bufs [][]byt
 			deliverSegments(r, from, payload, segSize)
 		}
 
-		if handoff == nil {
-			flush()
-			continue
-		}
 		next := handoff(buffers)
 		if len(next) != len(buffers) {
 			panic(fmt.Sprintf("udp: handoff returned %d buffers, need %d", len(next), len(buffers)))
 		}
-		for i := range next {
-			if len(next[i]) < bufSize {
-				panic(fmt.Sprintf("udp: handoff buffer is %d bytes, need %d", len(next[i]), bufSize))
+		if &next[0] != &buffers[0] {
+			// Replacements, point the recvmmsg iovecs at them
+			for i := range next {
+				if len(next[i]) < bufSize {
+					panic(fmt.Sprintf("udp: handoff buffer is %d bytes, need %d", len(next[i]), bufSize))
+				}
+				msgs[i].Hdr.Iov.Base = &next[i][0]
 			}
-			msgs[i].Hdr.Iov.Base = &next[i][0]
+			buffers = next
 		}
-		buffers = next
 	}
 }
 

udp/udp_tester.go

@@ -185,22 +185,29 @@ func (u *TesterConn) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, err
 func (u *TesterConn) ListenOut(r EncReader, flush func()) error {
 	// The batcher borrows plaintext decrypted in place inside p.Data
 	// until Flush, so the packet must stay alive across flush()
-	return u.listenOut(r, func(p *Packet) { flush() })
+	return u.ListenOutHandoff(r, func(bufs [][]byte) [][]byte {
+		flush()
+		return bufs
+	})
 }
 
-// ListenOutHandoff hands each packet's buffer over, the replacement goes back to the pool.
+// HandoffShape is one buffer per handoff; the pool resizes the replacements it gets back as it needs.
+func (u *TesterConn) HandoffShape() (n, size int) { return 1, MTU }
+
+// ListenOutHandoff hands each packet's buffer over as a batch of one, the buffer it gets back goes to the pool. Like
+// StdConn it carries on with the slice handoff returns, since a reader that keeps the packets keeps the one it was
+// given.
 func (u *TesterConn) ListenOutHandoff(r EncReader, handoff func(bufs [][]byte) [][]byte) error {
-	return u.listenOut(r, func(p *Packet) { p.Data = handoff([][]byte{p.Data})[0] })
-}
-
-func (u *TesterConn) listenOut(r EncReader, finish func(*Packet)) error {
+	bufs := make([][]byte, 1)
 	for {
 		select {
 		case <-u.done:
 			return os.ErrClosed
 		case p := <-u.RxPackets:
 			r(p.From, p.Data[:len(p.Data):len(p.Data)])
-			finish(p)
+			bufs[0] = p.Data
+			bufs = handoff(bufs)
+			p.Data = bufs[0]
 			p.Release()
 		}
 	}