udp: Fold buffer handoff into ListenOut

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
 }

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

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)

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 {

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 {

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]
 	}
 }
 

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]
 	}
 }
 

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 {

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))

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())

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]
 	}
 }
 

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 {