diff options
Diffstat (limited to 'src/peer.rs')
| -rw-r--r-- | src/peer.rs | 593 |
1 files changed, 593 insertions, 0 deletions
diff --git a/src/peer.rs b/src/peer.rs new file mode 100644 index 0000000..3dd344e --- /dev/null +++ b/src/peer.rs @@ -0,0 +1,593 @@ +//! 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::state::{AppState, HelloState, Liveness, PeerMembership, PeerState}; +use crate::store::{self, StoreError}; +use crate::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); + 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() +} + +async fn local_keepalive(state: &AppState) -> transport::Keepalive { + // AutoCommit::get_heads takes &mut self (it closes the pending transaction). + let heads_digest = heads_digest(&mut *state.doc.write().await); + transport::Keepalive { heads_digest } +} + +async fn heal_if_diverged( + state: &AppState, + peer: &str, + local: &transport::Keepalive, + remote: &transport::Keepalive, +) { + 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 ------------------------------------------------ + +#[instrument(skip(state, peer))] +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_count = peer_states.len(), + "Discovered and registered peer" + ); + } else { + trace!("Peer already registered"); + } +} + +#[instrument(skip(state), fields(peer = %peer))] +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 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!(hello_state = ?peer_state.hello_state, "not sending"); + false + } + }; + + if !claimed { + return Ok(()); + } + + info!("Sending HELLO"); + transmit_hello(state, peer, HELLO_RETRY_INITIAL).await +} + +#[instrument(skip(state), fields(peer = %peer))] +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!( + 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 {})) + .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; + } +} + +#[instrument(skip(state), fields(peer = %peer))] +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 }; + + let ready = state.is_ready().await; + + debug!(ready, "Prepared HELLO_REPLY"); + + state + .transport + .send_message(peer, &ParsedMessage::HelloReply(reply)) + .await + .context("Sending HELLO_REPLY")?; + + debug!("HELLO_REPLY sent"); + Ok(()) +} + +#[instrument(skip(state), fields(from_peer = %from_peer, advertised_peers = ?advertised_peers))] +async fn discover_peers(state: &AppState, from_peer: &str, advertised_peers: &[String]) { + 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 --------------------------------------------------------- + +#[instrument(skip(state), fields(peer = %peer, received_cluster = %cluster_id))] +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!("Handshake complete; sync state reset"); + } + + let pending = std::mem::take(&mut peer_state.pending_inbound_sync); + + info!(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()), + } +} + +// --- inbound dispatch -------------------------------------------------------- + +#[instrument(skip(state, message))] +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"); + } + + debug!(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!( + 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; + } + 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(()) +} |