enrich keepalive content

Add ready and syncs_with_you flags as well as actual heads, so the peer
- ready: same meaning as in the sync message
- syncs_with_you: sender counts recipient as member and in same epoch
- heads: allows to tell who's behind or if both sides have changes

5 files changed, 307 insertions, 44 deletions

diff --git a/src/state.rs b/src/state.rs
index 23892a9..72013b9 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -121,8 +121,11 @@ pub(crate) struct PeerState {
     /// Epoch the peer advertised in its most recent message.
     pub(crate) epoch: Option<Epoch>,
 
-    /// Whether the peer was ready according to its most recent sync message.
+    /// Whether the peer was ready according to its most recent sync message or keepalive.
     pub(crate) reports_ready: Option<bool>,
+
+    /// Whether the peer said, in its most recent keepalive, that it would sync with us.
+    pub(crate) syncs_with_us: Option<bool>,
 }
 
 /// One Automerge sync with a peer, from a reset to the next. Everything a reset must
@@ -326,6 +329,7 @@ impl PeerState {
             last_resync_at: None,
             epoch: None,
             reports_ready: None,
+            syncs_with_us: None,
         }
     }
 
@@ -431,8 +435,8 @@ impl PeerState {
     /// Self-healing check on PING/PONG. Returns true if a re-handshake should start.
     pub(crate) fn needs_heal(
         &mut self,
-        local: &transport::HeadsDigest,
-        remote: &transport::HeadsDigest,
+        local: &transport::HeadsSummary,
+        remote: &transport::HeadsSummary,
         timing: &Timing,
         now: Instant,
     ) -> bool {
@@ -792,11 +796,11 @@ impl AppState {
             }
         }
 
-        let keepalive = crate::swarm::peer::local_keepalive(self).await;
         for id in announce {
+            let keepalive = crate::swarm::peer::local_keepalive(self, id).await;
             if let Err(error) = self
                 .transport
-                .send_message(id, &transport::ParsedMessage::Ping(keepalive.clone()))
+                .send_message(id, &transport::ParsedMessage::Ping(keepalive))
                 .await
             {
                 debug!(peer = %id, error = %format!("{error:#}"), "Epoch announcement failed");
diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index 1e8b91d..f98bf17 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -6,7 +6,6 @@ use std::time::Duration;
 
 use anyhow::{Context, Result, anyhow, bail};
 use automerge::AutoCommit;
-use sha2::{Digest, Sha256};
 use tokio::sync::Mutex;
 use tokio::time::{Instant, MissedTickBehavior};
 use tracing::{debug, info, instrument, trace, warn};
@@ -98,23 +97,40 @@ async fn apply_forget_records(state: &AppState) {
 
 // --- keepalive and self-healing ----------------------------------------------
 
-fn heads_digest(doc: &mut AutoCommit) -> transport::HeadsDigest {
-    let mut heads = doc.get_heads();
-    heads.sort_unstable();
-    let mut hasher = Sha256::new();
-    for head in &heads {
-        hasher.update(head.0);
+/// Our keepalive for `peer`.
+pub(crate) async fn local_keepalive(state: &AppState, peer: &str) -> transport::Keepalive {
+    // AutoCommit::get_heads takes &mut self (it closes the pending transaction).
+    // Heads and epoch come from the same locked replica, so they always match.
+    let (heads, epoch) = {
+        let mut replica = state.doc.write().await;
+        (
+            transport::HeadsSummary::of(replica.doc.get_heads()),
+            replica.epoch,
+        )
+    };
+    let syncs_with_you = match state.get_peer_state(peer).await {
+        Some(peer_state) => {
+            let peer_state = peer_state.lock().await;
+            peer_state.liveness == Liveness::Active
+                && matches!(peer_state.standing, Standing::Member { .. })
+                && peer_state.epoch == Some(epoch)
+        }
+        None => false,
+    };
+    transport::Keepalive {
+        heads,
+        epoch,
+        ready: state.is_ready().await,
+        syncs_with_you,
     }
-    hasher.finalize().into()
 }
 
-pub(crate) async fn local_keepalive(state: &AppState) -> transport::Keepalive {
-    // AutoCommit::get_heads takes &mut self (it closes the pending transaction).
-    // Digest and epoch come from the same locked replica, so they always match.
-    let mut replica = state.doc.write().await;
-    transport::Keepalive {
-        heads_digest: heads_digest(&mut replica.doc),
-        epoch: replica.epoch,
+/// What a peer's keepalive says about it.
+async fn note_keepalive(state: &AppState, peer: &str, remote: &transport::Keepalive) {
+    if let Some(peer_state) = state.get_peer_state(peer).await {
+        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);
     }
 }
 
@@ -124,7 +140,7 @@ async fn heal_if_diverged(
     local: &transport::Keepalive,
     remote: &transport::Keepalive,
 ) {
-    // Digests of documents from different epochs are never equal and never converge.
+    // Heads of documents from different epochs are never equal and never converge.
     if local.epoch != remote.epoch {
         return;
     }
@@ -135,18 +151,13 @@ async fn heal_if_diverged(
 
     let heal = {
         let mut peer_state = peer_state_mutex.lock().await;
-        if local.heads_digest == remote.heads_digest {
+        if local.heads == remote.heads {
             // Same document: whatever answer we were waiting for, there is nothing left
             // to send. (Automerge answers only when it has something to say, so a sync
             // message to a peer that already has everything gets no reply at all.)
             peer_state.sync.reply_settled();
         }
-        peer_state.needs_heal(
-            &local.heads_digest,
-            &remote.heads_digest,
-            &state.timing,
-            Instant::now(),
-        )
+        peer_state.needs_heal(&local.heads, &remote.heads, &state.timing, Instant::now())
     };
 
     if heal {
@@ -432,7 +443,7 @@ pub(crate) async fn run_peer_worker(state: AppState) {
             }
 
             if ping_due {
-                let local = local_keepalive(&state).await;
+                let local = local_keepalive(&state, peer).await;
                 if let Err(error) = state
                     .transport
                     .send_message(peer, &ParsedMessage::Ping(local))
@@ -747,7 +758,8 @@ pub(crate) async fn handle_message(
             if reactivated || peer_swarm(state, &peer).await.is_none() {
                 ensure_handshake(state, &peer).await;
             }
-            let local = local_keepalive(state).await;
+            note_keepalive(state, &peer, &remote).await;
+            let local = local_keepalive(state, &peer).await;
             heal_if_diverged(state, &peer, &local, &remote).await;
             if let Err(error) = state
                 .transport
@@ -761,7 +773,8 @@ pub(crate) async fn handle_message(
             if reactivated || peer_swarm(state, &peer).await.is_none() {
                 ensure_handshake(state, &peer).await;
             }
-            let local = local_keepalive(state).await;
+            note_keepalive(state, &peer, &remote).await;
+            let local = local_keepalive(state, &peer).await;
             heal_if_diverged(state, &peer, &local, &remote).await;
         }
     }
diff --git a/src/swarm/source.rs b/src/swarm/source.rs
new file mode 100644
index 0000000..6be62b7
--- /dev/null
+++ b/src/swarm/source.rs
@@ -0,0 +1,169 @@
+//! The catch-up source: while a node catches up (joining, or rejoining a newer epoch,
+//! until its initial sync is done), it syncs with one peer only.
+//!
+//! Every peer would otherwise send us what we lack, all at once: a node with an empty
+//! document gets the whole document from each of its peers, and a node that is merely
+//! behind gets the missing changes from each. Automerge can't deduplicate across peers.
+//! So the first suitable peer we would sync with becomes the source; the others wait
+//! until the initial sync is done, and then only exchange heads. Messages from waiting
+//! peers are still applied; we just don't answer them yet.
+//!
+//! A source must be ready and willing to sync with us, as far as we know (its sync
+//! messages and keepalives say so). It counts as alive while any message from it
+//! arrives, so a slow transfer of a large document never looks like a stall. When it
+//! is quiet for `Timing::catch_up_probe`, we PING it; after three unanswered probes, or
+//! when it says it isn't ready, or twice in a row that it wouldn't sync with us (the
+//! first time may be a handshake still under way), the next peer takes over.
+
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
+use tokio::time::Instant;
+use tracing::{debug, info};
+
+use crate::state::{AppState, CatchUpSource, PeerState};
+use crate::swarm::sync::SyncEvent;
+
+/// Probes without an answer before the source is given up.
+const UNANSWERED_PROBES: u32 = 3;
+
+/// Keepalives in a row saying "I wouldn't sync with you" before the source is given up.
+const REFUSALS: u32 = 2;
+
+fn catching_up(state: &AppState) -> bool {
+    !state.initial_sync_done.load(Ordering::Acquire)
+}
+
+fn lock(state: &AppState) -> std::sync::MutexGuard<'_, Option<CatchUpSource>> {
+    state
+        .catch_up_source
+        .lock()
+        .unwrap_or_else(|poisoned| poisoned.into_inner())
+}
+
+/// Returns how long to hold back sync with `peer`, if at all: while catching up, with
+/// everyone but the source. Picks the source if there is none (or it went quiet).
+pub(crate) fn defer(state: &AppState, peer: &str, peer_state: &PeerState) -> Option<Duration> {
+    if !catching_up(state) {
+        return None;
+    }
+    let stall = state.timing.catch_up_probe * UNANSWERED_PROBES;
+    if peer_state.reports_ready == Some(false) || peer_state.syncs_with_us == Some(false) {
+        // Not a candidate; its next message re-evaluates this, the timer is a fallback.
+        return Some(stall);
+    }
+
+    let now = Instant::now();
+    let mut source = lock(state);
+    let quiet = match &*source {
+        Some(current) if current.peer == peer => return None,
+        Some(current) if now < current.last_heard + stall => {
+            return Some(current.last_heard + stall - now);
+        }
+        Some(current) => Some(current.peer.clone()),
+        None => None,
+    };
+    match quiet {
+        Some(previous) => info!(
+            peer = %peer,
+            previous = %previous,
+            quiet_secs = stall.as_secs_f32(),
+            "Catch-up source stopped answering; catching up from another peer"
+        ),
+        None => debug!(peer = %peer, "Catching up from this peer first"),
+    }
+    *source = Some(CatchUpSource {
+        peer: peer.to_owned(),
+        last_heard: now,
+        last_probe: None,
+        refusals: 0,
+    });
+    None
+}
+
+/// Any message from `peer` arrived.
+pub(crate) fn heard_from(state: &AppState, peer: &str) {
+    if let Some(current) = lock(state).as_mut().filter(|current| current.peer == peer) {
+        current.last_heard = Instant::now();
+    }
+}
+
+/// What `peer` said about itself, in a sync message (`syncs_with_us: None`) or a
+/// keepalive. Gives up the source if that rules it out.
+pub(crate) fn reported(state: &AppState, peer: &str, ready: bool, syncs_with_us: Option<bool>) {
+    if !catching_up(state) {
+        return;
+    }
+    let reason = {
+        let mut source = lock(state);
+        let Some(current) = source.as_mut().filter(|current| current.peer == peer) else {
+            return;
+        };
+        let reason = if !ready {
+            Some("it isn't ready itself")
+        } else {
+            match syncs_with_us {
+                Some(false) => {
+                    current.refusals += 1;
+                    (current.refusals >= REFUSALS).then_some("it wouldn't sync with us")
+                }
+                Some(true) => {
+                    current.refusals = 0;
+                    None
+                }
+                None => None,
+            }
+        };
+        if reason.is_some() {
+            *source = None;
+        }
+        reason
+    };
+    if let Some(reason) = reason {
+        debug!(peer = %peer, reason, "Giving up the catch-up source; trying another peer");
+        resync_everyone(state);
+    }
+}
+
+/// Whether to PING `peer` now: it is the source and has been quiet for a probe interval.
+pub(crate) fn probe_due(state: &AppState, peer: &str, now: Instant) -> bool {
+    if !catching_up(state) {
+        return false;
+    }
+    let probe = state.timing.catch_up_probe;
+    let mut source = lock(state);
+    let Some(current) = source.as_mut().filter(|current| current.peer == peer) else {
+        return false;
+    };
+    let due = now >= current.last_heard + probe
+        && current.last_probe.is_none_or(|last| now >= last + probe);
+    if due {
+        current.last_probe = Some(now);
+    }
+    due
+}
+
+/// The initial sync is done: no source any more, and the peers we held back only need
+/// to learn our heads now.
+pub(crate) fn finished(state: &AppState, from: &str) {
+    *lock(state) = None;
+    let state = state.clone();
+    let from = from.to_owned();
+    tokio::spawn(async move {
+        let peers: Vec<String> = state.peer_states.read().await.keys().cloned().collect();
+        for other in peers.into_iter().filter(|other| *other != from) {
+            let _ = state.sync_events.send(SyncEvent::PeerNeedsSync(other));
+        }
+    });
+}
+
+/// After giving up a source: let every peer be considered again.
+fn resync_everyone(state: &AppState) {
+    let state = state.clone();
+    tokio::spawn(async move {
+        let peers: Vec<String> = state.peer_states.read().await.keys().cloned().collect();
+        for other in peers {
+            let _ = state.sync_events.send(SyncEvent::PeerNeedsSync(other));
+        }
+    });
+}
diff --git a/src/transport/mod.rs b/src/transport/mod.rs
index eeba47f..aee3ee3 100644
--- a/src/transport/mod.rs
+++ b/src/transport/mod.rs
@@ -1,6 +1,8 @@
 use anyhow::Result;
 use async_trait::async_trait;
+use automerge::ChangeHash;
 use automerge::sync::Message;
+use sha2::{Digest, Sha256};
 use std::boxed::Box;
 use std::future::Future;
 use std::pin::Pin;
@@ -10,13 +12,54 @@ pub use crate::document::replica::Epoch;
 pub mod firewall;
 pub mod mycelium;
 
-pub type HeadsDigest = [u8; 32];
+/// 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;
 
+/// 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)]
+pub enum HeadsSummary {
+    /// Sorted.
+    Heads(Vec<ChangeHash>),
+    /// SHA-256 over the sorted heads.
+    Digest([u8; 32]),
+}
+
+impl HeadsSummary {
+    pub fn of(mut heads: Vec<ChangeHash>) -> Self {
+        heads.sort_unstable();
+        if heads.len() <= MAX_LISTED_HEADS {
+            return Self::Heads(heads);
+        }
+        let mut hasher = Sha256::new();
+        for head in &heads {
+            hasher.update(head.0);
+        }
+        Self::Digest(hasher.finalize().into())
+    }
+
+    /// The heads, if listed.
+    pub fn heads(&self) -> Option<&[ChangeHash]> {
+        match self {
+            Self::Heads(heads) => Some(heads),
+            Self::Digest(_) => None,
+        }
+    }
+}
+
+/// PING and PONG. Addressed to one peer: `syncs_with_you` is about the recipient.
 #[derive(Debug, Clone)]
 pub struct Keepalive {
-    pub heads_digest: HeadsDigest,
-    /// The digest is only comparable between nodes on the same epoch.
+    /// Only comparable between nodes on the same epoch.
+    pub heads: HeadsSummary,
     pub epoch: Epoch,
+    /// The sender is ready (same meaning as `SyncEnvelope::ready`).
+    pub ready: bool,
+    /// The sender would sync with the recipient: counts it as an active member on its
+    /// own epoch.
+    pub syncs_with_you: bool,
 }
 
 #[derive(Debug)]
diff --git a/src/transport/mycelium.rs b/src/transport/mycelium.rs
index 193beaa..9d1fb90 100644
--- a/src/transport/mycelium.rs
+++ b/src/transport/mycelium.rs
@@ -1,5 +1,6 @@
 use anyhow::{Context, Result, anyhow, bail};
 use async_trait::async_trait;
+use automerge::ChangeHash;
 use base64::engine::{Engine, general_purpose::STANDARD as BASE64};
 use reqwest::StatusCode;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
@@ -18,8 +19,8 @@ use tokio::process::{Child, Command};
 use url::Url;
 
 use crate::transport::EncodedMessage;
-use crate::transport::Keepalive;
 use crate::transport::{Epoch, HelloReply, HelloRequest, SyncEnvelope};
+use crate::transport::{HeadsSummary, Keepalive, MAX_LISTED_HEADS};
 use crate::transport::{MessageHandler, ParsedMessage, Transport};
 
 const MYCELIUM_POLL_TIMEOUT_SECS: u64 = 300;
@@ -764,30 +765,63 @@ fn epoch_from_shadow(shadow: EpochShadow) -> Result<Epoch, TransportError> {
     })
 }
 
+/// Either `heads` (listed) or `digest` (32 bytes, `heads` empty) is set.
 #[derive(Serialize, Deserialize)]
 struct KeepaliveShadow {
-    heads_digest: ByteBuf,
-    #[serde(default)]
+    heads: Vec<ByteBuf>,
+    digest: ByteBuf,
     epoch: EpochShadow,
+    ready: bool,
+    syncs_with_you: bool,
 }
 
 fn keepalive_to_bytes(msg: &Keepalive) -> Result<Vec<u8>, rmp_serde::encode::Error> {
+    let (heads, digest) = match &msg.heads {
+        HeadsSummary::Heads(heads) => (
+            heads
+                .iter()
+                .map(|head| ByteBuf::from(head.0.to_vec()))
+                .collect(),
+            ByteBuf::new(),
+        ),
+        HeadsSummary::Digest(digest) => (Vec::new(), ByteBuf::from(digest.to_vec())),
+    };
     rmp_serde::to_vec(&KeepaliveShadow {
-        heads_digest: ByteBuf::from(msg.heads_digest.to_vec()),
+        heads,
+        digest,
         epoch: epoch_to_shadow(&msg.epoch),
+        ready: msg.ready,
+        syncs_with_you: msg.syncs_with_you,
     })
 }
 
 fn bytes_to_keepalive(bytes: &[u8]) -> Result<Keepalive, TransportError> {
     let shadow: KeepaliveShadow = rmp_serde::from_slice(bytes)?;
-    let heads_digest = shadow
-        .heads_digest
-        .as_slice()
-        .try_into()
-        .map_err(|_| TransportError::InvalidPayload("heads digest must be 32 bytes".into()))?;
+    let heads = if shadow.digest.is_empty() {
+        if shadow.heads.len() > MAX_LISTED_HEADS {
+            return Err(TransportError::InvalidPayload(format!(
+                "more than {MAX_LISTED_HEADS} heads listed"
+            )));
+        }
+        let heads = shadow
+            .heads
+            .iter()
+            .map(|head| ChangeHash::try_from(head.as_slice()))
+            .collect::<Result<Vec<_>, _>>()
+            .map_err(|_| TransportError::InvalidPayload("a head must be 32 bytes".into()))?;
+        HeadsSummary::Heads(heads)
+    } else {
+        let digest =
+            shadow.digest.as_slice().try_into().map_err(|_| {
+                TransportError::InvalidPayload("heads digest must be 32 bytes".into())
+            })?;
+        HeadsSummary::Digest(digest)
+    };
     Ok(Keepalive {
-        heads_digest,
+        heads,
         epoch: epoch_from_shadow(shadow.epoch)?,
+        ready: shadow.ready,
+        syncs_with_you: shadow.syncs_with_you,
     })
 }