bump number of heads in keepalive messages

Instead of arbitrarily limiting the number to 16, allow filling
keepalive messages as long as they don't cause mycelium to chunk it.

5 files changed, 94 insertions, 11 deletions

diff --git a/src/api/peers.rs b/src/api/peers.rs
index ce41776..3317185 100644
--- a/src/api/peers.rs
+++ b/src/api/peers.rs
@@ -57,12 +57,18 @@ pub(crate) struct PeerStatus {
     epoch: Option<EpochInfo>,
     /// Whether that epoch equals ours, i.e. whether we sync with this peer.
     same_epoch: Option<bool>,
+    /// Number of heads in its last keepalive; absent before one, or when it had more
+    /// than 32 (then only their digest is sent).
+    heads: Option<usize>,
 }
 
 #[derive(Debug, Serialize, ToSchema)]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct PeerListResponse {
     local_epoch: EpochInfo,
+    /// Number of heads of our document. Several heads are concurrent changes no later
+    /// change has merged yet.
+    local_heads: usize,
     peers: Vec<PeerStatus>,
 }
 
@@ -122,12 +128,15 @@ pub(crate) async fn list_peers(
             has_pending_delivery: ps.sync.outbound.is_some(),
             epoch: ps.epoch.map(EpochInfo::from),
             same_epoch: ps.epoch.map(|epoch| epoch == local_epoch),
+            heads: ps.reported_heads,
         });
     }
 
     peers.sort_by(|a, b| a.public_key.cmp(&b.public_key));
+    let local_heads = state.doc.write().await.doc.get_heads().len();
     Ok(Json(PeerListResponse {
         local_epoch: local_epoch.into(),
+        local_heads,
         peers,
     }))
 }
diff --git a/src/state.rs b/src/state.rs
index 8ff8b9d..4d82150 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -126,6 +126,9 @@ pub(crate) struct PeerState {
 
     /// Whether the peer said, in its most recent keepalive, that it would sync with us.
     pub(crate) syncs_with_us: Option<bool>,
+
+    /// How many heads its most recent keepalive listed (None if it sent a digest).
+    pub(crate) reported_heads: Option<usize>,
 }
 
 /// One Automerge sync with a peer, from a reset to the next. Everything a reset must
@@ -330,6 +333,7 @@ impl PeerState {
             epoch: None,
             reports_ready: None,
             syncs_with_us: None,
+            reported_heads: None,
         }
     }
 
diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index 9fe9b07..368c969 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -131,6 +131,7 @@ async fn note_keepalive(state: &AppState, peer: &str, remote: &transport::Keepal
         let mut peer_state = peer_state.lock().await;
         peer_state.reports_ready = Some(remote.ready);
         peer_state.syncs_with_us = Some(remote.syncs_with_you);
+        peer_state.reported_heads = remote.heads.heads().map(<[_]>::len);
     }
     crate::swarm::source::reported(state, peer, remote.ready, Some(remote.syncs_with_you));
 }
diff --git a/src/transport/mod.rs b/src/transport/mod.rs
index aee3ee3..d7e7135 100644
--- a/src/transport/mod.rs
+++ b/src/transport/mod.rs
@@ -12,14 +12,16 @@ pub use crate::document::replica::Epoch;
 pub mod firewall;
 pub mod mycelium;
 
-/// Most heads listed in a keepalive; more are sent as a digest. 16 heads are about
-/// 520 bytes, well within one Mycelium chunk (1300 bytes).
-pub const MAX_LISTED_HEADS: usize = 16;
+/// Most heads listed in a keepalive; more are sent as a digest. A head takes 34 bytes on
+/// the wire, so 32 of them keep a keepalive within one Mycelium chunk (1300 bytes; see
+/// the size test in `mycelium`). There are rarely more heads than nodes that wrote
+/// concurrently since the last merging write.
+pub const MAX_LISTED_HEADS: usize = 32;
 
 /// A document's heads, compactly: the heads themselves while there are few, otherwise a
