From 06035252b2b8ba2d4e158ec4e8ce89a495ea0fbc Mon Sep 17 00:00:00 2001 From: JackDoan Date: Thu, 15 Oct 2026 17:02:42 -0500 Subject: [PATCH] udp: Fold buffer handoff into ListenOut --- control_lifecycle_test.go | 8 ++++---- flow_rx.go | 8 ++++---- flow_test.go | 11 ++--------- interface.go | 5 ++++- udp/conn.go | 13 +++++++------ udp/udp_darwin.go | 8 +++++--- udp/udp_generic.go | 8 +++++--- udp/udp_linux.go | 14 ++------------ udp/udp_linux_offloads_test.go | 2 +- udp/udp_linux_test.go | 2 +- udp/udp_rio_windows.go | 8 +++++--- udp/udp_tester.go | 18 +++++------------- 12 files changed, 45 insertions(+), 60 deletions(-) diff --git a/control_lifecycle_test.go b/control_lifecycle_test.go index 4bc75273..518063de 100644 --- a/control_lifecycle_test.go +++ b/control_lifecycle_test.go @@ -144,10 +144,10 @@ type fakeConn struct { rebinds int } -func (c *fakeConn) Rebind() error { c.rebinds++; return nil } -func (c *fakeConn) LocalAddr() (netip.AddrPort, error) { return netip.AddrPort{}, nil } -func (c *fakeConn) ListenOut(_ udp.EncReader, _ func()) error { return nil } -func (c *fakeConn) WriteTo(_ []byte, _ netip.AddrPort) error { return nil } +func (c *fakeConn) Rebind() error { c.rebinds++; return nil } +func (c *fakeConn) LocalAddr() (netip.AddrPort, error) { return netip.AddrPort{}, nil } +func (c *fakeConn) ListenOut(_ udp.EncReader, _ func([][]byte) [][]byte) error { return nil } +func (c *fakeConn) WriteTo(_ []byte, _ netip.AddrPort) error { return nil } func (c *fakeConn) WriteBatch(bufs [][]byte, _ []netip.AddrPort) (int, error) { return len(bufs), nil } diff --git a/flow_rx.go b/flow_rx.go index 88272e42..8ac16a0b 100644 --- a/flow_rx.go +++ b/flow_rx.go @@ -10,10 +10,10 @@ import ( "github.com/slackhq/nebula/udp" ) -// handoffConn is a udp conn that can hand each received batch over. The callback returns replacement buffers in the -// conn's HandoffShape. +// handoffConn is a udp conn that receives in batches of a fixed shape, so staged receive can have replacement buffers +// ready for every batch it keeps. type handoffConn interface { - ListenOutHandoff(udp.EncReader, func([][]byte) [][]byte) error + ListenOut(udp.EncReader, func([][]byte) [][]byte) error HandoffShape() (n, size int) } @@ -152,7 +152,7 @@ func (p *stagedRx) listen(li handoffConn) error { } } defer p.close() - return li.ListenOutHandoff(p.deliver, p.handoff) + return li.ListenOut(p.deliver, p.handoff) } // deliver is the listener's per-packet work: classify the packet and add it to the current batch. Nothing here diff --git a/flow_test.go b/flow_test.go index f4d07bb7..5877da2b 100644 --- a/flow_test.go +++ b/flow_test.go @@ -396,13 +396,6 @@ type rxConn struct { swaps int } -func (c *rxConn) ListenOut(r udp.EncReader, flush func()) error { - return c.ListenOutHandoff(r, func(bufs [][]byte) [][]byte { - flush() - return bufs - }) -} - func (c *rxConn) HandoffShape() (n, size int) { n = c.batchLen for _, s := range c.sizes { @@ -411,7 +404,7 @@ func (c *rxConn) HandoffShape() (n, size int) { return n, udp.MTU } -func (c *rxConn) ListenOutHandoff(r udp.EncReader, handoff func(bufs [][]byte) [][]byte) error { +func (c *rxConn) ListenOut(r udp.EncReader, handoff func(bufs [][]byte) [][]byte) error { n, _ := c.HandoffShape() bufs := make([][]byte, n) for i := range bufs { @@ -447,7 +440,7 @@ func (c *rxConn) ListenOutHandoff(r udp.EncReader, handoff func(bufs [][]byte) [ return os.ErrClosed } -// plainConn hides ListenOutHandoff, so it cannot hand off its batches. +// plainConn hides HandoffShape, so receive stays inline. type plainConn struct{ udp.Conn } var rxBatches = metrics.GetOrRegisterCounter("flow.rx.batches", nil) diff --git a/interface.go b/interface.go index 3e004afb..8e1e90ab 100644 --- a/interface.go +++ b/interface.go @@ -418,7 +418,10 @@ func (f *Interface) listenOut(i int) { f.readOutsidePackets(ViaSender{UdpAddr: fromUdpAddr}, payload, rxc) } - flusher := func() { f.flushOutside(rxc) } + flusher := func(bufs [][]byte) [][]byte { + f.flushOutside(rxc) + return bufs + } var err error if hc, ok := li.(handoffConn); ok && f.staged.workers > 0 { diff --git a/udp/conn.go b/udp/conn.go index 6d20cbbf..78c4ddc3 100644 --- a/udp/conn.go +++ b/udp/conn.go @@ -29,11 +29,12 @@ type Settings struct { type Conn interface { Rebind() error LocalAddr() (netip.AddrPort, error) - // ListenOut invokes r for each received packet. - // On batch-capable backends (recvmmsg), flush is called after each batch is fully delivered. - // Callers use it to flush per-batch accumulators such as TUN write coalescers. - // Single-packet backends call flush after each packet. flush must not be nil. - ListenOut(r EncReader, flush func()) error + // ListenOut invokes r for each received packet. The packets alias the receive buffers. + // After each batch is fully delivered (after each packet on single-packet backends), handoff takes the batch's + // buffers and returns the ones to receive the next batch into: the same ones once the caller is done with the + // packets, after flushing per-batch accumulators such as TUN write coalescers, or replacements of the same count + // and size to keep them. handoff must not be nil. + ListenOut(r EncReader, handoff func(bufs [][]byte) [][]byte) error WriteTo(b []byte, addr netip.AddrPort) error // WriteBatch sends a contiguous batch of packets, each with its own // destination. bufs and addrs must have the same length. Linux uses @@ -56,7 +57,7 @@ func (NoopConn) Rebind() error { func (NoopConn) LocalAddr() (netip.AddrPort, error) { return netip.AddrPort{}, nil } -func (NoopConn) ListenOut(_ EncReader, _ func()) error { +func (NoopConn) ListenOut(_ EncReader, _ func([][]byte) [][]byte) error { return nil } func (NoopConn) SupportsMultipleReaders() bool { diff --git a/udp/udp_darwin.go b/udp/udp_darwin.go index 2bb0cf0d..19c1619d 100644 --- a/udp/udp_darwin.go +++ b/udp/udp_darwin.go @@ -180,8 +180,9 @@ func NewUDPStatsEmitter(udpConns []Conn) func() { return func() {} } -func (u *StdConn) ListenOut(r EncReader, flush func()) error { - buffer := make([]byte, MTU) +func (u *StdConn) ListenOut(r EncReader, handoff func(bufs [][]byte) [][]byte) error { + bufs := [][]byte{make([]byte, MTU)} + buffer := bufs[0] for { // Just read one packet at a time @@ -195,7 +196,8 @@ func (u *StdConn) ListenOut(r EncReader, flush func()) error { } r(netip.AddrPortFrom(rua.Addr().Unmap(), rua.Port()), buffer[:n:n]) - flush() + bufs = handoff(bufs) + buffer = bufs[0] } } diff --git a/udp/udp_generic.go b/udp/udp_generic.go index 0ba1d412..de32cf4a 100644 --- a/udp/udp_generic.go +++ b/udp/udp_generic.go @@ -86,8 +86,9 @@ type rawMessage struct { Len uint32 } -func (u *GenericConn) ListenOut(r EncReader, flush func()) error { - buffer := make([]byte, MTU) +func (u *GenericConn) ListenOut(r EncReader, handoff func(bufs [][]byte) [][]byte) error { + bufs := [][]byte{make([]byte, MTU)} + buffer := bufs[0] var lastRecvErr time.Time @@ -107,7 +108,8 @@ func (u *GenericConn) ListenOut(r EncReader, flush func()) error { } r(netip.AddrPortFrom(rua.Addr().Unmap(), rua.Port()), buffer[:n:n]) - flush() + bufs = handoff(bufs) + buffer = bufs[0] } } diff --git a/udp/udp_linux.go b/udp/udp_linux.go index 7e72742c..b028ff1e 100644 --- a/udp/udp_linux.go +++ b/udp/udp_linux.go @@ -245,14 +245,7 @@ func getFrom(names [][]byte, i int, isV4 bool) netip.AddrPort { return netip.AddrPortFrom(ip.Unmap(), binary.BigEndian.Uint16(names[i][2:4])) } -func (u *StdConn) ListenOut(r EncReader, flush func()) error { - return u.ListenOutHandoff(r, func(bufs [][]byte) [][]byte { - flush() - return bufs - }) -} - -// HandoffShape is the count and size of the buffers ListenOutHandoff hands over after each batch. +// HandoffShape is the count and size of the buffers ListenOut hands over after each batch. func (u *StdConn) HandoffShape() (n, size int) { if u.groSupported { return u.batch, udpGROBufferSize @@ -260,10 +253,7 @@ func (u *StdConn) HandoffShape() (n, size int) { 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 { +func (u *StdConn) ListenOut(r EncReader, handoff func(bufs [][]byte) [][]byte) error { _, bufSize := u.HandoffShape() cmsgSpace := 0 if u.groSupported { diff --git a/udp/udp_linux_offloads_test.go b/udp/udp_linux_offloads_test.go index 80648ca0..db276c0c 100644 --- a/udp/udp_linux_offloads_test.go +++ b/udp/udp_linux_offloads_test.go @@ -226,7 +226,7 @@ func TestOffloadsDisabledRxDelivers(t *testing.T) { _ = sc.ListenOut(func(from netip.AddrPort, payload []byte) { // payload aliases the shared recv buffer row; copy before handing off. rxCh <- rxPkt{from, slices.Clone(payload)} - }, func() {}) + }, func(bufs [][]byte) [][]byte { return bufs }) }() tx, err := net.DialUDP("udp4", nil, net.UDPAddrFromAddrPort(addr)) diff --git a/udp/udp_linux_test.go b/udp/udp_linux_test.go index b99f89eb..fde55fc6 100644 --- a/udp/udp_linux_test.go +++ b/udp/udp_linux_test.go @@ -142,7 +142,7 @@ func runTeardownCase(t *testing.T, batch int, name string, traffic func(send net go func() { loopDone <- sc.ListenOut(func(netip.AddrPort, []byte) { received.Add(1) - }, func() {}) + }, func(bufs [][]byte) [][]byte { return bufs }) }() send, err := net.Dial("udp", addr.String()) diff --git a/udp/udp_rio_windows.go b/udp/udp_rio_windows.go index 435d71a6..b078b745 100644 --- a/udp/udp_rio_windows.go +++ b/udp/udp_rio_windows.go @@ -142,8 +142,9 @@ func (u *RIOConn) bind(l *slog.Logger, sa windows.Sockaddr) error { return nil } -func (u *RIOConn) ListenOut(r EncReader, flush func()) error { - buffer := make([]byte, MTU) +func (u *RIOConn) ListenOut(r EncReader, handoff func(bufs [][]byte) [][]byte) error { + bufs := [][]byte{make([]byte, MTU)} + buffer := bufs[0] var lastRecvErr time.Time @@ -164,7 +165,8 @@ func (u *RIOConn) ListenOut(r EncReader, flush func()) error { } r(netip.AddrPortFrom(netip.AddrFrom16(rua.Addr).Unmap(), (rua.Port>>8)|((rua.Port&0xff)<<8)), buffer[:n:n]) - flush() + bufs = handoff(bufs) + buffer = bufs[0] } } diff --git a/udp/udp_tester.go b/udp/udp_tester.go index 0c436e07..e6fa7809 100644 --- a/udp/udp_tester.go +++ b/udp/udp_tester.go @@ -182,22 +182,14 @@ func (u *TesterConn) WriteBatch(bufs [][]byte, addrs []netip.AddrPort) (int, err return written, nil } -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.ListenOutHandoff(r, func(bufs [][]byte) [][]byte { - flush() - return bufs - }) -} - // 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 { +// ListenOut 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. +// The batcher borrows plaintext decrypted in place inside p.Data until it flushes, inside handoff, so the packet +// must stay alive until handoff returns. +func (u *TesterConn) ListenOut(r EncReader, handoff func(bufs [][]byte) [][]byte) error { bufs := make([][]byte, 1) for { select {