links: when asking for a lot relays, do chunk

1 files changed, 72 insertions, 17 deletions

diff --git a/src/swarm/links.rs b/src/swarm/links.rs
index 4b89e94..7f8f0b0 100644
--- a/src/swarm/links.rs
+++ b/src/swarm/links.rs
@@ -61,6 +61,11 @@ const WATCH_EVERY: Duration = Duration::from_secs(2);
 const REPLY_WITHIN: Duration = Duration::from_secs(30);
 /// Members a relayed LINKS names at most, and that one LINKS_REPLY passes on at most.
 const MAX_RELAYED: usize = 64;
+/// Bytes of relayed links one LINKS_REPLY carries at most, as encoded (one member's
+/// alone may be more): what a member reports is usually small (about 12 links of 150
+/// bytes), but may be up to `MAX_LINKS` links, and 64 such would make a message of
+/// megabytes.
+const MAX_RELAYED_BYTES: usize = 256 * 1024;
 /// How long relayed links are held to be sent together, along the whole way: the
 /// asker's window (at most `MAX_BATCH`), each hop holding at least `MIN_BATCH`.
 pub(crate) const DEFAULT_BATCH: Duration = Duration::from_millis(250);
@@ -404,9 +409,9 @@ async fn relay_for(state: &AppState, member: &str, exclude: Option<&str>) -> Opt
 }
 
 /// Asks `relay` to ask `members` on our behalf (a LINKS naming them), with `subscribe`,
-/// passing their links on within `window`, through `hops` more nodes at most. One nonce
-/// per relay while its replies are taken, so its replies for every member we asked it
-/// for are taken.
+/// passing their links on within `window`, through `hops` more nodes at most: in LINKS
+/// of `MAX_RELAYED` members at most. One nonce per relay while its replies are taken, so
+/// its replies for every member we asked it for are taken.
 async fn ask_via(
     state: &AppState,
     relay: &str,
@@ -415,7 +420,6 @@ async fn ask_via(
     window: Duration,
     hops: u8,
 ) {
-    let members: Vec<String> = members.into_iter().take(MAX_RELAYED).collect();
     let nonce = {
         let now = Instant::now();
         let mut inner = state.links.inner.lock_or_recover();
@@ -439,17 +443,19 @@ async fn ask_via(
         entry.0
     };
     let (epoch, epoch_seq) = state.epoch_stamp();
-    let request = ParsedMessage::Links(LinksRequest {
-        epoch,
-        epoch_seq,
-        nonce,
-        subscribe_secs: u32::try_from(subscribe.as_secs()).unwrap_or(u32::MAX),
-        of: members,
-        batch_ms: u32::try_from(window.as_millis()).unwrap_or(u32::MAX),
-        hops,
-    });
-    if let Err(error) = state.transport.send_message(relay, &request).await {
-        debug!(relay = %relay, error = %format!("{error:#}"), "Relayed LINKS failed");
+    for members in members.chunks(MAX_RELAYED) {
+        let request = ParsedMessage::Links(LinksRequest {
+            epoch,
+            epoch_seq,
+            nonce,
+            subscribe_secs: u32::try_from(subscribe.as_secs()).unwrap_or(u32::MAX),
+            of: members.to_vec(),
+            batch_ms: u32::try_from(window.as_millis()).unwrap_or(u32::MAX),
+            hops,
+        });
+        if let Err(error) = state.transport.send_message(relay, &request).await {
+            debug!(relay = %relay, error = %format!("{error:#}"), "Relayed LINKS failed");
+        }
     }
 }
 
@@ -713,13 +719,13 @@ async fn pass_on(state: &AppState, asker: &str) {
     }
     let (epoch, epoch_seq) = state.epoch_stamp();
     for (nonce, relayed) in by_nonce {
-        for chunk in relayed.chunks(MAX_RELAYED) {
+        for part in by_size(relayed) {
             let reply = ParsedMessage::LinksReply(LinksReply {
                 epoch,
                 epoch_seq,
                 nonce,
                 links: Vec::new(),
-                relayed: chunk.to_vec(),
+                relayed: part,
             });
             if let Err(error) = state.transport.send_message(asker, &reply).await {
                 debug!(asker = %asker, error = %format!("{error:#}"), "Relayed LINKS_REPLY failed");
@@ -728,6 +734,26 @@ async fn pass_on(state: &AppState, asker: &str) {
     }
 }
 
+/// `relayed` in parts for a LINKS_REPLY each: `MAX_RELAYED` members' links at most, and
+/// `MAX_RELAYED_BYTES` as encoded (but for one member's alone).
+fn by_size(relayed: Vec<RelayedLinks>) -> Vec<Vec<RelayedLinks>> {
+    let mut parts: Vec<Vec<RelayedLinks>> = Vec::new();
+    let mut bytes = 0;
+    for links in relayed {
+        let size = rmp_serde::encode::to_vec(&links).map_or(0, |encoded| encoded.len());
+        let full = parts.last().is_none_or(|part| {
+            part.len() >= MAX_RELAYED || (!part.is_empty() && bytes + size > MAX_RELAYED_BYTES)
+        });
+        if full {
+            parts.push(Vec::new());
+            bytes = 0;
+        }
+        bytes += size;
+        parts.last_mut().expect("just made").push(links);
+    }
+    parts
+}
+
 /// A LINKS_REPLY from `peer`: its own links kept if it answers our latest LINKS to it,
 /// and those it relays if it answers our latest relayed one, while their replies are
 /// taken (a subscription's go on); then passed on to whoever asked us for them.
@@ -881,6 +907,35 @@ fn changed(old: &[Link], new: &[Link]) -> bool {
 mod tests {
     use super::*;
 
+    #[test]
+    fn relayed_links_are_sent_in_parts_by_count_and_size() {
+        let relayed = |links: usize| RelayedLinks {
+            member: "m".into(),
+            age_ms: 0,
+            links: (0..links)
+                .map(|i| link(&format!("peer-{i:0>70}"), Some(1500)))
+                .collect(),
+        };
+        // Small: by count.
+        let parts = by_size((0..130).map(|_| relayed(2)).collect());
+        assert_eq!(parts.iter().map(Vec::len).collect::<Vec<_>>(), [64, 64, 2]);
+        // Large: by size, each part within the bytes but for one alone.
+        let big = rmp_serde::encode::to_vec(&relayed(MAX_LINKS))
+            .unwrap()
+            .len();
+        assert!(20 * big > MAX_RELAYED_BYTES, "{big}");
+        let parts = by_size((0..20).map(|_| relayed(MAX_LINKS)).collect());
+        assert_eq!(parts.iter().map(Vec::len).sum::<usize>(), 20);
+        for part in &parts {
+            let bytes: usize = part
+                .iter()
+                .map(|links| rmp_serde::encode::to_vec(links).unwrap().len())
+                .sum();
+            assert!(part.len() == 1 || bytes <= MAX_RELAYED_BYTES, "{bytes}");
+        }
+        assert!(parts.len() > 1);
+    }
+
     fn link(peer: &str, rtt_us: Option<u32>) -> Link {
         Link {
             peer: peer.into(),