diff options
Diffstat (limited to 'src/swarm')
| -rw-r--r-- | src/swarm/bootstrap.rs | 104 | ||||
| -rw-r--r-- | src/swarm/mod.rs | 6 | ||||
| -rw-r--r-- | src/swarm/peer.rs | 675 | ||||
| -rw-r--r-- | src/swarm/sync.rs | 564 |
4 files changed, 1349 insertions, 0 deletions
diff --git a/src/swarm/bootstrap.rs b/src/swarm/bootstrap.rs new file mode 100644 index 0000000..2907e93 --- /dev/null +++ b/src/swarm/bootstrap.rs @@ -0,0 +1,104 @@ +//! Joining a cluster (or creating one as the genesis node). + +use std::sync::atomic::Ordering; +use std::time::Duration; + +use anyhow::{bail, Result}; +use tracing::{debug, error, info, warn}; +use uuid::Uuid; + +use crate::state::{AppState, PeerMembership}; +use crate::swarm::peer::{remember_peer, send_hello_once}; + +const BOOTSTRAP_TIMEOUT: Duration = Duration::from_secs(300); +const BOOTSTRAP_RETRY: Duration = Duration::from_millis(500); + +pub(crate) async fn bootstrap(state: &AppState, seeds: &[String]) -> Result<()> { + if seeds.is_empty() { + info!("Initializing genesis node (no seed peers configured)"); + state.initial_sync_done.store(true, Ordering::Release); + state.mark_ready(Uuid::new_v4().to_string()).await?; + return Ok(()); + } + + info!( + timeout_seconds = BOOTSTRAP_TIMEOUT.as_secs(), + "Joining existing cluster" + ); + + let mut peers = Vec::new(); + + for seed in seeds { + match state.transport.resolve_peer(seed).await { + Ok(peer) => { + debug!(seed = %seed, peer = %peer, "Resolved seed to peer"); + peers.push(peer) + } + Err(e) => warn!(seed = %seed, error = %e, "Ignoring seed argument"), + } + } + + if peers.is_empty() { + bail!("None of the configured seeds could be resolved"); + } + + for peer in peers { + remember_peer(state, &peer).await; + if let Err(error) = send_hello_once(state, &peer).await { + warn!(peer = %peer, error = %format!("{error:#}"), "Initial HELLO failed; will retry"); + } + } + + let deadline = tokio::time::Instant::now() + BOOTSTRAP_TIMEOUT; + loop { + if let Some(cluster_id) = state.cluster_id().await { + // A node joining a compacted cluster rejoins the newer epoch first; it is + // ready only once that rejoin verified the epoch. + if state.initial_sync_done.load(Ordering::Acquire) && !state.compaction.is_rejoining() { + state.mark_ready(cluster_id.clone()).await?; + info!(cluster_id = %cluster_id, "Node is ready"); + return Ok(()); + } + } + + if tokio::time::Instant::now() >= deadline { + let mut peer_diagnostics: Vec<String> = Vec::new(); + + let peer_states = state.peer_states.read().await.clone(); + + for (peer, peer_state_mutex) in &peer_states { + let peer_state = peer_state_mutex.lock().await; + let validated_cluster = match &peer_state.membership { + PeerMembership::Unknown => None, + PeerMembership::Member { cluster_id } => Some(cluster_id.as_str()), + }; + peer_diagnostics.push(format!( + "{peer}: validated_cluster={:?}, epoch={:?}, pending_inbound_sync={}", + validated_cluster, + peer_state.epoch, + peer_state.pending_inbound_sync.len() + )) + } + let local_cluster = state.cluster_id().await; + let local_epoch = state.local_epoch(); + let initial_sync_done = state.initial_sync_done.load(Ordering::Acquire); + error!( + local_cluster = ?local_cluster, + local_epoch = %local_epoch, + initial_sync_done, + peers = %peer_diagnostics.join("; "), + timeout_seconds = BOOTSTRAP_TIMEOUT.as_secs(), + "Bootstrap timed out" + ); + bail!( + "Bootstrap timed out after {:?}; local_cluster={:?}; local_epoch={}; initial_sync_done={}; peers=[{}]. Refusing to create a genesis document.", + BOOTSTRAP_TIMEOUT, + local_cluster, + local_epoch, + initial_sync_done, + peer_diagnostics.join("; ") + ); + } + tokio::time::sleep(BOOTSTRAP_RETRY).await; + } +} diff --git a/src/swarm/mod.rs b/src/swarm/mod.rs new file mode 100644 index 0000000..7316146 --- /dev/null +++ b/src/swarm/mod.rs @@ -0,0 +1,6 @@ +//! The swarm of peers: bootstrap, handshake, keepalive and dormancy (`peer`), and +//! Automerge synchronization with every member (`sync`). + +pub(crate) mod bootstrap; +pub(crate) mod peer; +pub(crate) mod sync; 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(()) +} diff --git a/src/swarm/sync.rs b/src/swarm/sync.rs new file mode 100644 index 0000000..115d5c0 --- /dev/null +++ b/src/swarm/sync.rs @@ -0,0 +1,564 @@ +//! Automerge synchronization: the sync worker, delivery with retries, inbound sync +//! messages, and the change publisher that feeds incremental streams. + +use std::collections::HashSet; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{anyhow, bail, Context, Result}; +use automerge::sync::SyncDoc; +use tokio::sync::mpsc; +use tokio::time::Instant; +use tracing::{debug, info, trace, warn}; + +use crate::state::{AppState, Liveness, PeerMembership, PeerState, PendingOutbound}; +use crate::swarm::peer::{peer_cluster, send_hello_once}; +use crate::transport::{self, ParsedMessage}; + +const INITIAL_SYNC_RETRY: Duration = Duration::from_millis(500); +const MAX_SYNC_RETRY: Duration = Duration::from_secs(30); + +#[derive(Debug)] +pub(crate) enum SyncEvent { + PeerValidated(String), + DocumentChanged, + MessageReceived(String), + PeerNeedsSync(String), + RetryDelivery(String), +} + +/// Store paths touched by one batch of document changes. +#[derive(Debug)] +pub(crate) struct ChangeSet { + pub(crate) paths: Vec<Vec<String>>, +} + +async fn deliver_sync_to_peer(state: &AppState, peer: &str, retry_delivery: bool) -> Result<()> { + let Some(peer_state_mutex) = state.get_peer_state(peer).await else { + trace!("Peer was forgotten; skipping synchronization"); + return Ok(()); + }; + + let validated_cluster = { + let peer_state = peer_state_mutex.lock().await; + + if matches!(peer_state.liveness, Liveness::Dormant { .. }) { + // Reactivation re-handshakes, and PeerValidated then restarts sync. + trace!("Skipping synchronization with dormant peer"); + return Ok(()); + } + + if peer_state + .epoch + .is_some_and(|epoch| epoch != state.local_epoch()) + { + // Unrelated document histories must never be synced. + trace!("Skipping synchronization with peer on another epoch"); + return Ok(()); + } + + match &peer_state.membership { + PeerMembership::Unknown => return Ok(()), + PeerMembership::Member { cluster_id } => cluster_id.clone(), + } + }; + + let cluster_id = state + .cluster_id() + .await + .ok_or_else(|| anyhow!("Cannot sync without an established cluster ID"))?; + + if validated_cluster != cluster_id { + bail!("Peer {peer} belongs to a different cluster"); + } + + let existing = { + let mut guard = peer_state_mutex.lock().await; + let peer_state = &mut *guard; + match &peer_state.pending_outbound_sync { + Some(entry) if retry_delivery => Some((entry.message.clone(), entry.next_attempt)), + Some(_) => { + peer_state.needs_generate = true; + return Ok(()); + } + None => None, + } + }; + + let encoded = match existing { + Some((message, next_attempt)) => { + if Instant::now() < next_attempt { + return Ok(()); + } + message + } + None if retry_delivery => return Ok(()), + None => { + let doc = Arc::clone(&state.doc); + let psm = Arc::clone(&peer_state_mutex); + let transport = Arc::clone(&state.transport); + let peer_owned = peer.to_owned(); + + let generated = tokio::task::spawn_blocking( + move || -> Result<Option<transport::EncodedMessage>> { + // Label the message with the epoch of the document it was generated + // from, read under the same lock. + let (epoch, message) = { + let mut replica = doc.blocking_write(); + let mut peer_state = psm.blocking_lock(); + peer_state.needs_generate = false; + let epoch = replica.epoch; + let sync_doc = replica.doc.sync(); + ( + epoch, + sync_doc.generate_sync_message(&mut peer_state.sync_state), + ) + }; + let Some(message) = message else { + return Ok(None); + }; + let envelope = transport::SyncEnvelope { + cluster_id, + epoch, + message, + }; + transport + .encode(&peer_owned, &ParsedMessage::Sync(envelope)) + .map(Some) + }, + ) + .await + .context("Sync generation task panicked")??; + + let Some(encoded) = generated else { + return Ok(()); + }; + + peer_state_mutex.lock().await.pending_outbound_sync = Some(PendingOutbound { + message: encoded.clone(), + retry_delay: INITIAL_SYNC_RETRY, + next_attempt: Instant::now(), + }); + + encoded + } + }; + + // No state locks are held across the send. + match state.transport.send_encoded(&encoded).await { + Ok(()) => { + let regenerate = { + let mut peer_state = peer_state_mutex.lock().await; + peer_state.pending_outbound_sync = None; + std::mem::take(&mut peer_state.needs_generate) + }; + + trace!(regenerate, "Outbound sync message accepted by Mycelium"); + + if regenerate { + state + .sync_events + .send(SyncEvent::PeerNeedsSync(peer.to_owned())) + .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?; + } + + Ok(()) + } + Err(error) => { + let delay = { + let mut peer_state = peer_state_mutex.lock().await; + + let Some(ref mut entry) = peer_state.pending_outbound_sync else { + debug!("Pending envelope discarded by a sync reset during delivery"); + return Ok(()); + }; + + entry.retry_delay = entry.retry_delay.saturating_mul(2).min(MAX_SYNC_RETRY); + entry.next_attempt = Instant::now() + entry.retry_delay; + entry.retry_delay + }; + + warn!( + retry_ms = delay.as_millis(), + error = %format!("{error:#}"), + "Sync delivery failed; retaining the same outbound envelope" + ); + + let tx = state.sync_events.clone(); + let peer = peer.to_owned(); + tokio::spawn(async move { + tokio::time::sleep(delay).await; + let _ = tx.send(SyncEvent::RetryDelivery(peer)); + }); + + Ok(()) + } + } +} + +pub(crate) async fn handle_sync( + state: &AppState, + peer: &str, + envelope: transport::SyncEnvelope, +) -> Result<()> { + debug!(peer = %peer, "Processing sync envelope"); + + let peer_state_mutex = state.get_or_create_peer_state(peer).await; + + let needs_hello = { + let mut peer_state = peer_state_mutex.lock().await; + + match &peer_state.membership { + PeerMembership::Unknown => { + let duplicate = peer_state.pending_inbound_sync.iter().any(|pending| { + pending.message == envelope.message + && pending.cluster_id == envelope.cluster_id + && pending.epoch == envelope.epoch + }); + if !duplicate { + peer_state.pending_inbound_sync.push(envelope.clone()); + info!( + peer = %peer, + pending_count = peer_state.pending_inbound_sync.len(), + "Queued sync envelope pending peer membership" + ); + } else { + trace!(peer = %peer, "Ignoring duplicate pending sync envelope"); + } + true + } + PeerMembership::Member { cluster_id } if cluster_id != &envelope.cluster_id => { + bail!( + "Sync cluster mismatch for peer {peer}: validated_cluster={cluster_id}, envelope_cluster={}", + envelope.cluster_id + ); + } + PeerMembership::Member { .. } => false, + } + }; + + if needs_hello { + // The peer treats us as a member but we don't know it, typically because we + // restarted. Our HELLO makes it reset its sync state and reply. + if let Err(error) = send_hello_once(state, peer).await { + warn!(peer = %peer, error = %format!("{error:#}"), "HELLO to unknown sync sender failed; will retry"); + } + return Ok(()); + } + + let expected_cluster = state + .cluster_id() + .await + .ok_or_else(|| anyhow!("Missing expected cluster ID after peer membership"))?; + if envelope.cluster_id != expected_cluster { + bail!( + "Incompatible cluster from {peer}: expected {}, received {}", + expected_cluster, + envelope.cluster_id + ); + } + + let epoch_label = envelope.epoch; + + let applied = { + let doc = Arc::clone(&state.doc); + let peer_state_mutex = Arc::clone(&peer_state_mutex); + let check_caught_up = !state.initial_sync_done.load(Ordering::Acquire); + let epoch = envelope.epoch; + let message = envelope.message; + let peer = peer.to_owned(); + + // Some(caught_up) when applied, None when the message belongs to another epoch. + tokio::task::spawn_blocking(move || -> Result<Option<bool>> { + let mut replica = doc.blocking_write(); + + // Authoritative check, under the same lock as the apply: the document may + // have switched epochs since the message was accepted. + if replica.epoch != epoch { + return Ok(None); + } + + let mut peer_state = peer_state_mutex.blocking_lock(); + + let heads_before = replica.doc.get_heads(); + let their_heads_before = peer_state.sync_state.their_heads.clone(); + + replica + .doc + .sync() + .receive_sync_message(&mut peer_state.sync_state, message) + .context("Automerge synchronization failed")?; + peer_state.last_sync_rx = Some(Instant::now()); + + let heads_after = replica.doc.get_heads(); + let unproductive = heads_before == heads_after + && peer_state.sync_state.their_heads == their_heads_before + && peer_state.sync_state.their_heads.as_ref() != Some(&heads_after); + break_sync_stall(&peer, &mut peer_state, unproductive); + + Ok(Some( + check_caught_up + && peer_state + .sync_state + .their_heads + .as_ref() + .is_some_and(|heads| { + heads + .iter() + .all(|h| replica.doc.get_change_by_hash(h).is_some()) + }), + )) + }) + .await + .context("Sync task panicked")?? + }; + + let Some(caught_up) = applied else { + debug!(peer = %peer, epoch = %epoch_label, "Dropped sync message from another epoch"); + return Ok(()); + }; + + if caught_up && !state.initial_sync_done.swap(true, Ordering::AcqRel) { + info!(peer = %peer, "Initial sync complete"); + } + + state.notify_store_changed(); + + state + .sync_events + .send(SyncEvent::MessageReceived(peer.to_owned())) + .map_err(|error| anyhow!("Failed to notify sync worker: {error}"))?; + + Ok(()) +} + +/// Exchanges in a row that changed nothing on either side before `sent_hashes` is +/// cleared, and before a full two-sided reset (re-handshake). +const STALL_RESEND_NEEDED: u32 = 2; +const STALL_CLEAR_SENT: u32 = 8; +const STALL_RESET: u32 = 32; + +/// Breaks sync livelocks. +/// +/// Automerge never resends a change it believes the peer has received (`sent_hashes`). +/// If that change was lost on the way (a dropped or undeliverable message, an epoch +/// check, a reset of only one side's state), the two sides keep exchanging messages +/// with differing heads and no changes, as fast as the transport allows. The keepalive +/// heal never fires then, because sync traffic keeps flowing. +/// +/// * A change the peer keeps asking for is sent again. +/// * After a few unproductive exchanges, forget what we sent: the next message carries +/// everything the peer's bloom filter lacks. +/// * If that doesn't help either, reset both sides through a fresh handshake. +fn break_sync_stall(peer: &str, peer_state: &mut PeerState, unproductive: bool) { + if !unproductive { + peer_state.stalled_exchanges = 0; + return; + } + peer_state.stalled_exchanges += 1; + + let sync_state = &mut peer_state.sync_state; + + // Not on the first unproductive exchange: when messages cross, the peer asks for a + // change that is merely still in flight, and resending it would only duplicate it. + if peer_state.stalled_exchanges >= STALL_RESEND_NEEDED { + if let Some(need) = &sync_state.their_need { + let before = sync_state.sent_hashes.len(); + sync_state.sent_hashes.retain(|hash| !need.contains(hash)); + let resend = before - sync_state.sent_hashes.len(); + if resend > 0 { + debug!(peer = %peer, resend, "Peer still needs changes we had sent; sending them again"); + } + } + } + + match peer_state.stalled_exchanges { + STALL_CLEAR_SENT => { + warn!( + peer = %peer, + sent = sync_state.sent_hashes.len(), + their_heads = ?sync_state.their_heads, + shared_heads = ?sync_state.shared_heads, + "Sync is stalled; forgetting which changes were sent" + ); + sync_state.sent_hashes.clear(); + } + STALL_RESET => { + warn!(peer = %peer, "Sync is still stalled; resetting it on both sides"); + peer_state.stalled_exchanges = 0; + peer_state.request_resync(); + } + _ => {} + } +} + +pub(crate) async fn run_sync_worker(state: AppState, mut rx: mpsc::UnboundedReceiver<SyncEvent>) { + let mut dirty_peers = HashSet::<String>::new(); + let mut retry_peers = HashSet::<String>::new(); + let mut document_dirty = false; + + info!("Starting synchronization worker"); + + loop { + if dirty_peers.is_empty() && retry_peers.is_empty() && !document_dirty { + match rx.recv().await { + Some(event) => match event { + SyncEvent::PeerValidated(peer) + | SyncEvent::MessageReceived(peer) + | SyncEvent::PeerNeedsSync(peer) => { + dirty_peers.insert(peer); + } + SyncEvent::RetryDelivery(peer) => { + retry_peers.insert(peer); + } + SyncEvent::DocumentChanged => { + document_dirty = true; + } + }, + None => { + warn!("Synchronization event channel closed"); + return; + } + } + } + + // Coalesce event bursts. + while let Ok(event) = rx.try_recv() { + match event { + SyncEvent::PeerValidated(peer) + | SyncEvent::MessageReceived(peer) + | SyncEvent::PeerNeedsSync(peer) => { + dirty_peers.insert(peer); + } + SyncEvent::RetryDelivery(peer) => { + retry_peers.insert(peer); + } + SyncEvent::DocumentChanged => { + document_dirty = true; + } + } + } + + // Let pending HTTP task locks wake up and complete their lifecycle transitions gracefully. + tokio::task::yield_now().await; + + if document_dirty { + document_dirty = false; + + // Lock and extract only the cloned string keys into a small transient vector + let peer_keys: Vec<String> = { + let peers_guard = state.peer_states.read().await; + peers_guard.keys().cloned().collect() + }; + + for peer in peer_keys { + if peer_cluster(&state, &peer).await.is_some() { + dirty_peers.insert(peer); + } + } + } + + // Retry existing envelopes separately from generating new ones. + let retries: Vec<String> = retry_peers.drain().collect(); + + for peer in retries { + if let Err(error) = deliver_sync_to_peer(&state, &peer, true).await { + warn!( + peer = %peer, + error = %format!("{error:#}"), + "Sync delivery retry failed" + ); + } + } + + let peers: Vec<String> = dirty_peers.drain().collect(); + + for peer in peers { + if let Err(error) = deliver_sync_to_peer(&state, &peer, false).await { + warn!( + peer = %peer, + error = %format!("{error:#}"), + "Synchronization processing failed" + ); + } + } + } +} + +// --- change publisher -------------------------------------------------------- + +fn patch_store_path(patch: &automerge::Patch) -> Vec<String> { + let mut path: Vec<String> = patch + .path + .iter() + .map(|(_, prop)| match prop { + automerge::Prop::Map(key) => key.clone(), + automerge::Prop::Seq(index) => index.to_string(), + }) + .collect(); + + match &patch.action { + automerge::PatchAction::PutMap { key, .. } + | automerge::PatchAction::DeleteMap { key, .. } => path.push(key.clone()), + // List, text and counter changes: the containing object changed. + _ => {} + } + path +} + +pub(crate) fn paths_related(a: &[String], b: &[String]) -> bool { + a.starts_with(b) || b.starts_with(a) +} + +pub(crate) async fn run_change_publisher(state: AppState) { + let mut changed = state.change_tx.subscribe(); + let mut last_heads = state.write_doc(|doc| doc.get_heads()).await; + + info!("Starting change publisher"); + + while changed.changed().await.is_ok() { + let before = last_heads.clone(); + + let (heads, paths, conflicts) = state + .write_doc(move |doc| { + let heads = doc.get_heads(); + if heads == before { + return (heads, Vec::new(), Vec::new()); + } + + // After an epoch switch the previous heads belong to the old document: + // report the root as changed, which every stream is related to. + if !before.iter().all(|h| doc.get_change_by_hash(h).is_some()) { + return (heads, vec![Vec::new()], Vec::new()); + } + + let mut paths = Vec::new(); + let mut conflicts = Vec::new(); + for patch in doc.diff(&before, &heads) { + let path = patch_store_path(&patch); + if matches!( + patch.action, + automerge::PatchAction::PutMap { conflict: true, .. } + ) { + conflicts.push(path.clone()); + } + paths.push(path); + } + (heads, paths, conflicts) + }) + .await; + + if heads == last_heads { + continue; + } + last_heads = heads; + + for path in &conflicts { + warn!(path = ?path, "Concurrent writes produced a conflict"); + } + + // Fails only when nobody is subscribed. + let _ = state.change_events.send(Arc::new(ChangeSet { paths })); + } +} |