relay: Let WriteBatch handle destination changes and drop the flush before responses

outside.go

@@ -77,12 +77,6 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont
 		f.messageMetrics.Rx(h.Type, h.Subtype, 1)
 	}
 
-	// Earlier forwards must leave before this packet can generate a response.
-	isMessageRelay := h.Type == header.Message && h.Subtype == header.MessageRelay
-	if !isMessageRelay {
-		f.flushRelayBatch(rxc)
-	}
-
 	// Unencrypted packets
 	switch h.Type {
 	case header.Handshake:
@@ -94,6 +88,9 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont
 		return
 	}
 
+	// Relay packets are special
+	isMessageRelay := (h.Type == header.Message && h.Subtype == header.MessageRelay)
+
 	var hostinfo *HostInfo
 	if isMessageRelay {
 		hostinfo = f.hostMap.QueryRelayIndex(h.RemoteIndex)
@@ -105,8 +102,6 @@ func (f *Interface) readOutsidePackets(via ViaSender, packet []byte, rxc *rxCont
 	// recvError if necessary
 	if hostinfo == nil || hostinfo.ConnectionState == nil {
 		if !via.IsRelayed {
-			// An unknown relay index can also produce a response.
-			f.flushRelayBatch(rxc)
 			f.maybeSendRecvError(via.UdpAddr, h.RemoteIndex)
 		}
 		return
@@ -240,7 +235,11 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender,
 				fwdBuf := packet[:0]
 				out, err := f.prepareSendVia(targetHI, targetRelay, signedPayload, rxc.nb, fwdBuf, true)
 				if err == nil {
-					f.queueRelayPacket(out, targetHI.GetRemote(), rxc)
+					// Borrows the receive buffer, the listener's flush callback drains it before the buffer is reused
+					rxc.relayBatch.Commit(out, targetHI.GetRemote())
+					if rxc.relayBatch.Len() >= batch.SendBatchCap {
+						f.flushSendBatch(rxc.relayBatch, rxc.q)
+					}
 				}
 			case TerminalType:
 				hostinfo.logger(f.l).Error("Unexpected Relay Type of Terminal")
@@ -266,27 +265,6 @@ func (f *Interface) handleOutsideRelayPacket(hostinfo *HostInfo, via ViaSender,
 	}
 }
 
-// queueRelayPacket borrows a receive buffer until the batch is flushed. Destination changes and the send batch
-// limit end a run, preserving arrival order while allowing WriteBatch to use UDP segmentation offload.
-func (f *Interface) queueRelayPacket(packet []byte, remote netip.AddrPort, rxc *rxContext) {
-	if rxc.relayBatch == nil {
-		rxc.relayBatch = batch.NewSendBatch(f.writers[rxc.q], batch.SendBatchCap, 0)
-	} else if rxc.relayRemote != remote {
-		f.flushRelayBatch(rxc)
-	}
-	rxc.relayRemote = remote
-	rxc.relayBatch.Commit(packet, remote)
-	if rxc.relayBatch.Len() >= batch.SendBatchCap {
-		f.flushRelayBatch(rxc)
-	}
-}
-
-func (f *Interface) flushRelayBatch(rxc *rxContext) {
-	if rxc.relayBatch != nil {
-		f.flushSendBatch(rxc.relayBatch, rxc.q)
-	}
-}
-
 // closeTunnel closes a tunnel locally, it does not send a closeTunnel packet to the remote
 func (f *Interface) closeTunnel(hostInfo *HostInfo) {
 	// Hold the hostmap lock across both the delete and the lighthouse eviction,