diff options
Diffstat (limited to 'src/state.rs')
| -rw-r--r-- | src/state.rs | 483 |
1 files changed, 483 insertions, 0 deletions
diff --git a/src/state.rs b/src/state.rs new file mode 100644 index 0000000..3fe1790 --- /dev/null +++ b/src/state.rs @@ -0,0 +1,483 @@ +//! Shared node state: the document, lifecycle, per-peer state and timing. + +use std::collections::HashMap; +use std::sync::atomic::AtomicBool; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{bail, Result}; +use automerge::{sync, AutoCommit}; +use tokio::sync::{broadcast, mpsc, watch, Mutex, RwLock}; +use tokio::time::Instant; +use tracing::{debug, error, warn}; + +use crate::firewall::Firewall; +use crate::sync::{ChangeSet, SyncEvent}; +use crate::transport; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum NodeLifecycle { + // TODO maybe rename discovered_cluster_id to cluster_id if that doesn't impair readability. + Joining { + discovered_cluster_id: Option<String>, + }, + Ready { + cluster_id: String, + }, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(crate) enum PeerMembership { + #[default] + Unknown, + Member { + cluster_id: String, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Liveness { + Active, + /// No inbound message for Timing::dormant_after. Kept indefinitely, probed slowly. + Dormant { since: Instant }, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) enum HelloState { + #[default] + NotSent, + /// A send is in flight. `next_delay` is how long to wait after it before retrying. + Sending { next_delay: Duration }, + /// HELLO was sent (or the send failed); no HELLO_REPLY yet. + AwaitingReply { + next_attempt: Instant, + next_delay: Duration, + }, + /// HELLO_REPLY received. No more retries. + Established, +} + +#[derive(Clone)] +pub(crate) struct PendingOutbound { + pub(crate) message: transport::EncodedMessage, + pub(crate) retry_delay: Duration, + pub(crate) next_attempt: Instant, +} + +#[derive(Clone)] +pub(crate) struct PeerState { + pub(crate) membership: PeerMembership, + pub(crate) sync_state: sync::State, + pub(crate) hello_state: HelloState, + + /// The peer sent HELLO; we still owe it a HELLO_REPLY (e.g. we had no cluster ID yet, + /// or sending failed). + pub(crate) reply_owed: bool, + + /// Received from this peer, waiting to be processed. + pub(crate) pending_inbound_sync: Vec<transport::SyncEnvelope>, + + /// Generated for this peer, waiting for delivery. + pub(crate) pending_outbound_sync: Option<PendingOutbound>, + + /// A sync pass was requested while an envelope was pending. Once that envelope + /// is delivered, the peer must be re-enqueued so the deferred changes go out. + pub(crate) needs_generate: bool, + + pub(crate) liveness: Liveness, + pub(crate) added_at: Instant, + pub(crate) last_seen: Option<Instant>, + pub(crate) last_ping_at: Option<Instant>, + + /// Last time a sync message from this peer was applied. + pub(crate) last_sync_rx: Option<Instant>, + + /// Last self-healing re-handshake (cooldown). + pub(crate) last_resync_at: Option<Instant>, +} + +impl Default for PeerState { + fn default() -> Self { + Self { + membership: PeerMembership::default(), + sync_state: sync::State::new(), + hello_state: HelloState::default(), + reply_owed: false, + pending_inbound_sync: Vec::new(), + pending_outbound_sync: None, + needs_generate: false, + liveness: Liveness::Active, + added_at: Instant::now(), + last_seen: None, + last_ping_at: None, + last_sync_rx: None, + last_resync_at: None, + } + } +} + +impl PeerState { + /// Start Automerge sync with this peer from scratch. Any envelope generated against + /// the old state is discarded; a stale RetryDelivery event is then a no-op. + pub(crate) fn reset_sync(&mut self) { + self.sync_state = sync::State::new(); + self.pending_outbound_sync = None; + self.needs_generate = false; + } + + fn quiet_for(&self, now: Instant) -> Duration { + now.saturating_duration_since(self.last_seen.unwrap_or(self.added_at)) + } + + /// Records inbound activity. Returns true if the peer was dormant. + pub(crate) fn mark_seen(&mut self, now: Instant) -> bool { + self.last_seen = Some(now); + + if matches!(self.liveness, Liveness::Dormant { .. }) { + self.liveness = Liveness::Active; + // Force a fresh handshake; its reply resets sync on both sides. + if !matches!(self.hello_state, HelloState::Sending { .. }) { + self.hello_state = HelloState::NotSent; + } + return true; + } + false + } + + /// 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 { + return false; + } + + self.liveness = Liveness::Dormant { since: now }; + self.reset_sync(); + + if !matches!(self.hello_state, HelloState::Sending { .. }) { + self.hello_state = HelloState::AwaitingReply { + next_attempt: now + timing.dormant_probe, + next_delay: timing.dormant_probe, + }; + } + true + } + + /// True (and records the ping) if an established, quiet peer needs a keepalive. + pub(crate) fn take_keepalive_due(&mut self, timing: &Timing, now: Instant) -> bool { + let due = self.liveness == Liveness::Active + && self.hello_state == HelloState::Established + && self.quiet_for(now) >= timing.keepalive + && self + .last_ping_at + .is_none_or(|t| now.saturating_duration_since(t) >= timing.keepalive); + + if due { + self.last_ping_at = Some(now); + } + due + } + + /// Self-healing check on PING/PONG. Returns true if a re-handshake should start. + pub(crate) fn needs_heal( + &mut self, + local: &transport::HeadsDigest, + remote: &transport::HeadsDigest, + timing: &Timing, + now: Instant, + ) -> bool { + if local == remote + || self.liveness != Liveness::Active + || self.hello_state != HelloState::Established + { + return false; + } + + let sync_idle = now.saturating_duration_since(self.last_sync_rx.unwrap_or(self.added_at)) + >= timing.resync_after; + let cooled_down = self + .last_resync_at + .is_none_or(|t| now.saturating_duration_since(t) >= timing.resync_after); + + if !(sync_idle && cooled_down) { + return false; + } + + self.last_resync_at = Some(now); + self.hello_state = HelloState::NotSent; // establish_peer resets sync on the reply + true + } + + /// Heard from recently. A returning node counts as alive from its first message, + /// before the handshake finishes, so a cluster-wide forget can't make it flap. + pub(crate) fn is_alive(&self, timing: &Timing, now: Instant) -> bool { + self.liveness == Liveness::Active + && self + .last_seen + .is_some_and(|t| now.saturating_duration_since(t) < timing.dormant_after) + } +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct Timing { + pub(crate) keepalive: Duration, + pub(crate) dormant_after: Duration, + pub(crate) resync_after: Duration, + pub(crate) dormant_probe: Duration, +} + +impl Timing { + // Recommended values + // Environment | Keepalive | Notes + // --------------------------------------------+-----------+------------------------------------------------------------ + // Test harness | 1-2 s | Dormancy after about 4–8 s and healing within seconds keep tests short. + // LAN or small cluster (up to ~50 nodes) | 15-30 s | Under 3.5 messages/s per node at 50 nodes and 30 s. Current setting. + // WAN, mobile or metered links, or 100+ nodes | 60-120 s | At 100 nodes and 30 s it’s about 6–7 messages/s per node. Fine, but rarely worth it on battery or metered links. + // Several hundred nodes or more | - | The full mesh becomes the problem in itself, not the interval. Message counts grow O(N²) across the cluster, which calls for partial views (SWIM-style random probing and syncing with a subset of peers). + pub(crate) fn from_keepalive(keepalive: Duration) -> Self { + Self { + keepalive, + dormant_after: keepalive * 4, // dormant_after >= 4* keepalive to tolerate three lost PINGs in a row + resync_after: keepalive * 2, // resync_after should be >= 2* keepalive to guarantee at least one digest exchange while sync is idle + dormant_probe: keepalive * 10, // >= 10*: dormant peers cost almost nothing, but still rejoin within minutes + } + } +} + +#[derive(Clone)] +pub(crate) struct AppState { + pub(crate) doc: Arc<RwLock<AutoCommit>>, + + pub(crate) lifecycle: Arc<RwLock<NodeLifecycle>>, + + // Serializes cluster-ID establishment, peer membership, and readiness transitions. + //[concurrent]cluster_validation: Arc<Mutex<()>>, + pub(crate) peer_states: Arc<RwLock<HashMap<String, Arc<Mutex<PeerState>>>>>, + pub(crate) transport: Arc<dyn transport::Transport>, + pub(crate) firewall: Arc<Firewall>, + + pub(crate) sync_events: mpsc::UnboundedSender<SyncEvent>, + + // TODO see key_watchers.md ("Use a per-key notification registry instead of the global `watch<u64>`") + pub(crate) change_tx: watch::Sender<u64>, + + pub(crate) initial_sync_done: Arc<AtomicBool>, + + pub(crate) change_events: broadcast::Sender<Arc<ChangeSet>>, + + pub(crate) timing: Timing, +} + +impl AppState { + pub(crate) fn local_id(&self) -> &str { + self.transport.local_id() + } + + pub(crate) async fn is_ready(&self) -> bool { + matches!(*self.lifecycle.read().await, NodeLifecycle::Ready { .. }) + } + + pub(crate) async fn cluster_id(&self) -> Option<String> { + match &*self.lifecycle.read().await { + NodeLifecycle::Ready { cluster_id } => Some(cluster_id.clone()), + NodeLifecycle::Joining { + discovered_cluster_id, + } => discovered_cluster_id.clone(), + } + } + + // Record the expected cluster without making the node Ready. Bootstrap owns readiness. + pub(crate) async fn set_expected_cluster_id(&self, cluster_id: String) -> Result<()> { + //[concurrent]let _guard = self.cluster_validation.lock().await; + let mut lifecycle = self.lifecycle.write().await; + match &*lifecycle { + NodeLifecycle::Joining { + discovered_cluster_id: Some(existing), + } if existing != &cluster_id => { + bail!( + "Conflicting expected cluster IDs: expected {}, received {}", + existing, + cluster_id + ); + } + NodeLifecycle::Joining { + discovered_cluster_id: None, + } => { + let old_lifecycle = lifecycle.clone(); + *lifecycle = NodeLifecycle::Joining { + discovered_cluster_id: Some(cluster_id), + }; + debug!("Node transitioned from {old_lifecycle:?} {lifecycle:?}"); + Ok(()) + } + NodeLifecycle::Joining { + discovered_cluster_id: Some(_), + } => { + warn!("Node already Joining with the same identity"); + Ok(()) + } + NodeLifecycle::Ready { + cluster_id: existing, + } if existing == &cluster_id => { + warn!("Node already Ready with the same identity"); + Ok(()) + } + NodeLifecycle::Ready { + cluster_id: existing, + } => { + bail!( + "Cannot change established cluster ID: existing {}, received {}", + existing, + cluster_id + ); + } + } + } + + // Only bootstrap should call this, after validating document metadata. + pub(crate) async fn mark_ready(&self, cluster_id: String) -> Result<()> { + //[concurrent]let _guard = self.cluster_validation.lock().await; + let mut lifecycle = self.lifecycle.write().await; + match &*lifecycle { + NodeLifecycle::Joining { + discovered_cluster_id: Some(expected), + } if expected != &cluster_id => { + bail!( + "Cannot mark node ready: expected cluster {}, received {}", + expected, + cluster_id + ); + } + NodeLifecycle::Joining { .. } => { + let old_lifecycle = lifecycle.clone(); + *lifecycle = NodeLifecycle::Ready { cluster_id }; + debug!("Node transitioned from {old_lifecycle:?} {lifecycle:?}"); + } + NodeLifecycle::Ready { + cluster_id: existing, + } if existing == &cluster_id => { + debug!("Node already ready with the same identity"); + } + NodeLifecycle::Ready { + cluster_id: existing, + } => { + bail!( + "Cannot change established cluster ID: existing {}, received {}", + existing, + cluster_id + ); + } + } + Ok(()) + } + + /// Looks up a peer without creating it. Background work (sync, retries, + /// keepalives) must use this so a forgotten peer isn't resurrected. + pub(crate) async fn get_peer_state(&self, pk: &str) -> Option<Arc<Mutex<PeerState>>> { + self.peer_states.read().await.get(pk).cloned() + } + + pub(crate) async fn get_or_create_peer_state(&self, pk: &str) -> Arc<Mutex<PeerState>> { + let mut peers = self.peer_states.write().await; + + Arc::clone( + peers + .entry(pk.to_owned()) + .or_insert_with(|| Arc::new(Mutex::new(PeerState::default()))), + ) + } + + pub(crate) fn notify_store_changed(&self) { + self.change_tx.send_modify(|version| { + *version = version.wrapping_add(1); + }); + } + + /// Reads the document on the blocking pool; `finish` runs after the lock is released. + pub(crate) async fn render<T, R>( + &self, + read: impl FnOnce(&AutoCommit) -> T + Send + 'static, + finish: impl FnOnce(T) -> R + Send + 'static, + ) -> R + where + R: Send + 'static, + { + let doc = Arc::clone(&self.doc); + tokio::task::spawn_blocking(move || { + let data = { + let guard = doc.blocking_read(); + read(&guard) + }; + finish(data) + }) + .await + .expect("document read task panicked") + } + + /// Mutates the document on the blocking pool. + pub(crate) async fn write_doc<T>( + &self, + f: impl FnOnce(&mut AutoCommit) -> T + Send + 'static, + ) -> T + where + T: Send + 'static, + { + let doc = Arc::clone(&self.doc); + tokio::task::spawn_blocking(move || f(&mut doc.blocking_write())) + .await + .expect("document write task panicked") + } + + /// Mutates the document on the blocking pool as one Automerge change, then + /// publishes it. Write and notification run in the same blocking task, so a + /// cancelled request can't leave a committed write unannounced. + pub(crate) async fn mutate_doc<T, E>( + &self, + f: impl FnOnce(&mut AutoCommit) -> Result<T, E> + Send + 'static, + ) -> Result<T, E> + where + T: Send + 'static, + E: Send + 'static, + { + let state = self.clone(); + + tokio::task::spawn_blocking(move || { + let result = { + let mut doc = state.doc.blocking_write(); + let result = f(&mut doc); + match &result { + // One API write = one Automerge change. + Ok(_) => { + doc.commit(); + } + // Discard partial ops (e.g. old value deleted, new one not written). + Err(_) => { + doc.rollback(); + } + } + result + }; // lock released before notifying + + if result.is_ok() { + state.publish_change(); + } + result + }) + .await + .expect("document write task panicked") + } + + pub(crate) fn publish_change(&self) { + self.notify_store_changed(); + if let Err(error) = self.sync_events.send(SyncEvent::DocumentChanged) { + // The write is committed; failing the request would misreport it. + error!(error = %error, "Sync worker is gone; change will not be replicated"); + } + } +} + +pub(crate) fn unix_millis() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} |