keepalive sends head samples alongside digest

When there are more than 32 heads that can be send in a single keepalive
message, the digest is accompanied by 32 heads, a different window each
time, so after a few keepalives every head has been listed.

4 files changed, 117 insertions, 47 deletions

diff --git a/src/swarm/relay.rs b/src/swarm/relay.rs
index 7c56c90..087e028 100644
--- a/src/swarm/relay.rs
+++ b/src/swarm/relay.rs
@@ -128,20 +128,26 @@ pub(crate) async fn pull_if_behind(
     if ours == remote.heads {
         return false;
     }
-    let missing = match &remote.heads {
-        HeadsSummary::Heads(heads) => {
-            // AutoCommit's lookups take &mut self.
-            let mut replica = state.doc.write().await;
-            let missing = heads
-                .iter()
-                .any(|head| replica.doc.get_change_by_hash(head).is_none());
-            if !missing {
-                return false;
-            }
-            Some(heads.clone())
-        }
-        // Can't tell which heads are missing; the digest differs from ours, though.
-        HeadsSummary::Digest(_) => None,
+    // Those of the listed heads we lack. With a digest only a sample is listed; if we
+    // have all of it, a later keepalive lists others.
+    let missing: Vec<_> = {
+        // AutoCommit's lookups take &mut self.
+        let mut replica = state.doc.write().await;
+        remote
+            .heads
+            .listed()
+            .iter()
+            .filter(|head| replica.doc.get_change_by_hash(head).is_none())
+            .copied()
+            .collect()
+    };
+    if missing.is_empty() {
+        return false;
+    }
+    // All heads, when we know them; otherwise those we lack.
+    let theirs = match remote.heads.heads() {
+        Some(heads) => heads.to_vec(),
+        None => missing,
     };
 
     let Some(peer_state) = state.get_peer_state(peer).await else {
@@ -166,13 +172,8 @@ pub(crate) async fn pull_if_behind(
         pulling.waiting.remove(peer);
     }
 
-    match missing {
-        // Automerge asks for whatever these heads need that we don't have.
-        Some(heads) => peer_state.lock().await.sync.pull_from(heads),
-        // A fresh session's opening carries our heads and bloom filter; the answer
-        // brings what we lack.
-        None => peer_state.lock().await.reset_sync(),
-    }
+    // Automerge asks for whatever these heads need that we don't have.
+    peer_state.lock().await.sync.pull_from(theirs);
     debug!(peer = %peer, "Neighbour has changes we lack; pulling");
     let _ = state
         .sync_events
diff --git a/src/transport/mod.rs b/src/transport/mod.rs
index 614ac48..a040993 100644
--- a/src/transport/mod.rs
+++ b/src/transport/mod.rs
@@ -7,6 +7,7 @@ use sha2::{Digest, Sha256};
 use std::boxed::Box;
 use std::future::Future;
 use std::pin::Pin;
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 pub use crate::document::replica::Epoch;
 
@@ -27,23 +28,42 @@ pub type RunId = [u8; 16];
 pub const MAX_LISTED_HEADS: usize = 32;
 
 /// A document's heads, compactly: the heads themselves while there are few, otherwise a
-/// digest of them. Two summaries are equal exactly when the heads are, whatever form
-/// each took (nodes with different limits still agree).
+/// digest of them with some 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>),
-    /// SHA-256 over the sorted heads.
-    Digest([u8; 32]),
+    Digest {
+        /// SHA-256 over the sorted heads.
+        digest: [u8; 32],
+        /// `MAX_LISTED_HEADS` of the heads, another window of them each time: a receiver
+        /// pulls those it lacks, and over a few keepalives every head is listed, instead
+        /// of starting the sync over to find out what it lacks.
+        sample: Vec<ChangeHash>,
+    },
 }
 
+/// Where the next sample starts (`HeadsSummary::of`).
+static SAMPLE_OFFSET: AtomicUsize = AtomicUsize::new(0);
+
 impl HeadsSummary {
     pub fn of(mut heads: Vec<ChangeHash>) -> Self {
         heads.sort_unstable();
         if heads.len() <= MAX_LISTED_HEADS {
-            Self::Heads(heads)
-        } else {
-            Self::Digest(digest_of(&heads))
+            return Self::Heads(heads);
+        }
+        let start = SAMPLE_OFFSET.fetch_add(MAX_LISTED_HEADS, Ordering::Relaxed) % heads.len();
+        let sample = heads
+            .iter()
+            .cycle()
+            .skip(start)
+            .take(MAX_LISTED_HEADS)
+            .copied()
+            .collect();
+        Self::Digest {
+            digest: digest_of(&heads),
+            sample,
         }
     }
 
@@ -51,15 +71,23 @@ impl HeadsSummary {
     pub fn digest(&self) -> [u8; 32] {
         match self {
             Self::Heads(heads) => digest_of(heads),
-            Self::Digest(digest) => *digest,
+            Self::Digest { digest, .. } => *digest,
         }
     }
 
-    /// The heads, if listed.
+    /// The heads, if all are listed.
     pub fn heads(&self) -> Option<&[ChangeHash]> {
         match self {
             Self::Heads(heads) => Some(heads),
-            Self::Digest(_) => None,
+            Self::Digest { .. } => None,
+        }
+    }
+
+    /// The heads listed: all of them, or a sample.
+    pub fn listed(&self) -> &[ChangeHash] {
+        match self {
+            Self::Heads(heads) => heads,
+            Self::Digest { sample, .. } => sample,
         }
     }
 }
@@ -369,7 +397,16 @@ mod tests {
             let mut shuffled = heads.clone();
             rng.shuffle(&mut shuffled);
             assert_eq!(summary, HeadsSummary::of(shuffled));
-            assert_eq!(summary, HeadsSummary::Digest(summary.digest()));
+            assert_eq!(
+                summary,
+                HeadsSummary::Digest {
+                    digest: summary.digest(),
+                    sample: Vec::new(),
+                }
+            );
+            // A sample is some of the heads.
+            assert!(summary.listed().iter().all(|head| heads.contains(head)));
+            assert_eq!(summary.listed().len(), n.min(MAX_LISTED_HEADS));
 
             let mut other = heads.clone();
             match rng.below(3) {
@@ -386,4 +423,22 @@ mod tests {
             assert_ne!(summary, HeadsSummary::of(other));
         });
     }
+
+    /// Successive samples cover every head.
+    #[test]
+    fn samples_rotate_over_all_heads() {
+        let heads: Vec<ChangeHash> = (0..100u8).map(|n| ChangeHash([n; 32])).collect();
+        let mut seen = std::collections::HashSet::new();
+        for _ in 0..heads.len().div_ceil(MAX_LISTED_HEADS) + 1 {
+            seen.extend(HeadsSummary::of(heads.clone()).listed().iter().copied());
+        }
+        // Other tests sample concurrently and move the window too, so allow a few more.
+        for _ in 0..100 {
+            if seen.len() == heads.len() {
+                break;
+            }
+            seen.extend(HeadsSummary::of(heads.clone()).listed().iter().copied());
+        }
+        assert_eq!(seen.len(), heads.len());
+    }
 }
diff --git a/src/transport/mycelium/mod.rs b/src/transport/mycelium/mod.rs
index 585b5ee..495f75e 100644
--- a/src/transport/mycelium/mod.rs
+++ b/src/transport/mycelium/mod.rs
@@ -529,9 +529,10 @@ mod tests {
 
     #[test]
     fn largest_keepalive_fits_in_one_chunk() {
+        // A digest with a full sample: the most a keepalive carries.
         let keepalive = Keepalive {
-            heads: HeadsSummary::Heads(
-                (0..MAX_LISTED_HEADS as u8)
+            heads: HeadsSummary::of(
+                (0..=MAX_LISTED_HEADS as u8)
                     .map(|n| ChangeHash([n; 32]))
                     .collect(),
             ),
@@ -624,7 +625,10 @@ mod tests {
     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());
+        let digest = HeadsSummary::Digest {
+            digest: listed.digest(),
+            sample: Vec::new(),
+        };
         assert_eq!(listed, digest);
         assert_ne!(listed, HeadsSummary::of(heads[..2].to_vec()));
     }
diff --git a/src/transport/wire.rs b/src/transport/wire.rs
index 11d36cf..3c80b11 100644
--- a/src/transport/wire.rs
+++ b/src/transport/wire.rs
@@ -14,7 +14,7 @@ use super::{MessageKind, ParsedMessage};
 
 /// The wire format's version, the first byte of every payload. Bump it with every change
 /// to a message type or a codec (`golden_bytes` fails then).
-pub(crate) const VERSION: u8 = 1;
+pub(crate) const VERSION: u8 = 2;
 
 /// What the second byte of a payload says.
 fn tag(kind: MessageKind) -> u8 {
@@ -200,7 +200,10 @@ pub(crate) mod codec {
                     heads.iter().map(|head| Bytes::new(&head.0)).collect(),
                     Bytes::new(&[]),
                 ),
-                HeadsSummary::Digest(digest) => (Vec::new(), Bytes::new(digest)),
+                HeadsSummary::Digest { digest, sample } => (
+                    sample.iter().map(|head| Bytes::new(&head.0)).collect(),
+                    Bytes::new(digest),
+                ),
             };
             (listed, digest).serialize(s)
         }
@@ -209,21 +212,24 @@ pub(crate) mod codec {
             d: D,
         ) -> Result<HeadsSummary, D::Error> {
             let (listed, digest): (Vec<ByteBuf>, ByteBuf) = Deserialize::deserialize(d)?;
-            if !digest.is_empty() {
-                return Ok(HeadsSummary::Digest(fixed(&digest, "a heads digest")?));
-            }
             if listed.len() > MAX_LISTED_HEADS {
                 return Err(D::Error::custom(format!(
                     "more than {MAX_LISTED_HEADS} heads listed"
                 )));
             }
-            listed
+            let listed = listed
                 .iter()
                 .map(|head| -> Result<ChangeHash, D::Error> {
                     Ok(ChangeHash(fixed(head, "a head")?))
                 })
-                .collect::<Result<Vec<_>, _>>()
-                .map(HeadsSummary::Heads)
+                .collect::<Result<Vec<_>, _>>()?;
+            if digest.is_empty() {
+                return Ok(HeadsSummary::Heads(listed));
+            }
+            Ok(HeadsSummary::Digest {
+                digest: fixed(&digest, "a heads digest")?,
+                sample: listed,
+            })
         }
     }
 
@@ -317,7 +323,10 @@ mod tests {
     #[test]
     fn keepalives_keep_their_kind_and_heads() {
         let keepalive = Keepalive {
-            heads: HeadsSummary::Digest([8; 32]),
+            heads: HeadsSummary::Digest {
+                digest: [8; 32],
+                sample: vec![ChangeHash([9; 32])],
+            },
             epoch: Epoch::GENESIS,
             epoch_seq: 0,
             ready: true,
@@ -328,6 +337,7 @@ mod tests {
             panic!("not a PONG");
         };
         assert_eq!(decoded.heads, keepalive.heads);
+        assert_eq!(decoded.heads.listed(), [ChangeHash([9; 32])]);
     }
 
     #[test]
@@ -401,7 +411,7 @@ mod tests {
         });
         assert_eq!(
             hex(&encode(&hello).unwrap()),
-            format!("010097a5737761726d{epoch_hex}01c403010203c3{}07", run("04"))
+            format!("020097a5737761726d{epoch_hex}01c403010203c3{}07", run("04"))
         );
 
         let reply = ParsedMessage::HelloReply(HelloReply {
@@ -418,7 +428,7 @@ mod tests {
         assert_eq!(
             hex(&encode(&reply).unwrap()),
             format!(
-                "010199a5737761726d91a170{epoch_hex}01c400c2{}05{}",
+                "020199a5737761726d91a170{epoch_hex}01c400c2{}05{}",
                 run("04"),
                 run("06")
             )
@@ -433,7 +443,7 @@ mod tests {
         });
         assert_eq!(
             hex(&encode(&ping).unwrap()),
-            format!("0103959290c400{epoch_hex}01c3c2")
+            format!("0203959290c400{epoch_hex}01c3c2")
         );
     }
 }