split state moduleHEADmaster

1 files changed, 404 insertions, 0 deletions

diff --git a/src/state/mod.rs b/src/state/mod.rs
new file mode 100644
index 0000000..c5a6885
--- /dev/null
+++ b/src/state/mod.rs
@@ -0,0 +1,404 @@
+//! Shared node state: the document, lifecycle, per-peer state and timing.
+//!
+//! * `peer`: what a node knows about each peer (`PeerState`): standing, liveness,
+//!   handshake, incidents.
+//! * `session`: one Automerge sync with a peer (`SyncSession`).
+//! * `timing`: the intervals the protocol runs on (`Timing`).
+//! * `read`: serving reads, live or from the pre-rejoin copy.
+//! * `host`: the node as the remake module sees it (`remake::Host`).
+
+mod host;
+mod peer;
+mod read;
+mod session;
+mod timing;
+
+pub(crate) use peer::{HelloState, Incidents, Liveness, PeerState, Standing};
+pub(crate) use read::{ReadSource, StaleView};
+pub(crate) use session::{Sending, SyncSession};
+pub(crate) use timing::Timing;
+
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::time::Duration;
+
+use anyhow::{Result, bail};
+use automerge::AutoCommit;
+use tokio::sync::{Mutex, RwLock, broadcast, mpsc, watch};
+use tokio::time::Instant;
+use tracing::{debug, error, warn};
+
+use crate::document::replica::{Epoch, Replica};
+use crate::remake::{Host, RemakeState};
+use crate::swarm::sync::{ChangeSet, SyncEvent};
+use crate::swarm::throughput::ThroughputTable;
+use crate::transport;
+use crate::transport::firewall::Firewall;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) enum NodeLifecycle {
+    // TODO maybe rename discovered_swarm_id to swarm_id if that doesn't impair readability.
+    Joining { discovered_swarm_id: Option<String> },
+    Ready { swarm_id: String },
+}
+
+#[derive(Clone)]
+pub(crate) struct AppState {
+    /// The document and its epoch. Derefs to `AutoCommit`.
+    pub(crate) doc: Arc<RwLock<Replica>>,
+
+    /// Copy of `doc.epoch` for cheap, lock-free reads (handshakes, keepalives, /peers).
+    /// Only ever changed while holding `doc`'s write lock, together with `doc.epoch`.
+    pub(crate) epoch_tx: watch::Sender<Epoch>,
+
+    pub(crate) lifecycle: Arc<RwLock<NodeLifecycle>>,
+
+    // Serializes swarm-ID establishment, peer membership, and readiness transitions.
+    //[concurrent]swarm_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>,
+    /// Messages and bytes exchanged, per peer (`transport::traffic`).
+    pub(crate) traffic: Arc<transport::traffic::Traffic>,
+    /// Node-wide figures for /admin/stats (`swarm::stats`).
+    pub(crate) stats: Arc<crate::swarm::stats::NodeStats>,
+
+    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,
+
+    /// Per-node remake state (prebuilt snapshots, switch status, archive).
+    pub(crate) remake: Arc<RemakeState>,
+
+    /// Pre-shared key every member must prove in the handshake; None admits any node
+    /// that reaches us.
+    pub(crate) swarm_key: Option<Arc<crate::swarm::auth::SwarmKey>>,
+
+    /// Random per run of this node; our HELLOs carry it and replies must echo it.
+    pub(crate) run_id: transport::RunId,
+
+    /// Target number of neighbours (`--neighbours`); up to this many + 1 members the
+    /// swarm is a full mesh.
+    pub(crate) neighbours_k: usize,
+    /// Our alias (`--alias`), published in the member list (`swarm::members`).
+    pub(crate) alias: Option<String>,
+    /// Members alive within this window vote on remakes (`--electorate-window-days`).
+    pub(crate) electorate_window: Duration,
+    /// Chunked transfers served and fetched (`swarm::blob`).
+    pub(crate) blobs: Arc<crate::swarm::blob::Blobs>,
+    /// HTTP API users who may write protocol data through the store API
+    /// (`--protocol-writers`).
+    pub(crate) protocol_writers: Arc<std::collections::HashSet<String>>,
+    /// This node's own files, and how much it read and wrote them.
+    pub(crate) files: Arc<crate::files::Files>,
+    /// This node's document on disk (`--document-file`), if configured.
+    pub(crate) document_file: Option<Arc<crate::document::file::DocumentFile>>,
+    /// What this node remembers across restarts (`--state-file`), if configured.
+    pub(crate) state_file: Option<Arc<crate::swarm::state_file::StateFile>>,
+    /// The effective configuration with sources, imported values left out
+    /// (`config::Loaded::report`), for `GET /admin/config`.
+    pub(crate) config: Arc<serde_json::Value>,
+    /// Our own connectivity record and stability tier (`swarm::stability`).
+    pub(crate) stability: Arc<std::sync::Mutex<crate::swarm::stability::Stability>>,
+    /// Round-trip times to peers (`swarm::rtt`).
+    pub(crate) rtt: Arc<std::sync::Mutex<crate::swarm::rtt::RttTable>>,
+    /// Sync throughput with peers (`swarm::throughput`).
+    pub(crate) throughput: Arc<std::sync::Mutex<ThroughputTable>>,
+    /// The neighbour we are pulling changes from, and since when (`swarm::relay`).
+    pub(crate) pulling: Arc<std::sync::Mutex<Option<(String, Instant)>>>,
+    /// An announcement round is scheduled, for changes from these peers
+    /// (`swarm::relay`).
+    pub(crate) announce_pending: Arc<std::sync::Mutex<Option<std::collections::HashSet<String>>>>,
+    /// Our neighbours, once known (`swarm::topology`).
+    pub(crate) neighbours: Arc<std::sync::RwLock<Option<crate::swarm::topology::NeighbourView>>>,
+    /// Something the neighbours are computed from may have changed: the member list or a
+    /// peer's reachability (`swarm::topology::run_worker`).
+    pub(crate) topology_changed: Arc<tokio::sync::Notify>,
+
+    /// While rejoining a newer epoch, the previous document, served read-only (and
+    /// marked stale) instead of answering 503.
+    pub(crate) stale: Arc<std::sync::RwLock<Option<StaleView>>>,
+
+    /// While the initial sync isn't done: the one peer we catch up from.
+    pub(crate) catch_up_source: Arc<std::sync::Mutex<Option<CatchUpSource>>>,
+    /// Since when a closer peer has had a head start to become the source.
+    pub(crate) catch_up_preference: Arc<std::sync::Mutex<Option<Instant>>>,
+    /// A catch-up source given up for not answering our sync messages, and until when it
+    /// isn't chosen again (`swarm::source`).
+    pub(crate) catch_up_avoided: Arc<std::sync::Mutex<Option<(String, Instant)>>>,
+    /// The node's wall clock (`clock`).
+    pub(crate) clock: Arc<crate::clock::Clock>,
+}
+
+/// The peer a catching-up node syncs with first (`swarm::source`).
+#[derive(Clone, Debug)]
+pub(crate) struct CatchUpSource {
+    pub(crate) peer: String,
+    /// When it was chosen or last sent us anything.
+    pub(crate) last_heard: Instant,
+    /// When it was chosen or last sent us a sync message: progress, not just liveness.
+    pub(crate) last_sync: Instant,
+    /// Our sync messages to it in a row whose answer didn't come in time.
+    pub(crate) reply_timeouts: u32,
+    /// When we last warned that catching up from it stalls.
+    pub(crate) last_warned: Option<Instant>,
+    /// Our last probe PING to it.
+    pub(crate) last_probe: Option<Instant>,
+    /// Its keepalives in a row saying it wouldn't sync with us.
+    pub(crate) refusals: u32,
+}
+
+impl AppState {
+    pub(crate) fn local_id(&self) -> &str {
+        self.transport.local_id()
+    }
+
+    /// Peers that can serve our epoch's data: active, ready members on our epoch
+    /// (`PeerSummary::serves_epoch`).
+    pub(crate) async fn active_peers_on_local_epoch(&self) -> Vec<String> {
+        let epoch = self.local_epoch();
+        self.peer_summaries()
+            .await
+            .into_iter()
+            .filter(|peer| peer.serves_epoch(epoch))
+            .map(|peer| peer.id)
+            .collect()
+    }
+
+    pub(crate) fn lock_rtt(&self) -> std::sync::MutexGuard<'_, crate::swarm::rtt::RttTable> {
+        self.rtt
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    /// The local epoch. Code that labels or applies sync messages must use the epoch
+    /// of the locked `Replica` instead, so label and document can't disagree.
+    pub(crate) fn local_epoch(&self) -> Epoch {
+        *self.epoch_tx.borrow()
+    }
+
+    /// The local epoch and `Replica::epoch_seq`, read together.
+    pub(crate) async fn epoch_stamp(&self) -> (Epoch, u32) {
+        let replica = self.doc.read().await;
+        (replica.epoch, replica.epoch_seq)
+    }
+
+    pub(crate) async fn is_ready(&self) -> bool {
+        matches!(*self.lifecycle.read().await, NodeLifecycle::Ready { .. })
+    }
+
+    pub(crate) async fn swarm_id(&self) -> Option<String> {
+        match &*self.lifecycle.read().await {
+            NodeLifecycle::Ready { swarm_id } => Some(swarm_id.clone()),
+            NodeLifecycle::Joining {
+                discovered_swarm_id,
+            } => discovered_swarm_id.clone(),
+        }
+    }
+
+    // Record the expected swarm without making the node Ready. Bootstrap owns readiness.
+    pub(crate) async fn set_expected_swarm_id(&self, swarm_id: String) -> Result<()> {
+        //[concurrent]let _guard = self.swarm_validation.lock().await;
+        let mut lifecycle = self.lifecycle.write().await;
+        match &*lifecycle {
+            NodeLifecycle::Joining {
+                discovered_swarm_id: Some(existing),
+            } if existing != &swarm_id => {
+                bail!(
+                    "Conflicting expected swarm IDs: expected {}, received {}",
+                    existing,
+                    swarm_id
+                );
+            }
+            NodeLifecycle::Joining {
+                discovered_swarm_id: None,
+            } => {
+                let old_lifecycle = lifecycle.clone();
+                *lifecycle = NodeLifecycle::Joining {
+                    discovered_swarm_id: Some(swarm_id),
+                };
+                debug!("Node transitioned from {old_lifecycle:?} {lifecycle:?}");
+                Ok(())
+            }
+            NodeLifecycle::Joining {
+                discovered_swarm_id: Some(_),
+            } => {
+                warn!("Node already Joining with the same identity");
+                Ok(())
+            }
+            NodeLifecycle::Ready { swarm_id: existing } if existing == &swarm_id => {
+                warn!("Node already Ready with the same identity");
+                Ok(())
+            }
+            NodeLifecycle::Ready { swarm_id: existing } => {
+                bail!(
+                    "Cannot change established swarm ID: existing {}, received {}",
+                    existing,
+                    swarm_id
+                );
+            }
+        }
+    }
+
+    // Only bootstrap should call this, after validating document metadata.
+    pub(crate) async fn mark_ready(&self, swarm_id: String) -> Result<()> {
+        //[concurrent]let _guard = self.swarm_validation.lock().await;
+        let mut lifecycle = self.lifecycle.write().await;
+        match &*lifecycle {
+            NodeLifecycle::Joining {
+                discovered_swarm_id: Some(expected),
+            } if expected != &swarm_id => {
+                bail!(
+                    "Cannot mark node ready: expected swarm {}, received {}",
+                    expected,
+                    swarm_id
+                );
+            }
+            NodeLifecycle::Joining { .. } => {
+                let old_lifecycle = lifecycle.clone();
+                *lifecycle = NodeLifecycle::Ready { swarm_id };
+                debug!("Node transitioned from {old_lifecycle:?} {lifecycle:?}");
+            }
+            NodeLifecycle::Ready { swarm_id: existing } if existing == &swarm_id => {
+                debug!("Node already ready with the same identity");
+            }
+            NodeLifecycle::Ready { swarm_id: existing } => {
+                bail!(
+                    "Cannot change established swarm ID: existing {}, received {}",
+                    existing,
+                    swarm_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(|| self.new_peer_state(pk)),
+        )
+    }
+
+    /// State for a peer we haven't known so far (not inserted).
+    pub(crate) fn new_peer_state(&self, pk: &str) -> Arc<Mutex<PeerState>> {
+        let leads = crate::swarm::pair::leads(self.local_id(), pk);
+        Arc::new(Mutex::new(PeerState::new(leads, &self.timing)))
+    }
+
+    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.doc)
+            };
+            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 || {
+            let mut replica = doc.blocking_write();
+            let out = f(&mut replica.doc);
+            replica.refresh_heads();
+            out
+        })
+        .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 replica = state.doc.blocking_write();
+                let result = f(&mut replica.doc);
+                match &result {
+                    // One API write = one Automerge change.
+                    Ok(_) => {
+                        // Stamped with our clock (milliseconds), so receivers can tell
+                        // how long the change took to reach them (`swarm::stats`).
+                        replica.doc.commit_with(
+                            automerge::transaction::CommitOptions::default()
+                                .with_time(state.clock.now_millis() as i64),
+                        );
+                    }
+                    // Discard partial ops (e.g. old value deleted, new one not written).
+                    Err(_) => {
+                        replica.doc.rollback();
+                    }
+                }
+                replica.refresh_heads();
+                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");
+        }
+    }
+}