1 files changed, 66 insertions, 11 deletions

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());
+    }
 }