1 files changed, 112 insertions, 0 deletions
diff --git a/src/swarm/peer/keepalive.rs b/src/swarm/peer/keepalive.rs
new file mode 100644
index 0000000..4fcad7e
--- /dev/null
+++ b/src/swarm/peer/keepalive.rs
@@ -0,0 +1,112 @@
+//! Keepalives (PING / PONG) and the self-healing re-handshake when they show our
+//! documents differ while sync is idle.
+
+use std::sync::atomic::Ordering;
+
+use anyhow::Result;
+use tokio::time::Instant;
+use tracing::warn;
+
+use super::handshake::ensure_handshake;
+use crate::state::{AppState, Liveness, Standing};
+use crate::swarm::sync::SyncEvent;
+use crate::swarm::topology;
+use crate::transport::{self, ParsedMessage};
+
+/// Sends `peer` a PING with our keepalive, noting the time for its round trip.
+pub(crate) async fn send_ping(state: &AppState, peer: &str) -> Result<()> {
+    let keepalive = local_keepalive(state, peer).await;
+    state.lock_rtt().ping_sent(peer, Instant::now());
+    state
+        .transport
+        .send_message(peer, &ParsedMessage::Ping(keepalive))
+        .await
+}
+
+/// Our keepalive for `peer`.
+pub(crate) async fn local_keepalive(state: &AppState, peer: &str) -> transport::Keepalive {
+    // Heads and epoch come from the same locked replica, so they always match. The heads
+    // are cached (`Replica`), so the read lock is enough.
+    let (heads, epoch, epoch_seq) = {
+        let replica = state.doc.read().await;
+        (
+            transport::HeadsSummary::of(replica.heads().to_vec()),
+            replica.epoch,
+            replica.epoch_seq,
+        )
+    };
+    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,
+        epoch_seq,
+        ready: state.is_ready().await,
+        syncs_with_you,
+    }
+}
+
+/// What a peer's keepalive says about it.
+pub(super) 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;
+        let became_ready = remote.ready && peer_state.reports_ready != Some(true);
+        peer_state.reports_ready = Some(remote.ready);
+        peer_state.syncs_with_us = Some(remote.syncs_with_you);
+        peer_state.reported_heads = remote.heads.heads().map(<[_]>::len);
+        drop(peer_state);
+        crate::swarm::relay::pull_if_behind(state, peer, remote).await;
+        if became_ready {
+            // We may have left the opening to it while it wasn't ready, and nothing else
+            // would make us look again.
+            let _ = state
+                .sync_events
+                .send(SyncEvent::PeerNeedsSync(peer.to_owned()));
+        }
+    }
+    crate::swarm::source::reported(state, peer, remote.ready, Some(remote.syncs_with_you));
+}
+
+pub(super) async fn heal_if_diverged(
+    state: &AppState,
+    peer: &str,
+    local: &transport::Keepalive,
+    remote: &transport::Keepalive,
+) {
+    // Heads of documents from different epochs are never equal and never converge.
+    // 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;
+    }
+
+    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
+        return;
+    };
+
+    let heal = {
+        let mut peer_state = peer_state_mutex.lock().await;
+        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();
+        }
+        // While we catch up, we hold back sync with everyone but the source on purpose:
+        // differing heads and idle sync are expected then, not a reason to heal.
+        let catching_up = !state.initial_sync_done.load(Ordering::Acquire);
+        !catching_up
+            && peer_state.needs_heal(&local.heads, &remote.heads, &state.timing, Instant::now())
+    };
+
+    if heal {
+        warn!(peer = %peer, "Documents differ and sync is idle; re-handshaking to restart sync");
+        ensure_handshake(state, peer).await;
+    }
+}