sync with neighbours

6 files changed, 171 insertions, 4 deletions

diff --git a/src/api/peers.rs b/src/api/peers.rs
index 6a9a271..d381395 100644
--- a/src/api/peers.rs
+++ b/src/api/peers.rs
@@ -60,6 +60,8 @@ pub(crate) struct PeerStatus {
     /// Number of heads in its last keepalive; absent before one, or when it had more
     /// than 32 (then only their digest is sent).
     heads: Option<usize>,
+    /// We keep it alive and sync with it on our own; other peers we only answer.
+    neighbour: bool,
 }
 
 #[derive(Debug, Serialize, ToSchema)]
@@ -106,6 +108,7 @@ pub(crate) async fn list_peers(
     let mut peers = Vec::with_capacity(snapshot.len());
 
     for (public_key, peer_state_mutex) in snapshot {
+        let neighbour = crate::swarm::topology::is_neighbour(&state, &public_key);
         let ps = peer_state_mutex.lock().await;
 
         let (status, dormant_for_secs) = match ps.liveness {
@@ -129,6 +132,7 @@ pub(crate) async fn list_peers(
             epoch: ps.epoch.map(EpochInfo::from),
             same_epoch: ps.epoch.map(|epoch| epoch == local_epoch),
             heads: ps.reported_heads,
+            neighbour,
         });
     }
 
diff --git a/src/main.rs b/src/main.rs
index 6082e67..a688e04 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -91,6 +91,11 @@ struct Args {
     #[arg(long, value_name = "SECONDS", default_value_t = 60)]
     keepalive_secs: u64,
 
+    /// Neighbours each node keeps (at least 2). Swarms of up to this many + 1 nodes are
+    /// a full mesh; larger ones sync over several hops.
+    #[arg(long, value_name = "COUNT", default_value_t = 12)]
+    neighbours: usize,
+
     /// Test hook: start on a non-genesis epoch of this generation (with a fixed dummy
     /// id), to check that nodes on different epochs refuse to sync.
     #[arg(long, hide = true, requires = "debug_api", value_name = "GENERATION")]
@@ -186,6 +191,8 @@ async fn main() -> Result<()> {
         remake: Arc::new(remake::RemakeState::default()),
         swarm_key,
         run_id: uuid::Uuid::new_v4().into_bytes(),
+        neighbours_k: args.neighbours.max(2),
+        neighbours: Default::default(),
         stale: Default::default(),
         catch_up_source: Default::default(),
     };
diff --git a/src/state.rs b/src/state.rs
index cfa5ddf..64fe6a4 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -321,6 +321,14 @@ impl PeerState {
         false
     }
 
+    /// No more HELLO retries (the peer isn't our neighbour any more). A HELLO in flight
+    /// finishes; a later handshake starts over.
+    pub(crate) fn stop_hello_retries(&mut self) {
+        if matches!(self.hello_state, HelloState::AwaitingReply { .. }) {
+            self.hello_state = HelloState::NotSent;
+        }
+    }
+
     /// Moves a silent peer to Dormant. Returns true on the transition.
     pub(crate) fn check_dormancy(&mut self, timing: &Timing, now: Instant) -> bool {
         if self.liveness != Liveness::Active || self.quiet_for(now) < timing.dormant_after {
@@ -495,6 +503,12 @@ pub(crate) struct AppState {
     /// Random per run of this node; our HELLOs carry it and replies must echo it.
     pub(crate) run_id: transport::RunId,
 
+    /// Target number of neighbours (`--neighbours`); up to this many + 1 members the
+    /// swarm is a full mesh.
+    pub(crate) neighbours_k: usize,
+    /// Our neighbours, once known (`swarm::topology`).
+    pub(crate) neighbours: Arc<std::sync::RwLock<Option<crate::swarm::topology::NeighbourView>>>,
+
     /// While rejoining a newer epoch, the previous document, served read-only (and
     /// marked stale) instead of answering 503.
     pub(crate) stale: Arc<std::sync::RwLock<Option<StaleView>>>,
@@ -751,6 +765,10 @@ impl AppState {
         let local = self.local_epoch();
         let mut announce = Vec::new();
         for (id, peer_state_mutex) in &peers {
+            // Others learn it from their own neighbours.
+            if !crate::swarm::topology::is_neighbour(self, id) {
+                continue;
+            }
             let mut peer_state = peer_state_mutex.lock().await;
             // What it said about syncing with us was about our old epoch.
             peer_state.syncs_with_us = None;
diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index 3db4710..7ab704f 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -15,6 +15,7 @@ use crate::document::replica::{Epoch, EpochRelation};
 use crate::document::store::{self, StoreError};
 use crate::state::{AppState, HelloState, Liveness, PeerState, Standing};
 use crate::swarm::sync::{SyncEvent, handle_sync};
+use crate::swarm::topology;
 use crate::transport::{self, ParsedMessage};
 
 const PEER_TICK: Duration = Duration::from_millis(500);
@@ -150,7 +151,8 @@ async fn heal_if_diverged(
     remote: &transport::Keepalive,
 ) {
     // Heads of documents from different epochs are never equal and never converge.
-    if local.epoch != remote.epoch {
+    // And a peer that isn't our neighbour isn't ours to keep in step.
+    if local.epoch != remote.epoch || !topology::is_neighbour(state, peer) {
         return;
     }
 
@@ -320,6 +322,7 @@ async fn send_hello_reply(state: &AppState, peer: &str, run_id: transport::RunId
         let peer_state = peer_state_mutex.lock().await;
         if peer_state.liveness == Liveness::Active
             && peer_state.hello_state == HelloState::Established
+            && topology::is_neighbour(state, &known)
         {
             peers.push(known);
         }
@@ -361,6 +364,9 @@ async fn discover_peers(state: &AppState, from_peer: &str, advertised_peers: &[S
 
         remember_peer(state, discovered_peer).await;
 
+        if !topology::is_neighbour(state, discovered_peer) {
+            continue;
+        }
         if let Err(error) = send_hello_once(state, discovered_peer).await {
             warn!(
                 peer = %discovered_peer,
@@ -371,6 +377,26 @@ async fn discover_peers(state: &AppState, from_peer: &str, advertised_peers: &[S
     }
 }
 
+/// New neighbours get our HELLO (we may not have known them before: the member list is
+/// how a partial mesh discovers them). Former neighbours get no more HELLOs, keepalives
+/// or sync openings from us; their session is reset, so we don't keep pushing changes
+/// to them either. They are still answered if they contact us.
+async fn apply_neighbour_change(state: &AppState, change: topology::NeighbourChange) {
+    for peer in &change.added {
+        remember_peer(state, peer).await;
+        if let Err(error) = send_hello_once(state, peer).await {
+            warn!(peer = %peer, error = %format!("{error:#}"), "HELLO to new neighbour failed; will retry");
+        }
+    }
+    for peer in &change.removed {
+        if let Some(peer_state) = state.get_peer_state(peer).await {
+            let mut peer_state = peer_state.lock().await;
+            peer_state.stop_hello_retries();
+            peer_state.reset_sync();
+        }
+    }
+}
+
 // --- peer worker -------------------------------------------------------------
 
 pub(crate) async fn run_peer_worker(state: AppState) {
@@ -397,11 +423,17 @@ pub(crate) async fn run_peer_worker(state: AppState) {
 
         for (peer, peer_state_mutex) in &peers {
             let now = Instant::now();
+            // Only neighbours are kept alive and judged; others we only answer.
+            let neighbour = topology::is_neighbour(&state, peer);
 
             let (became_dormant, ping_due, unanswered) = {
                 let mut peer_state = peer_state_mutex.lock().await;
-                let became_dormant = ready && peer_state.check_dormancy(&state.timing, now);
-                let ping_due = peer_state.take_keepalive_due(&state.timing, now);
+                let became_dormant =
+                    ready && neighbour && peer_state.check_dormancy(&state.timing, now);
+                let ping_due = neighbour && peer_state.take_keepalive_due(&state.timing, now);
+                if !neighbour {
+                    peer_state.stop_hello_retries();
+                }
                 let unanswered = peer_state.take_overdue_sync_reply(now);
                 (became_dormant, ping_due, unanswered)
             };
@@ -443,6 +475,9 @@ pub(crate) async fn run_peer_worker(state: AppState) {
         if last_forget_check.is_none_or(|t| t.elapsed() >= FORGET_CHECK_INTERVAL) {
             apply_forget_records(&state).await;
             crate::swarm::members::ensure_registered(&state).await;
+            if let Some(change) = topology::refresh(&state).await {
+                apply_neighbour_change(&state, change).await;
+            }
             last_forget_check = Some(Instant::now());
         }
     }
diff --git a/src/swarm/sync.rs b/src/swarm/sync.rs
index e05d275..b8d16aa 100644
--- a/src/swarm/sync.rs
+++ b/src/swarm/sync.rs
@@ -85,6 +85,14 @@ async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> {
             return Ok(());
         }
 
+        if peer_state.sync.state.their_heads.is_none()
+            && !crate::swarm::topology::is_neighbour(state, peer)
+        {
+            // We only answer peers that aren't our neighbours.
+            trace!("Not opening sync with a peer that isn't our neighbour");
+            return Ok(());
+        }
+
         swarm_id.clone()
     };
 
@@ -444,8 +452,12 @@ pub(crate) async fn run_sync_worker(state: AppState, mut rx: mpsc::UnboundedRece
                 peers_guard.keys().cloned().collect()
             };
 
+            // Our own changes go to our neighbours; they pass them on. Others only get
+            // answers to what they send us.
             for peer in peer_keys {
-                if peer_swarm(&state, &peer).await.is_some() {
+                if peer_swarm(&state, &peer).await.is_some()
+                    && crate::swarm::topology::is_neighbour(&state, &peer)
+                {
                     dirty_peers.insert(peer);
                 }
             }
diff --git a/src/swarm/topology.rs b/src/swarm/topology.rs
index 0ec8b14..f097cdd 100644
--- a/src/swarm/topology.rs
+++ b/src/swarm/topology.rs
@@ -18,6 +18,97 @@
 use std::collections::{BTreeMap, BTreeSet};
 
 use sha2::{Digest, Sha256};
+use tracing::info;
+
+use crate::state::AppState;
+use crate::swarm::members;
+
+/// This node's current neighbours, and the member list they were computed from.
+#[derive(Debug, Default)]
+pub(crate) struct NeighbourView {
+    members: Vec<String>,
+    k: usize,
+    neighbours: BTreeSet<String>,
+}
+
+/// Whether we initiate contact with `peer`: HELLOs, keepalives, dormancy checks, sync
+/// openings. Everyone is a neighbour until the view is known (while joining, or before
+/// any member registered), so joining works as in a full mesh. We always *answer*
+/// anyone who contacts us.
+pub(crate) fn is_neighbour(state: &AppState, peer: &str) -> bool {
+    match &*state
+        .neighbours
+        .read()
+        .unwrap_or_else(|poisoned| poisoned.into_inner())
+    {
+        Some(view) => view.neighbours.contains(peer),
+        None => true,
+    }
+}
+
+/// How the neighbours changed.
+pub(crate) struct NeighbourChange {
+    pub(crate) added: Vec<String>,
+    pub(crate) removed: Vec<String>,
+}
+
+/// Recomputes the neighbours once ready, if the member list (or `k`) changed.
+pub(crate) async fn refresh(state: &AppState) -> Option<NeighbourChange> {
+    if !state.is_ready().await {
+        return None;
+    }
+    let current = members::members(&state.doc.read().await.doc);
+    let k = state.neighbours_k;
+    let previous = {
+        let view = state
+            .neighbours
+            .read()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        if view
+            .as_ref()
+            .is_some_and(|view| view.members == current && view.k == k)
+        {
+            return None;
+        }
+        view.as_ref().map(|view| view.neighbours.clone())
+    };
+
+    // Before the first view everyone known counted as a neighbour.
+    let previous = match previous {
+        Some(previous) => previous,
+        None => state.peer_states.read().await.keys().cloned().collect(),
+    };
+    let computed = neighbours(&current, state.local_id(), k);
+    let change = NeighbourChange {
+        added: computed
+            .iter()
+            .filter(|m| !previous.contains(*m))
+            .cloned()
+            .collect(),
+        removed: previous
+            .iter()
+            .filter(|m| !computed.contains(*m))
+            .cloned()
+            .collect(),
+    };
+    info!(
+        members = current.len(),
+        k,
+        neighbours = computed.len(),
+        added = change.added.len(),
+        removed = change.removed.len(),
+        "Neighbours computed from the member list"
+    );
+    *state
+        .neighbours
+        .write()
+        .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(NeighbourView {
+        members: current,
+        k,
+        neighbours: computed,
+    });
+    Some(change)
+}
 
 const RING_LABEL: &[u8] = b"continuum/topology-ring/1";
 const PAIR_LABEL: &[u8] = b"continuum/topology-pair/1";