prefer closer sources to catch up

3 files changed, 49 insertions, 3 deletions

diff --git a/src/main.rs b/src/main.rs
index f5f21fe..15f5c0e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -197,6 +197,7 @@ async fn main() -> Result<()> {
         rtt: Default::default(),
         stale: Default::default(),
         catch_up_source: Default::default(),
+        catch_up_preference: Default::default(),
     };
 
     trace!("Application state initialized");
diff --git a/src/state.rs b/src/state.rs
index 2bc9264..6021baa 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -549,6 +549,8 @@ pub(crate) struct AppState {
 
     /// While the initial sync isn't done: the one peer we catch up from.
     pub(crate) catch_up_source: Arc<std::sync::Mutex<Option<CatchUpSource>>>,
+    /// Since when a closer peer has had a head start to become the source.
+    pub(crate) catch_up_preference: Arc<std::sync::Mutex<Option<Instant>>>,
 }
 
 /// The peer a catching-up node syncs with first (`swarm::source`).
@@ -1029,6 +1031,10 @@ impl Host for AppState {
             .catch_up_source
             .lock()
             .unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
+        *self
+            .catch_up_preference
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
     }
 
     fn subscribe_changes(&self) -> broadcast::Receiver<Arc<ChangeSet>> {
diff --git a/src/swarm/source.rs b/src/swarm/source.rs
index 6be62b7..292cfc8 100644
--- a/src/swarm/source.rs
+++ b/src/swarm/source.rs
@@ -4,8 +4,9 @@
 //! 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
+//! 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
@@ -24,12 +25,23 @@ 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)
 }
@@ -63,14 +75,40 @@ pub(crate) fn defer(state: &AppState, peer: &str, peer_state: &PeerState) -> Opt
         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 mine = rtts.get(peer);
+        let closer = rtts.all().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, "Catching up from this peer first"),
+        None => debug!(peer = %peer, rtt_ms, "Catching up from this peer first"),
     }
     *source = Some(CatchUpSource {
         peer: peer.to_owned(),
@@ -147,6 +185,7 @@ pub(crate) fn probe_due(state: &AppState, peer: &str, now: Instant) -> bool {
 /// 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 {