1 files changed, 210 insertions, 0 deletions
diff --git a/src/swarm/source.rs b/src/swarm/source.rs
new file mode 100644
index 0000000..378c40a
--- /dev/null
+++ b/src/swarm/source.rs
@@ -0,0 +1,210 @@
+//! 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 one suitable peer becomes the source, preferably the closest by round-trip time
+//! (a clearly closer one gets a 1 s head start to claim it, else the first to ask);
+//! 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;
+
+/// How long a clearly closer peer (by round-trip time) gets to claim the source role
+/// before whoever asks next takes it.
+const CLOSER_PEER_HEAD_START: Duration = Duration::from_secs(1);
+
+/// 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 lock_preference(state: &AppState) -> std::sync::MutexGuard<'_, Option<Instant>> {
+    state
+        .catch_up_preference
+        .lock()
+        .unwrap_or_else(|poisoned| poisoned.into_inner())
+}
+
+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,
+    };
+
+    // Prefer the closest peer: the source sends us everything we lack, so its distance
+    // matters most. If another peer we measured is clearly closer (by a fifth), it gets
+    // a short head start to claim the role; it is asked in the same sync pass. After
+    // that, whoever asks first takes it (the closer one may not be suitable after all).
+    let rtt = {
+        let rtts = state.lock_rtt();
+        let fresh_for = state.timing.rtt_fresh_for();
+        let mine = rtts.effective(peer, now, fresh_for);
+        let closer = rtts.all_effective(now, fresh_for).any(|(other, rtt)| {
+            other != peer
+                && quiet.as_deref() != Some(other)
+                && mine.is_none_or(|mine| rtt * 5 < mine * 4)
+        });
+        if closer {
+            let mut preference = lock_preference(state);
+            let since = *preference.get_or_insert(now);
+            if now < since + CLOSER_PEER_HEAD_START {
+                return Some(since + CLOSER_PEER_HEAD_START - now);
+            }
+        }
+        mine
+    };
+    *lock_preference(state) = None;
+    let rtt_ms = rtt.map(|rtt| rtt.as_millis() as u64);
+
+    match quiet {
+        Some(previous) => info!(
+            peer = %peer,
+            previous = %previous,
+            rtt_ms,
+            quiet_secs = stall.as_secs_f32(),
+            "Catch-up source stopped answering; catching up from another peer"
+        ),
+        None => debug!(peer = %peer, rtt_ms, "Catching up from this peer first"),
+    }
+    state.stats.catch_up_source_chosen();
+    *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;
+    *lock_preference(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.
+pub(crate) 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));
+        }
+    });
+}