diff --git a/flow.go b/flow.go index 4a2e5b54..6d0b1c9b 100644 --- a/flow.go +++ b/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() } diff --git a/flow_rx.go b/flow_rx.go index be36e7d2..88272e42 100644 --- a/flow_rx.go +++ b/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() } diff --git a/flow_test.go b/flow_test.go index 2553bb01..f4d07bb7 100644 --- a/flow_test.go +++ b/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) diff --git a/udp/udp_linux.go b/udp/udp_linux.go index 6eecdc8e..7e72742c 100644 --- a/udp/udp_linux.go +++ b/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 } } diff --git a/udp/udp_tester.go b/udp/udp_tester.go index 557d4584..0c436e07 100644 --- a/udp/udp_tester.go +++ b/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() } }