1 files changed, 320 insertions, 0 deletions
diff --git a/src/remake/local.rs b/src/remake/local.rs
new file mode 100644
index 0000000..35a4ff1
--- /dev/null
+++ b/src/remake/local.rs
@@ -0,0 +1,320 @@
+//! Per-node, in-memory remake state: snapshots built ahead of the switch, switch
+//! progress for reports, a rejoin in progress, quarantined changes, and archived
+//! documents of previous epochs.
+//!
+//! Nothing here is replicated or durable. After a restart the node simply has no
+//! archive or quarantine, and rebuilds snapshots when needed.
+
+use std::collections::HashMap;
+use std::sync::{Mutex, MutexGuard};
+
+use automerge::AutoCommit;
+use serde::Serialize;
+use tokio::time::Instant;
+
+use super::records::{unix_millis, Retention};
+use super::snapshot::Snapshot;
+use crate::document::replica::Epoch;
+
+#[derive(Default)]
+pub(crate) struct RemakeState {
+    // A std mutex: every access is short and never spans an await.
+    inner: Mutex<Inner>,
+}
+
+#[derive(Default)]
+struct Inner {
+    /// Built by the coordinator at commit time; consumed by the local switch.
+    prebuilt: HashMap<String, Snapshot>,
+    /// When this node first saw a remake committed.
+    first_seen: HashMap<String, Instant>,
+    status: HashMap<String, LocalStatus>,
+    archive: Vec<Archived>,
+    /// Last attempt to catch up with a newer epoch.
+    catch_up: Option<CatchUp>,
+    /// A rejoin in progress: the node syncs an empty document of a newer epoch.
+    rejoin: Option<Rejoin>,
+    quarantine: Vec<Quarantined>,
+}
+
+/// Catching up with a newer epoch: attempts and the last failure.
+#[derive(Clone, Debug)]
+pub(crate) struct CatchUp {
+    pub(crate) target: Epoch,
+    pub(crate) attempts: u32,
+    pub(crate) last_attempt: Instant,
+    pub(crate) error: Option<String>,
+}
+
+/// The node left `previous` and is syncing `target` from scratch.
+pub(crate) struct Rejoin {
+    pub(crate) previous: Epoch,
+    pub(crate) target: Epoch,
+    /// The document of `previous`, until it is replayed, quarantined or restored.
+    pub(crate) old: AutoCommit,
+    pub(crate) started: Instant,
+    /// The node served requests before; it does again once the rejoin is over.
+    pub(crate) was_ready: bool,
+}
+
+/// Changes of this node that were not replayed into the current epoch.
+pub(crate) struct Quarantined {
+    pub(crate) id: String,
+    /// The epoch the changes were made in, and the one the node rejoined.
+    pub(crate) epoch: Epoch,
+    pub(crate) target: Epoch,
+    pub(crate) created_at: u64,
+    pub(crate) reason: String,
+    /// Store paths the changes touched (minimal, prefix-free).
+    pub(crate) paths: Vec<Vec<String>>,
+    /// The old document; `paths` are read from its merged view.
+    pub(crate) document: AutoCommit,
+}
+
+#[derive(Clone, Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub(crate) struct CatchUpReport {
+    pub(crate) target: String,
+    pub(crate) attempts: u32,
+    /// True while the node syncs the target epoch (not serving requests).
+    pub(crate) rejoining: bool,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) error: Option<String>,
+}
+
+/// What happened to a remake on this node.
+#[derive(Clone, Debug, Default, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub(crate) struct LocalStatus {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) switched_at: Option<u64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) replayed_paths: Option<usize>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) failed_paths: Option<usize>,
+    /// ID of the quarantine entry holding this node's changes that weren't replayed.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) quarantine: Option<String>,
+    /// Why the switch did not happen (yet).
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) error: Option<String>,
+}
+
+struct Archived {
+    epoch: Epoch,
+    // Kept for auditing and recovery by hand; not read by the protocol.
+    #[allow(dead_code)]
+    doc: AutoCommit,
+    archived_at: u64,
+    retention: Retention,
+}
+
+/// Summary of an archived epoch.
+#[derive(Clone, Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub(crate) struct ArchivedEpoch {
+    pub(crate) generation: u64,
+    pub(crate) id: String,
+    pub(crate) archived_at: u64,
+}
+
+impl RemakeState {
+    fn lock(&self) -> MutexGuard<'_, Inner> {
+        self.inner
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    pub(crate) fn store_prebuilt(&self, id: &str, snapshot: Snapshot) {
+        self.lock().prebuilt.insert(id.to_owned(), snapshot);
+    }
+
+    pub(crate) fn take_prebuilt(&self, id: &str) -> Option<Snapshot> {
+        self.lock().prebuilt.remove(id)
+    }
+
+    /// When this node first saw the remake committed (now, on the first call).
+    pub(crate) fn first_seen(&self, id: &str) -> Instant {
+        *self
+            .lock()
+            .first_seen
+            .entry(id.to_owned())
+            .or_insert_with(Instant::now)
+    }
+
+    pub(crate) fn status(&self, id: &str) -> Option<LocalStatus> {
+        self.lock().status.get(id).cloned()
+    }
+
+    pub(crate) fn set_status(&self, id: &str, status: LocalStatus) {
+        self.lock().status.insert(id.to_owned(), status);
+    }
+
+    /// Keeps the previous epoch's document, then drops what has outlived its retention.
+    pub(crate) fn archive(
+        &self,
+        epoch: Epoch,
+        doc: AutoCommit,
+        retention: Retention,
+        current_generation: u64,
+    ) {
+        let mut inner = self.lock();
+        inner.archive.push(Archived {
+            epoch,
+            doc,
+            archived_at: unix_millis(),
+            retention,
+        });
+        prune(&mut inner.archive, current_generation);
+    }
+
+    pub(crate) fn prune_archive(&self, current_generation: u64) {
+        prune(&mut self.lock().archive, current_generation);
+    }
+
+    pub(crate) fn catch_up(&self) -> Option<CatchUp> {
+        self.lock().catch_up.clone()
+    }
+
+    pub(crate) fn set_catch_up(&self, catch_up: Option<CatchUp>) {
+        self.lock().catch_up = catch_up;
+    }
+
+    pub(crate) fn catch_up_report(&self) -> Option<CatchUpReport> {
+        let inner = self.lock();
+        let rejoining = inner.rejoin.is_some();
+        inner.catch_up.as_ref().map(|catch_up| CatchUpReport {
+            target: catch_up.target.to_string(),
+            attempts: catch_up.attempts,
+            rejoining,
+            error: catch_up.error.clone(),
+        })
+    }
+
+    pub(crate) fn is_rejoining(&self) -> bool {
+        self.lock().rejoin.is_some()
+    }
+
+    /// Progress of the rejoin in progress: (target, started).
+    pub(crate) fn rejoin_progress(&self) -> Option<(Epoch, Instant)> {
+        self.lock()
+            .rejoin
+            .as_ref()
+            .map(|rejoin| (rejoin.target, rejoin.started))
+    }
+
+    pub(crate) fn start_rejoin(&self, rejoin: Rejoin) {
+        self.lock().rejoin = Some(rejoin);
+    }
+
+    pub(crate) fn take_rejoin(&self) -> Option<Rejoin> {
+        self.lock().rejoin.take()
+    }
+
+    pub(crate) fn add_quarantine(&self, quarantined: Quarantined) {
+        self.lock().quarantine.push(quarantined);
+    }
+
+    /// Runs `f` on the quarantined entry with this ID.
+    pub(crate) fn with_quarantined<T>(
+        &self,
+        id: &str,
+        f: impl FnOnce(&Quarantined) -> T,
+    ) -> Option<T> {
+        self.lock()
+            .quarantine
+            .iter()
+            .find(|entry| entry.id == id)
+            .map(f)
+    }
+
+    pub(crate) fn list_quarantine<T>(&self, f: impl Fn(&Quarantined) -> T) -> Vec<T> {
+        self.lock().quarantine.iter().map(f).collect()
+    }
+
+    pub(crate) fn take_quarantined(&self, id: &str) -> Option<Quarantined> {
+        let mut inner = self.lock();
+        let index = inner.quarantine.iter().position(|entry| entry.id == id)?;
+        Some(inner.quarantine.remove(index))
+    }
+
+    /// Puts an entry back, e.g. after a failed replay.
+    pub(crate) fn restore_quarantined(&self, quarantined: Quarantined) {
+        self.lock().quarantine.push(quarantined);
+    }
+
+    pub(crate) fn archived_epochs(&self) -> Vec<ArchivedEpoch> {
+        self.lock()
+            .archive
+            .iter()
+            .map(|archived| ArchivedEpoch {
+                generation: archived.epoch.generation,
+                id: archived.epoch.id_hex(),
+                archived_at: archived.archived_at,
+            })
+            .collect()
+    }
+}
+
+/// An archive is kept until every limit of its retention has passed.
+fn prune(archive: &mut Vec<Archived>, current_generation: u64) {
+    let now = unix_millis();
+    archive.retain(|archived| {
+        archived.retention.covers(
+            current_generation.saturating_sub(archived.epoch.generation),
+            now.saturating_sub(archived.archived_at),
+        )
+    });
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn archive_retention() {
+        let remake = RemakeState::default();
+        let epoch = |generation| Epoch {
+            generation,
+            id: [generation as u8; 32],
+        };
+
+        // Kept by epoch count only.
+        remake.archive(
+            epoch(0),
+            AutoCommit::new(),
+            Retention {
+                epochs: Some(1),
+                seconds: None,
+            },
+            1,
+        );
+        assert_eq!(remake.archived_epochs().len(), 1);
+        remake.prune_archive(2);
+        assert!(remake.archived_epochs().is_empty());
+
+        // Kept by time while the epoch count has passed.
+        remake.archive(
+            epoch(1),
+            AutoCommit::new(),
+            Retention {
+                epochs: Some(0),
+                seconds: Some(3600),
+            },
+            5,
+        );
+        assert_eq!(remake.archived_epochs().len(), 1);
+
+        // No limits at all: not kept.
+        remake.archive(
+            epoch(2),
+            AutoCommit::new(),
+            Retention {
+                epochs: None,
+                seconds: None,
+            },
+            3,
+        );
+        assert_eq!(remake.archived_epochs().len(), 1);
+    }
+}