udp: Fold buffer handoff into ListenOut

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 {