-/// digest of them. Either way, two summaries are equal exactly when the heads are (the
-/// number of heads decides the form, so equal heads always take the same one).
-#[derive(Debug, Clone, PartialEq, Eq)]
+/// digest of them. Two summaries are equal exactly when the heads are, whatever form
+/// each took (nodes with different limits still agree).
+#[derive(Debug, Clone)]
 pub enum HeadsSummary {
     /// Sorted.
     Heads(Vec<ChangeHash>),
@@ -31,13 +33,18 @@ impl HeadsSummary {
     pub fn of(mut heads: Vec<ChangeHash>) -> Self {
         heads.sort_unstable();
         if heads.len() <= MAX_LISTED_HEADS {
-            return Self::Heads(heads);
+            Self::Heads(heads)
+        } else {
+            Self::Digest(digest_of(&heads))
         }
-        let mut hasher = Sha256::new();
-        for head in &heads {
-            hasher.update(head.0);
+    }
+
+    /// The digest of the heads, computed if they are listed.
+    pub fn digest(&self) -> [u8; 32] {
+        match self {
+            Self::Heads(heads) => digest_of(heads),
+            Self::Digest(digest) => *digest,
         }
-        Self::Digest(hasher.finalize().into())
     }
 
     /// The heads, if listed.
@@ -49,6 +56,26 @@ impl HeadsSummary {
     }
 }
 
+impl PartialEq for HeadsSummary {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::Heads(a), Self::Heads(b)) => a == b,
+            _ => self.digest() == other.digest(),
+        }
+    }
+}
+
+impl Eq for HeadsSummary {}
+
+/// SHA-256 over sorted heads.
+fn digest_of(sorted_heads: &[ChangeHash]) -> [u8; 32] {
+    let mut hasher = Sha256::new();
+    for head in sorted_heads {
+        hasher.update(head.0);
+    }
+    hasher.finalize().into()
+}
+
 /// PING and PONG. Addressed to one peer: `syncs_with_you` is about the recipient.
 #[derive(Debug, Clone)]
 pub struct Keepalive {
diff --git a/src/transport/mycelium.rs b/src/transport/mycelium.rs
index 9d1fb90..2683f82 100644
--- a/src/transport/mycelium.rs
+++ b/src/transport/mycelium.rs
@@ -973,3 +973,45 @@ pub enum TransportError {
         status: reqwest::StatusCode,
     },
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Mycelium splits messages into chunks of this size.
+    const MYCELIUM_CHUNK: usize = 1300;
+
+    #[test]
+    fn largest_keepalive_fits_in_one_chunk() {
+        let keepalive = Keepalive {
+            heads: HeadsSummary::Heads(
+                (0..MAX_LISTED_HEADS as u8)
+                    .map(|n| ChangeHash([n; 32]))
+                    .collect(),
+            ),
+            epoch: Epoch {
+                generation: u64::MAX,
+                id: [0xff; 32],
+            },
+            ready: true,
+            syncs_with_you: true,
+        };
+        let payload = keepalive_to_bytes(&keepalive).unwrap();
+        // Mycelium's message data: a length byte, the topic, then our payload.
+        let size = 1 + TOPIC_PONG.len().max(TOPIC_PING.len()) + payload.len();
+        assert!(
+            size + 50 <= MYCELIUM_CHUNK,
+            "keepalive with {MAX_LISTED_HEADS} heads takes {size} bytes"
+        );
+        assert_eq!(bytes_to_keepalive(&payload).unwrap().heads, keepalive.heads);
+    }
+
+    #[test]
+    fn listed_heads_equal_their_digest() {
+        let heads: Vec<_> = (0..3u8).map(|n| ChangeHash([n; 32])).collect();
+        let listed = HeadsSummary::of(heads.clone());
+        let digest = HeadsSummary::Digest(listed.digest());
+        assert_eq!(listed, digest);
+        assert_ne!(listed, HeadsSummary::of(heads[..2].to_vec()));
+    }
+}