initial commit

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)
+}