reorganize all the modules

1 files changed, 675 insertions, 0 deletions

diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
new file mode 100644
index 0000000..eb97a7e
--- /dev/null
+++ b/src/swarm/peer.rs
@@ -0,0 +1,675 @@
+//! Peers: handshake (HELLO / HELLO_REPLY), liveness and keepalives, self-healing,
+//! cluster-wide forget records, and dispatch of inbound messages.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use anyhow::{anyhow, bail, Context, Result};
+use automerge::AutoCommit;
+use sha2::{Digest, Sha256};
+use tokio::sync::Mutex;
+use tokio::time::{Instant, MissedTickBehavior};
+use tracing::{debug, info, instrument, trace, warn};
+
+use crate::document::replica::{Epoch, EpochRelation};
+use crate::document::store::{self, StoreError};
+use crate::state::{AppState, HelloState, Liveness, PeerMembership, PeerState};
+use crate::swarm::sync::{handle_sync, SyncEvent};
+use crate::transport::{self, ParsedMessage};
+
+const HELLO_RETRY_INITIAL: Duration = Duration::from_secs(1);
+const HELLO_RETRY_MAX: Duration = Duration::from_secs(30);
+
+const PEER_TICK: Duration = Duration::from_millis(500);
+
+const META_ROOT: &str = "_continuum";
+const FORGET_CHECK_INTERVAL: Duration = Duration::from_secs(5);
+
+// --- forget records ----------------------------------------------------------
+
+fn forgotten_dir() -> Vec<String> {
+    vec![META_ROOT.to_owned(), "forgotten".to_owned()]
+}
+
+pub(crate) fn forgotten_path(pk: &str) -> Vec<String> {
+    let mut path = forgotten_dir();
+    path.push(pk.to_owned());
+    path
+}
+
+fn forgotten_peers(doc: &AutoCommit) -> Vec<String> {
+    store::keys_at(doc, &forgotten_dir()).unwrap_or_default()
+}
+
+async fn apply_forget_records(state: &AppState) {
+    // Before the initial sync we can't judge the records, and we have nothing to delete.
+    if !state.is_ready().await {
+        return;
+    }
+
+    let forgotten = forgotten_peers(&state.doc.read().await.doc);
+    if forgotten.is_empty() {
+        return;
+    }
+
+    let now = Instant::now();
+    let mut cancel = Vec::new();
+
+    for pk in forgotten {
+        if pk == state.local_id() {
+            cancel.push(pk); // we are evidently alive
+            continue;
+        }
+
+        let Some(peer_state_mutex) = state.get_peer_state(&pk).await else {
+            continue; // already forgotten here
+        };
+
+        let alive = peer_state_mutex.lock().await.is_alive(&state.timing, now);
+
+        if alive {
+            cancel.push(pk);
+        } else if state.peer_states.write().await.remove(&pk).is_some() {
+            info!(peer = %pk, "Forgetting peer per cluster-wide record");
+        }
+    }
+
+    if cancel.is_empty() {
+        return;
+    }
+
+    // Individual failures are logged inside; NotFound means another node already
+    // cancelled the record.
+    let _ = state
+        .mutate_doc(move |doc| {
+            for pk in &cancel {
+                match store::delete_path(&mut *doc, &forgotten_path(pk), false) {
+                    Ok(()) => info!(peer = %pk, "Peer is alive; cancelling cluster-wide forget"),
+                    Err(StoreError::NotFound) => {}
+                    Err(error) => {
+                        warn!(peer = %pk, error = %error, "Failed to cancel forget record")
+                    }
+                }
+            }
+            Ok::<(), StoreError>(())
+        })
+        .await;
+}
+
+// --- 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);
+    }
+    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,
+    }
+}
+
+async fn heal_if_diverged(
+    state: &AppState,
+    peer: &str,
+    local: &transport::Keepalive,
+    remote: &transport::Keepalive,
+) {
+    // Digests of documents from different epochs are never equal and never converge.
+    if local.epoch != remote.epoch {
+        return;
+    }
+
+    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
+        return;
+    };
+
+    let heal = peer_state_mutex.lock().await.needs_heal(
+        &local.heads_digest,
+        &remote.heads_digest,
+        &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;
+    }
+}
+
+// --- peer discovery and HELLO ------------------------------------------------
+
+pub(crate) async fn remember_peer(state: &AppState, peer: &str) {
+    if peer.is_empty() {
+        warn!("Ignoring empty peer public key");
+        return;
+    }
+
+    if peer == state.local_id() {
+        debug!("Ignoring own public key");
+        return;
+    }
+
+    let mut peer_states = state.peer_states.write().await;
+
+    if !peer_states.contains_key(peer) {
+        peer_states.entry(peer.to_owned()).or_default();
+        info!(
+            peer = %peer,
+            peer_count = peer_states.len(),
+            "Discovered peer"
+        );
+    } else {
+        trace!(peer = %peer, "Peer already registered");
+    }
+}
+
+pub(crate) async fn send_hello_once(state: &AppState, peer: &str) -> Result<()> {
+    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
+        trace!(peer = %peer, "Peer was forgotten; not sending HELLO");
+        return Ok(());
+    };
+
+    let claimed = {
+        let mut peer_state = peer_state_mutex.lock().await;
+        if peer_state.hello_state == HelloState::NotSent {
+            peer_state.hello_state = HelloState::Sending {
+                next_delay: HELLO_RETRY_INITIAL,
+            };
+            true
+        } else {
+            trace!(peer = %peer, hello_state = ?peer_state.hello_state, "not sending");
+            false
+        }
+    };
+
+    if !claimed {
+        return Ok(());
+    }
+
+    debug!(to = %peer, "Sending HELLO");
+    transmit_hello(state, peer, HELLO_RETRY_INITIAL).await
+}
+
+async fn retry_hello_if_due(state: &AppState, peer: &str) -> Result<()> {
+    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
+        return Ok(());
+    };
+
+    let delay = {
+        let mut peer_state = peer_state_mutex.lock().await;
+        match peer_state.hello_state {
+            HelloState::AwaitingReply {
+                next_attempt,
+                next_delay,
+            } if Instant::now() >= next_attempt => {
+                peer_state.hello_state = HelloState::Sending { next_delay };
+                next_delay
+            }
+            _ => return Ok(()),
+        }
+    };
+
+    debug!(
+        peer = %peer,
+        next_retry_ms = delay.as_millis(),
+        "No HELLO_REPLY yet; re-sending HELLO"
+    );
+    transmit_hello(state, peer, delay).await
+}
+
+/// Sends a HELLO without holding any lock, then schedules the next retry.
+/// Success and failure both end in AwaitingReply: Mycelium accepting the message
+/// says nothing about whether the peer received it.
+async fn transmit_hello(state: &AppState, peer: &str, delay: Duration) -> Result<()> {
+    let result = state
+        .transport
+        .send_message(
+            peer,
+            &ParsedMessage::Hello(transport::HelloRequest {
+                epoch: state.local_epoch(),
+            }),
+        )
+        .await;
+
+    if let Some(peer_state_mutex) = state.get_peer_state(peer).await {
+        let mut peer_state = peer_state_mutex.lock().await;
+
+        if matches!(peer_state.hello_state, HelloState::Sending { .. }) {
+            let (wait, next_delay) = match peer_state.liveness {
+                Liveness::Dormant { .. } => {
+                    (state.timing.dormant_probe, state.timing.dormant_probe)
+                }
+                Liveness::Active => (delay, delay.saturating_mul(2).min(HELLO_RETRY_MAX)),
+            };
+            peer_state.hello_state = HelloState::AwaitingReply {
+                next_attempt: Instant::now() + wait,
+                next_delay,
+            };
+        }
+    }
+
+    result.context("Sending HELLO")
+}
+
+async fn ensure_handshake(state: &AppState, peer: &str) {
+    if let Err(error) = send_hello_once(state, peer).await {
+        warn!(peer = %peer, error = %format!("{error:#}"), "HELLO failed; will retry");
+    }
+}
+
+/// Sends a HELLO_REPLY if one is owed and our cluster ID is known.
+/// Otherwise it stays owed and the peer worker tries again.
+async fn send_owed_hello_reply(state: &AppState, peer: &str) {
+    if state.cluster_id().await.is_none() {
+        return;
+    }
+
+    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
+        return;
+    };
+
+    let owed = std::mem::take(&mut peer_state_mutex.lock().await.reply_owed);
+    if !owed {
+        return;
+    }
+
+    if let Err(error) = send_hello_reply(state, peer).await {
+        warn!(peer = %peer, error = %format!("{error:#}"), "HELLO_REPLY failed; will retry");
+        peer_state_mutex.lock().await.reply_owed = true;
+    }
+}
+
+async fn flush_owed_hello_replies(state: &AppState) {
+    let peers: Vec<String> = state.peer_states.read().await.keys().cloned().collect();
+    for peer in &peers {
+        send_owed_hello_reply(state, peer).await;
+    }
+}
+
+async fn send_hello_reply(state: &AppState, peer: &str) -> Result<()> {
+    let cluster_id = state
+        .cluster_id()
+        .await
+        .ok_or_else(|| anyhow!("Cannot send HELLO_REPLY without an established cluster ID"))?;
+
+    let candidates: Vec<(String, Arc<Mutex<PeerState>>)> = state
+        .peer_states
+        .read()
+        .await
+        .iter()
+        .filter(|(known, _)| known.as_str() != peer)
+        .map(|(known, m)| (known.clone(), Arc::clone(m)))
+        .collect();
+
+    let mut peers = Vec::new();
+    for (known, peer_state_mutex) in candidates {
+        let peer_state = peer_state_mutex.lock().await;
+        if peer_state.liveness == Liveness::Active
+            && peer_state.hello_state == HelloState::Established
+        {
+            peers.push(known);
+        }
+    }
+
+    let reply = transport::HelloReply {
+        cluster_id,
+        peers,
+        epoch: state.local_epoch(),
+    };
+
+    let ready = state.is_ready().await;
+
+    debug!(to = %peer, ready, "Sending HELLO_REPLY");
+
+    state
+        .transport
+        .send_message(peer, &ParsedMessage::HelloReply(reply))
+        .await
+        .context("Sending HELLO_REPLY")?;
+
+    trace!(to = %peer, "HELLO_REPLY sent");
+    Ok(())
+}
+
+async fn discover_peers(state: &AppState, from_peer: &str, advertised_peers: &[String]) {
+    // TODO here we could add a trace t oshwo advertised_peers
+    for discovered_peer in advertised_peers {
+        if discovered_peer == from_peer {
+            continue;
+        }
+
+        remember_peer(state, discovered_peer).await;
+
+        if let Err(error) = send_hello_once(state, discovered_peer).await {
+            warn!(
+                peer = %discovered_peer,
+                error = %format!("{error:#}"),
+                "Failed to initiate HELLO with discovered peer"
+            );
+        }
+    }
+}
+
+// --- peer worker -------------------------------------------------------------
+
+pub(crate) async fn run_peer_worker(state: AppState) {
+    let mut tick = tokio::time::interval(PEER_TICK);
+    tick.set_missed_tick_behavior(MissedTickBehavior::Delay);
+
+    info!("Starting peer maintenance worker");
+
+    let mut last_forget_check: Option<Instant> = None;
+
+    loop {
+        tick.tick().await;
+
+        // While joining, keep handshakes fast; dormancy would slow probes to minutes.
+        let ready = state.is_ready().await;
+
+        let peers: Vec<(String, Arc<Mutex<PeerState>>)> = state
+            .peer_states
+            .read()
+            .await
+            .iter()
+            .map(|(pk, m)| (pk.clone(), Arc::clone(m)))
+            .collect();
+
+        for (peer, peer_state_mutex) in &peers {
+            let now = Instant::now();
+
+            let (became_dormant, ping_due) = {
+                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);
+                (became_dormant, ping_due)
+            };
+
+            if became_dormant {
+                warn!(
+                    peer = %peer,
+                    silent_secs = state.timing.dormant_after.as_secs(),
+                    probe_secs = state.timing.dormant_probe.as_secs(),
+                    "Peer went dormant"
+                );
+            }
+
+            if ping_due {
+                let local = local_keepalive(&state).await;
+                if let Err(error) = state
+                    .transport
+                    .send_message(peer, &ParsedMessage::Ping(local))
+                    .await
+                {
+                    debug!(peer = %peer, error = %format!("{error:#}"), "Keepalive PING failed");
+                }
+            }
+
+            if let Err(error) = retry_hello_if_due(&state, peer).await {
+                warn!(peer = %peer, error = %format!("{error:#}"), "HELLO retry failed");
+            }
+
+            send_owed_hello_reply(&state, peer).await;
+        }
+
+        if last_forget_check.is_none_or(|t| t.elapsed() >= FORGET_CHECK_INTERVAL) {
+            apply_forget_records(&state).await;
+            last_forget_check = Some(Instant::now());
+        }
+    }
+}
+
+// --- peer membership ---------------------------------------------------------
+
+async fn establish_peer(
+    state: &AppState,
+    peer: &str,
+    cluster_id: String,
+) -> Result<Vec<transport::SyncEnvelope>> {
+    let expected_cluster = state.cluster_id().await;
+
+    if let Some(expected) = &expected_cluster {
+        if expected != &cluster_id {
+            bail!(
+                "Incompatible cluster at {peer}: expected {}, received {}",
+                expected,
+                cluster_id
+            );
+        }
+    }
+
+    let peer_state_mutex = state.get_or_create_peer_state(peer).await;
+    let mut peer_state = peer_state_mutex.lock().await;
+
+    if let PeerMembership::Member {
+        cluster_id: existing,
+    } = &peer_state.membership
+    {
+        if existing != &cluster_id {
+            bail!(
+                "Peer {peer} changed cluster identity from {} to {}",
+                existing,
+                cluster_id
+            );
+        }
+    }
+
+    // Lock order: peer mutex -> lifecycle.
+    if expected_cluster.is_none() {
+        state.set_expected_cluster_id(cluster_id.clone()).await?;
+    }
+
+    if matches!(peer_state.membership, PeerMembership::Unknown) {
+        peer_state.membership = PeerMembership::Member {
+            cluster_id: cluster_id.clone(),
+        };
+    }
+
+    // First reply completes the handshake: stop retries and sync from a clean slate.
+    // Duplicate replies (caused by our own retries) don't reset again.
+    if peer_state.hello_state != HelloState::Established {
+        peer_state.hello_state = HelloState::Established;
+        peer_state.reset_sync();
+        info!(peer = %peer, "Handshake complete; sync state reset");
+    }
+
+    let pending = std::mem::take(&mut peer_state.pending_inbound_sync);
+
+    info!(peer = %peer, cluster_id = %cluster_id, pending_messages = pending.len(), "Peer membership established");
+    Ok(pending)
+}
+
+#[instrument(skip(state), fields(peer = %peer))]
+pub(crate) async fn peer_cluster(state: &AppState, peer: &str) -> Option<String> {
+    let peer_state_mutex = state.get_peer_state(peer).await?;
+    let peer_state = peer_state_mutex.lock().await;
+    match &peer_state.membership {
+        PeerMembership::Unknown => None,
+        PeerMembership::Member { cluster_id } => Some(cluster_id.clone()),
+    }
+}
+
+// --- epochs ------------------------------------------------------------------
+
+/// Records the epoch a peer advertised and compares it with ours. Returns true when both
+/// are on the same epoch, the only case in which the documents may be synced.
+async fn observe_peer_epoch(state: &AppState, peer: &str, remote: Epoch) -> bool {
+    let local = state.local_epoch();
+    let relation = local.relation_to(&remote);
+
+    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
+        return relation == EpochRelation::Same;
+    };
+
+    let previous = {
+        let mut peer_state = peer_state_mutex.lock().await;
+        let previous = peer_state.epoch.replace(remote);
+        if previous.is_some_and(|previous| previous != remote) {
+            // The peer switched documents; our sync state describes its old one. Count
+            // the restart as sync activity, so the keepalive heal doesn't re-handshake
+            // (resetting both sides again) while this fresh sync is under way.
+            peer_state.reset_sync();
+            peer_state.last_sync_rx = Some(Instant::now());
+        }
+        previous
+    };
+
+    // Log transitions only, not every keepalive.
+    if previous != Some(remote) {
+        match relation {
+            EpochRelation::Same => {
+                if previous.is_some() {
+                    info!(peer = %peer, epoch = %remote, "Peer is on our epoch again; resuming sync");
+                    let _ = state
+                        .sync_events
+                        .send(SyncEvent::PeerNeedsSync(peer.to_owned()));
+                }
+            }
+            // Expected after a compaction: either side switches or catches up.
+            EpochRelation::PeerAhead | EpochRelation::PeerBehind => {
+                info!(
+                    peer = %peer,
+                    local_epoch = %local,
+                    peer_epoch = %remote,
+                    relation = ?relation,
+                    "Peer is on another epoch; not syncing with it until both are on the same"
+                );
+            }
+            EpochRelation::Diverged => {
+                warn!(
+                    peer = %peer,
+                    local_epoch = %local,
+                    peer_epoch = %remote,
+                    "Peer is on a different epoch of the same generation; not syncing with it"
+                );
+            }
+        }
+    }
+
+    relation == EpochRelation::Same
+}
+
+// --- inbound dispatch --------------------------------------------------------
+
+pub(crate) async fn handle_message(
+    state: &AppState,
+    peer: String,
+    message: transport::ParsedMessage,
+) -> Result<()> {
+    if peer == state.local_id() {
+        debug!("Ignoring message from self");
+        return Ok(());
+    }
+
+    remember_peer(state, &peer).await;
+
+    // Any inbound message proves liveness.
+    let reactivated = state
+        .get_or_create_peer_state(&peer)
+        .await
+        .lock()
+        .await
+        .mark_seen(Instant::now());
+
+    if reactivated {
+        info!(peer = %peer, "Dormant peer is active again; re-handshaking");
+    }
+
+    let same_epoch = observe_peer_epoch(state, &peer, message.epoch()).await;
+
+    //debug!(from = %peer, message = ?message, "Received message");
+
+    match message {
+        transport::ParsedMessage::Hello(_) => {
+            // A HELLO means the peer is (re)starting its handshake with us, e.g. after a
+            // restart. Whatever our sync state believed about it is stale.
+            {
+                let peer_state_mutex = state.get_or_create_peer_state(&peer).await;
+                let mut peer_state = peer_state_mutex.lock().await;
+                peer_state.reset_sync();
+                peer_state.reply_owed = true;
+            }
+
+            // Sends now if we know our cluster ID; otherwise stays owed.
+            send_owed_hello_reply(state, &peer).await;
+
+            // Our own HELLO goes out regardless of whether the reply could be sent.
+            if let Err(error) = send_hello_once(state, &peer).await {
+                warn!(peer = %peer, error = %format!("{error:#}"), "HELLO failed; will retry");
+            }
+
+            // Restart sync with the fresh state (no-op until the peer is a validated member).
+            let _ = state
+                .sync_events
+                .send(SyncEvent::PeerNeedsSync(peer.clone()));
+        }
+        transport::ParsedMessage::HelloReply(reply) => {
+            let pending = establish_peer(state, &peer, reply.cluster_id.clone()).await?;
+
+            discover_peers(state, &peer, &reply.peers).await;
+
+            // Our cluster ID may have just become known: pay any replies we owe.
+            flush_owed_hello_replies(state).await;
+
+            if !pending.is_empty() {
+                debug!(
+                    to = %peer,
+                    pending_messages = pending.len(),
+                    "Replaying pending sync messages"
+                );
+                for envelope in pending {
+                    if let Err(error) = handle_sync(state, &peer, envelope).await {
+                        warn!(peer = %peer, error = %format!("{error:#}"), "Discarding invalid pending sync envelope");
+                    }
+                }
+            }
+
+            state
+                .sync_events
+                .send(SyncEvent::PeerValidated(peer.clone()))
+                .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?;
+        }
+        transport::ParsedMessage::Sync(envelope) => {
+            if reactivated {
+                ensure_handshake(state, &peer).await;
+            }
+            if !same_epoch {
+                debug!(peer = %peer, epoch = %envelope.epoch, "Dropping sync message from another epoch");
+                return Ok(());
+            }
+            handle_sync(state, &peer, envelope).await?;
+        }
+        transport::ParsedMessage::Ping(remote) => {
+            if reactivated || peer_cluster(state, &peer).await.is_none() {
+                ensure_handshake(state, &peer).await;
+            }
+            let local = local_keepalive(state).await;
+            heal_if_diverged(state, &peer, &local, &remote).await;
+            if let Err(error) = state
+                .transport
+                .send_message(&peer, &ParsedMessage::Pong(local))
+                .await
+            {
+                debug!(peer = %peer, error = %format!("{error:#}"), "PONG failed");
+            }
+        }
+        transport::ParsedMessage::Pong(remote) => {
+            if reactivated || peer_cluster(state, &peer).await.is_none() {
+                ensure_handshake(state, &peer).await;
+            }
+            let local = local_keepalive(state).await;
+            heal_if_diverged(state, &peer, &local, &remote).await;
+        }
+    }
+
+    Ok(())
+}