split state moduleHEADmaster

1 files changed, 276 insertions, 0 deletions

diff --git a/src/state/host.rs b/src/state/host.rs
new file mode 100644
index 0000000..863cc70
--- /dev/null
+++ b/src/state/host.rs
@@ -0,0 +1,276 @@
+//! The node as the remake module sees it (`remake::Host`), and what follows an epoch
+//! switch.
+
+use std::future::Future;
+use std::sync::Arc;
+use std::time::Duration;
+
+use anyhow::Result;
+use automerge::{AutoCommit, ChangeHash};
+use tokio::sync::{Mutex, broadcast};
+use tokio::time::Instant;
+use tracing::debug;
+
+use super::{AppState, HelloState, Liveness, NodeLifecycle, PeerState, StaleView, Standing};
+use crate::document::replica::Epoch;
+use crate::remake::{Host, PeerSummary, RemakeState};
+use crate::swarm::sync::{ChangeSet, SyncEvent};
+
+impl AppState {
+    /// After the document was replaced: every peer's sync state refers to the old
+    /// document, so start over with all of them. Peers still on the old epoch are
+    /// skipped by the sync worker until they switch too.
+    async fn after_epoch_switch(&self) {
+        let peers: Vec<(String, Arc<Mutex<PeerState>>)> = self
+            .peer_states
+            .read()
+            .await
+            .iter()
+            .map(|(id, peer)| (id.clone(), Arc::clone(peer)))
+            .collect();
+
+        // Sync states were reset under the document lock by the swap itself.
+        self.publish_change();
+
+        // Tell peers about the new epoch now instead of with the next keepalive (up to a
+        // keepalive interval later). Ready peers already on it learn it from our sync
+        // opening, which goes out right away; everyone else gets a PING, so peers that
+        // never saw the commit start catching up (we don't sync with them, so nothing
+        // else would tell them).
+        let local = self.local_epoch();
+        let mut announce = Vec::new();
+        for (id, peer_state_mutex) in &peers {
+            // Others learn it from their own neighbours.
+            if !crate::swarm::topology::is_neighbour(self, id) {
+                continue;
+            }
+            let mut peer_state = peer_state_mutex.lock().await;
+            // What it said about syncing with us was about our old epoch.
+            peer_state.syncs_with_us = None;
+            let gets_our_opening = peer_state.liveness == Liveness::Active
+                && matches!(peer_state.standing, Standing::Member { .. })
+                && peer_state.epoch == Some(local)
+                && peer_state.reports_ready == Some(true);
+            if !gets_our_opening {
+                announce.push(id);
+            }
+        }
+
+        for id in announce {
+            if let Err(error) = crate::swarm::peer::send_ping(self, id).await {
+                debug!(peer = %id, error = %format!("{error:#}"), "Epoch announcement failed");
+            }
+        }
+
+        for (id, _) in peers {
+            let _ = self.sync_events.send(SyncEvent::PeerNeedsSync(id));
+        }
+    }
+}
+
+/// The node as seen by the remake module.
+impl Host for AppState {
+    fn remake(&self) -> &RemakeState {
+        &self.remake
+    }
+
+    fn clock(&self) -> Arc<crate::clock::Clock> {
+        Arc::clone(&self.clock)
+    }
+
+    fn node_id(&self) -> String {
+        self.local_id().to_owned()
+    }
+
+    fn current_epoch(&self) -> Epoch {
+        self.local_epoch()
+    }
+
+    fn dormant_after(&self) -> Duration {
+        self.timing.dormant_after
+    }
+
+    fn electorate(&self, doc: &AutoCommit, now_ms: u64) -> Vec<String> {
+        crate::swarm::members::electorate(doc, now_ms, self.electorate_window)
+    }
+
+    fn fetch_change(
+        &self,
+        hash: ChangeHash,
+    ) -> impl Future<Output = Result<automerge::Change>> + Send {
+        let state = self.clone();
+        async move {
+            let peers = state.active_peers_on_local_epoch().await;
+            crate::swarm::blob::fetch_change(&state, hash, peers).await
+        }
+    }
+
+    fn fetch_content(&self, hash: [u8; 32]) -> impl Future<Output = Result<Vec<u8>>> + Send {
+        let state = self.clone();
+        async move {
+            let peers = state.active_peers_on_local_epoch().await;
+            crate::swarm::blob::fetch_content(&state, hash, peers).await
+        }
+    }
+
+    fn request_sync(&self) {
+        crate::swarm::source::resync_everyone(self);
+    }
+
+    fn is_node_ready(&self) -> impl Future<Output = bool> + Send {
+        let state = self.clone();
+        async move { state.is_ready().await }
+    }
+
+    fn peer_summaries(&self) -> impl Future<Output = Vec<PeerSummary>> + Send {
+        let state = self.clone();
+        async move {
+            let peers: Vec<(String, Arc<Mutex<PeerState>>)> = state
+                .peer_states
+                .read()
+                .await
+                .iter()
+                .map(|(id, peer)| (id.clone(), Arc::clone(peer)))
+                .collect();
+
+            let mut summaries = Vec::with_capacity(peers.len());
+            for (id, peer) in peers {
+                let peer = peer.lock().await;
+                summaries.push(PeerSummary {
+                    id,
+                    active: peer.liveness == Liveness::Active,
+                    member: matches!(peer.standing, Standing::Member { .. })
+                        && peer.hello_state == HelloState::Established,
+                    ready: peer.reports_ready,
+                    epoch: peer.epoch,
+                });
+            }
+            summaries
+        }
+    }
+
+    fn suspend(&self) -> impl Future<Output = bool> + Send {
+        let state = self.clone();
+        async move {
+            let mut lifecycle = state.lifecycle.write().await;
+            let NodeLifecycle::Ready { swarm_id } = &*lifecycle else {
+                return false;
+            };
+            let swarm_id = swarm_id.clone();
+            debug!(swarm_id = %swarm_id, "Node suspended (rejoining)");
+            *lifecycle = NodeLifecycle::Joining {
+                discovered_swarm_id: Some(swarm_id),
+            };
+            true
+        }
+    }
+
+    fn resume(&self) -> impl Future<Output = ()> + Send {
+        let state = self.clone();
+        async move {
+            let mut lifecycle = state.lifecycle.write().await;
+            if let NodeLifecycle::Joining {
+                discovered_swarm_id: Some(swarm_id),
+            } = &*lifecycle
+            {
+                let swarm_id = swarm_id.clone();
+                debug!(swarm_id = %swarm_id, "Node resumed");
+                *lifecycle = NodeLifecycle::Ready { swarm_id };
+            }
+        }
+    }
+
+    fn set_stale_view(&self, view: Option<(Epoch, AutoCommit)>) {
+        self.set_stale(view.map(|(epoch, doc)| StaleView {
+            epoch,
+            doc: Arc::new(doc),
+        }));
+    }
+
+    fn initial_sync_done(&self) -> bool {
+        self.initial_sync_done
+            .load(std::sync::atomic::Ordering::Acquire)
+    }
+
+    fn reset_initial_sync(&self) {
+        self.initial_sync_done
+            .store(false, std::sync::atomic::Ordering::Release);
+        *self
+            .catch_up_source
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
+        *self
+            .catch_up_preference
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
+    }
+
+    fn subscribe_changes(&self) -> broadcast::Receiver<Arc<ChangeSet>> {
+        self.change_events.subscribe()
+    }
+
+    fn read_doc<T, F>(&self, f: F) -> impl Future<Output = T> + Send
+    where
+        F: FnOnce(&AutoCommit) -> T + Send + 'static,
+        T: Send + 'static,
+    {
+        let state = self.clone();
+        async move { state.render(f, |value| value).await }
+    }
+
+    fn commit_doc<T, E, F>(&self, f: F) -> impl Future<Output = Result<T, E>> + Send
+    where
+        F: FnOnce(&mut AutoCommit) -> Result<T, E> + Send + 'static,
+        T: Send + 'static,
+        E: Send + 'static,
+    {
+        let state = self.clone();
+        async move { state.mutate_doc(f).await }
+    }
+
+    fn switch_epoch<E, F>(
+        &self,
+        epoch: Epoch,
+        f: F,
+    ) -> impl Future<Output = Result<AutoCommit, E>> + Send
+    where
+        F: FnOnce(&mut AutoCommit) -> Result<AutoCommit, E> + Send + 'static,
+        E: Send + 'static,
+    {
+        let state = self.clone();
+        async move {
+            let swapped = {
+                let state = state.clone();
+                tokio::task::spawn_blocking(move || -> Result<AutoCommit, E> {
+                    let mut replica = state.doc.blocking_write();
+                    let new_doc = f(&mut replica.doc)?;
+                    let old_doc = std::mem::replace(&mut replica.doc, new_doc);
+                    replica.epoch = epoch;
+                    replica.epoch_seq += 1;
+                    replica.refresh_heads();
+                    // Updated under the same lock, so the cached epoch never disagrees
+                    // with the document for anyone holding the lock.
+                    state.epoch_tx.send_replace(epoch);
+                    // Reset every peer's sync state before the lock is released, so no
+                    // message is generated from or applied to a state that describes the
+                    // old document. (Lock order doc -> peer, as in sync generation.)
+                    let now = Instant::now();
+                    for peer in state.peer_states.blocking_read().values() {
+                        let mut peer = peer.blocking_lock();
+                        peer.reset_sync();
+                        // A fresh start, not an idle sync: no heal re-handshake for it.
+                        peer.last_sync_rx = Some(now);
+                    }
+                    Ok(old_doc)
+                })
+                .await
+                .expect("epoch switch task panicked")
+            };
+
+            if swapped.is_ok() {
+                state.after_epoch_switch().await;
+            }
+            swapped
+        }
+    }
+}