refactor(iroh): Re-batch datagrams inside RelayTransport instead of the ActiveRelayActor (#3421)

## Description

This does the re-batching inside `RelayTransport` by storing a
`pending_item: Option<RelayRecvDatagram>`.
When we `poll_recv`, we first try to use that pending item instead of
polling a new one.

Once we've got a pending item, we try to split off as much from it as
can possibly fit into our receive buffer and handle that.

## Breaking Changes

- `iroh_relay::protos::relay::Datagrams::take_segments` now return
`Datagrams` instead of `Option<Datagrams>`, not distinguishing the case
where `Datagrams` might be empty.

## Notes

Probably needs some test.
I'd love to test two `RelayTransport`s talking to each other with
different `max_transmit_segments`/`max_receive_segments`, but I'm not
sure I can make such a test setup happen easily.

## Change checklist
<!-- Remove any that are not relevant. -->
- [x] Self-review.
- [x] Documentation updates following the [style
guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text),
if relevant.

Author: matheus23

iroh-relay/src/protos/relay.rs

@@ -157,32 +157,26 @@ impl<T: AsRef<[u8]>> From<T> for Datagrams {
 }
 
 impl Datagrams {
-    /// Splits the current datagram into at maximum `num_segments` segments, potentially returning
+    /// Splits the current datagram into at maximum `num_segments` segments, returning
     /// the batch with at most `num_segments` and leaving only the rest in `self`.
     ///
     /// Calling this on a datagram batch that only contains a single datagram (`segment_size == None`)
-    /// will result in returning essentially `Some(self.clone())`, while making `self` empty afterwards.
+    /// will result in returning essentially a clone of `self`, while making `self` empty afterwards.
     ///
     /// Calling this on a datagram batch with e.g. 15 datagrams with `num_segments == 10` will
-    /// result in returning `Some(datagram_batch)` where that `datagram_batch` contains the first
-    /// 10 datagrams and `self` contains the remaining 5 datagrams.
+    /// result in returning a datagram batch that contains the first 10 datagrams and leave `self`
+    /// containing the remaining 5 datagrams.
     ///
-    /// Calling this on a datagram batch that doesn't contain `num_segments` datagrams, but less
-    /// will result in making `self` empty and returning essentially a clone of `self`.
-    ///
-    /// Calling this on an empty datagram batch (i.e. one where `contents.is_empty()`) will return `None`.
-    pub fn take_segments(&mut self, num_segments: usize) -> Option<Datagrams> {
-        if self.contents.is_empty() {
-            return None;
-        }
-
+    /// Calling this on a datagram batch with less than `num_segments` datagrams will result in
+    /// making `self` empty and returning essentially a clone of `self`.
+    pub fn take_segments(&mut self, num_segments: usize) -> Datagrams {
         let Some(segment_size) = self.segment_size else {
             let contents = std::mem::take(&mut self.contents);
-            return Some(Datagrams {
+            return Datagrams {
                 ecn: self.ecn,
                 segment_size: None,
                 contents,
-            });
+            };
         };
 
         let usize_segment_size = usize::from(u16::from(segment_size));
@@ -193,11 +187,17 @@ impl Datagrams {
 
         let is_datagram_batch = num_segments > 1 && usize_segment_size < contents.len();
 
-        Some(Datagrams {
+        // If this left our batch with only one more datagram, then remove the segment size
+        // to uphold the invariant that single-datagram batches don't have a segment size set.
+        if self.contents.len() <= usize_segment_size {
+            self.segment_size = None;
+        }
+
+        Datagrams {
             ecn: self.ecn,
             segment_size: is_datagram_batch.then_some(segment_size),
             contents,
-        })
+        }
     }
 
     fn write_to<O: BufMut>(&self, mut dst: O) -> O {

iroh/src/magicsock.rs

@@ -23,7 +23,7 @@ use std::{
     pin::Pin,
     sync::{
         Arc, Mutex, RwLock,
-        atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
+        atomic::{AtomicBool, AtomicU64, Ordering},
     },
     task::{Context, Poll},
 };
@@ -1380,7 +1380,6 @@ impl Handle {
 
         let my_relay = Watchable::new(None);
         let ipv6_reported = Arc::new(AtomicBool::new(ipv6_reported));
-        let max_receive_segments = Arc::new(AtomicUsize::new(1));
 
         let relay_transport = RelayTransport::new(RelayActorConfig {
             my_relay: my_relay.clone(),
@@ -1389,7 +1388,6 @@ impl Handle {
             dns_resolver: dns_resolver.clone(),
             proxy_url: proxy_url.clone(),
             ipv6_reported: ipv6_reported.clone(),
-            max_receive_segments: max_receive_segments.clone(),
             #[cfg(any(test, feature = "test-utils"))]
             insecure_skip_relay_cert_verify,
             metrics: metrics.magicsock.clone(),
@@ -1401,9 +1399,9 @@ impl Handle {
         let ipv6 = ip_transports.iter().any(|t| t.bind_addr().is_ipv6());
 
         #[cfg(not(wasm_browser))]
-        let transports = Transports::new(ip_transports, relay_transports, max_receive_segments);
+        let transports = Transports::new(ip_transports, relay_transports);
         #[cfg(wasm_browser)]
-        let transports = Transports::new(relay_transports, max_receive_segments);
+        let transports = Transports::new(relay_transports);
 
         let (disco, disco_receiver) = DiscoState::new(secret_encryption_key);
 

iroh/src/magicsock/transports.rs

@@ -32,7 +32,6 @@ pub(crate) struct Transports {
     ip: Vec<IpTransport>,
     relay: Vec<RelayTransport>,
 
-    max_receive_segments: Arc<AtomicUsize>,
     poll_recv_counter: AtomicUsize,
 }
 
@@ -62,13 +61,11 @@ impl Transports {
     pub(crate) fn new(
         #[cfg(not(wasm_browser))] ip: Vec<IpTransport>,
         relay: Vec<RelayTransport>,
-        max_receive_segments: Arc<AtomicUsize>,
     ) -> Self {
         Self {
             #[cfg(not(wasm_browser))]
             ip,
             relay,
-            max_receive_segments,
             poll_recv_counter: Default::default(),
         }
     }
@@ -199,7 +196,6 @@ impl Transports {
 
     #[cfg(not(wasm_browser))]
     pub(crate) fn max_receive_segments(&self) -> usize {
-        use std::sync::atomic::Ordering::Relaxed;
         // `max_receive_segments` controls the size of the `RecvMeta` buffer
         // that quinn creates. Having buffers slightly bigger than necessary
         // isn't terrible, and makes sure a single socket can read the maximum
@@ -208,9 +204,7 @@ impl Transports {
         // and it's impossible and unnecessary to be refactored that way.
 
         let res = self.ip.iter().map(|t| t.max_receive_segments()).max();
-        let segments = res.unwrap_or(1);
-        self.max_receive_segments.store(segments, Relaxed);
-        segments
+        res.unwrap_or(1)
     }
 
     #[cfg(wasm_browser)]

iroh/src/magicsock/transports/relay.rs

@@ -29,6 +29,8 @@ pub(crate) struct RelayTransport {
     relay_datagram_recv_queue: mpsc::Receiver<RelayRecvDatagram>,
     /// Channel on which to send datagrams via a relay server.
     relay_datagram_send_channel: mpsc::Sender<RelaySendItem>,
+    /// A datagram from the last poll_recv that didn't quite fit our buffers.
+    pending_item: Option<RelayRecvDatagram>,
     actor_sender: mpsc::Sender<RelayActorMessage>,
     _actor_handle: AbortOnDropHandle<()>,
     my_relay: Watchable<Option<RelayUrl>>,
@@ -60,6 +62,7 @@ impl RelayTransport {
         Self {
             relay_datagram_recv_queue: relay_datagram_recv_rx,
             relay_datagram_send_channel: relay_datagram_send_tx,
+            pending_item: None,
             actor_sender,
             _actor_handle: actor_handle,
             my_relay,
@@ -86,7 +89,7 @@ impl RelayTransport {
             .zip(metas.iter_mut())
             .zip(source_addrs.iter_mut())
         {
-            let dm = match self.relay_datagram_recv_queue.poll_recv(cx) {
+            let dm = match self.poll_recv_queue(cx) {
                 Poll::Ready(Some(recv)) => recv,
                 Poll::Ready(None) => {
                     error!("relay_recv_channel closed");
@@ -100,6 +103,23 @@ impl RelayTransport {
                 }
             };
 
+            // This *tries* to make the datagrams fit into our buffer by re-batching them.
+            let num_segments = dm
+                .datagrams
+                .segment_size
+                .map_or(1, |ss| buf_out.len() / u16::from(ss) as usize);
+            let datagrams = dm.datagrams.take_segments(num_segments);
+            let empty_after = dm.datagrams.contents.is_empty();
+            let dm = RelayRecvDatagram {
+                datagrams,
+                src: dm.src,
+                url: dm.url.clone(),
+            };
+            // take_segments can leave `self.pending_item` empty, in that case we clear it
+            if empty_after {
+                self.pending_item = None;
+            }
+
             if buf_out.len() < dm.datagrams.contents.len() {
                 // Our receive buffer isn't big enough to process this datagram.
                 // Continuing would cause a panic.
@@ -119,6 +139,7 @@ impl RelayTransport {
                 // that's essentially bigger than our configured `max_udp_payload_size`.
                 // In that case we drop it and let MTU discovery take over.
             }
+
             buf_out[..dm.datagrams.contents.len()].copy_from_slice(&dm.datagrams.contents);
             meta_out.len = dm.datagrams.contents.len();
             meta_out.stride = dm
@@ -156,6 +177,28 @@ impl RelayTransport {
             sender: self.actor_sender.clone(),
         }
     }
+
+    /// Makes sure we have a pending item stored, if not, it'll poll a new one from the queue.
+    ///
+    /// Returns a mutable reference to the stored pending item.
+    #[inline]
+    fn poll_recv_queue<'a>(
+        &'a mut self,
+        cx: &mut Context,
+    ) -> Poll<Option<&'a mut RelayRecvDatagram>> {
+        // Borrow checker doesn't quite understand an if let Some(_)... here
+        if self.pending_item.is_some() {
+            return Poll::Ready(self.pending_item.as_mut());
+        }
+
+        let item = match self.relay_datagram_recv_queue.poll_recv(cx) {
+            Poll::Ready(Some(item)) => item,
+            Poll::Ready(None) => return Poll::Ready(None),
+            Poll::Pending => return Poll::Pending,
+        };
+
+        Poll::Ready(Some(self.pending_item.insert(item)))
+    }
 }
 
 #[derive(Debug)]