allow to persist document

2 files changed, 363 insertions, 0 deletions

diff --git a/src/document/file.rs b/src/document/file.rs
new file mode 100644
index 0000000..a773502
--- /dev/null
+++ b/src/document/file.rs
@@ -0,0 +1,362 @@
+//! The document file (`--document-file`): this node's document on disk, so a restart
+//! resumes with it instead of starting empty.
+//!
+//! Format: the line `continuum-document/1`, a header line (JSON: the actor IDs this
+//! node wrote with, so a later rejoin replays those writes as its own), then records,
+//! each a little-endian `u32` length and that many bytes: first a full save of the
+//! document, then increments (the changes after the heads written before). A record cut
+//! short by a crash is dropped when reading.
+//!
+//! Written by a worker (`run`), only while the node is ready and not rejoining (a rejoin
+//! keeps the old file until it is done):
+//! * every `--document-append-secs`, what changed is appended;
+//! * the whole file is rewritten, atomically, when the epoch or this node's actors
+//!   changed (a switch, a rejoin, a restart), and when the appended increments outgrew
+//!   the full save, at most every `--document-rewrite-secs`;
+//! * at shutdown, what changed is appended.
+
+use std::collections::HashSet;
+use std::path::{Path, PathBuf};
+use std::sync::{Arc, Mutex, MutexGuard};
+use std::time::Duration;
+
+use anyhow::{Context, Result, anyhow, bail};
+use automerge::{ActorId, AutoCommit, ChangeHash};
+use serde::{Deserialize, Serialize};
+use tokio::time::Instant;
+use tracing::{debug, error, warn};
+
+use crate::document::replica::Epoch;
+use crate::files::{FileKind, Files};
+use crate::state::AppState;
+
+const MAGIC: &[u8] = b"continuum-document/1\n";
+
+#[derive(Debug, Default, Serialize, Deserialize)]
+struct Header {
+    /// Hex actor IDs.
+    #[serde(default)]
+    actors: Vec<String>,
+}
+
+/// A document file as read at start.
+pub(crate) struct Loaded {
+    /// Under a fresh actor.
+    pub(crate) doc: AutoCommit,
+    /// The actors this node wrote with before.
+    pub(crate) actors: Vec<ActorId>,
+}
+
+/// Reads the document file; None if there is none. An error if it can't be read or
+/// isn't one: the caller decides (`--force badDocumentFile`).
+pub(crate) fn load(path: &Path, files: &Files) -> Result<Option<Loaded>> {
+    let Some(bytes) = files
+        .read(FileKind::Document, path)
+        .with_context(|| format!("Reading the document file {}", path.display()))?
+    else {
+        return Ok(None);
+    };
+    let (loaded, dropped) =
+        parse(&bytes).with_context(|| format!("The document file {}", path.display()))?;
+    if dropped > 0 {
+        warn!(
+            path = %path.display(),
+            bytes = dropped,
+            "Dropped a record cut short at the end of the document file (a crash while appending)"
+        );
+    }
+    Ok(Some(loaded))
+}
+
+/// The file's content, and how many bytes at its end were cut short.
+fn parse(bytes: &[u8]) -> Result<(Loaded, usize)> {
+    let rest = bytes
+        .strip_prefix(MAGIC)
+        .ok_or_else(|| anyhow!("not a continuum document file"))?;
+    let line_end = rest
+        .iter()
+        .position(|byte| *byte == b'\n')
+        .ok_or_else(|| anyhow!("no header"))?;
+    let header: Header = serde_json::from_slice(&rest[..line_end]).context("header")?;
+    let actors = header
+        .actors
+        .iter()
+        .map(|hex| {
+            ActorId::try_from(hex.as_str())
+                .map_err(|_| anyhow!("invalid actor {hex:?} in the header"))
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    let mut records = &rest[line_end + 1..];
+    let mut doc: Option<AutoCommit> = None;
+    while !records.is_empty() {
+        let Some((length, after)) = records.split_first_chunk::<4>() else {
+            break;
+        };
+        let length = u32::from_le_bytes(*length) as usize;
+        if after.len() < length {
+            break;
+        }
+        let (record, after) = after.split_at(length);
+        match &mut doc {
+            None => doc = Some(AutoCommit::load(record).context("the full save")?),
+            Some(doc) => {
+                doc.load_incremental(record).context("an increment")?;
+            }
+        }
+        records = after;
+    }
+    let Some(doc) = doc else {
+        bail!("no full save");
+    };
+    let loaded = Loaded {
+        doc: doc.with_actor(ActorId::random()),
+        actors,
+    };
+    Ok((loaded, records.len()))
+}
+
+fn record(bytes: &[u8]) -> Vec<u8> {
+    let mut out = Vec::with_capacity(bytes.len() + 4);
+    out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
+    out.extend_from_slice(bytes);
+    out
+}
+
+/// A whole file: header and one full save.
+fn whole(actors: &[ActorId], save: &[u8]) -> Result<Vec<u8>> {
+    let header = Header {
+        actors: actors.iter().map(ActorId::to_hex_string).collect(),
+    };
+    let mut out = MAGIC.to_vec();
+    out.extend(serde_json::to_vec(&header)?);
+    out.push(b'\n');
+    out.extend(record(save));
+    Ok(out)
+}
+
+/// The document file and what was last written to it.
+pub(crate) struct DocumentFile {
+    path: PathBuf,
+    files: Arc<Files>,
+    append_every: Duration,
+    rewrite_every: Duration,
+    // A std mutex: never held across an await.
+    written: Mutex<Option<Written>>,
+}
+
+#[derive(Debug, Clone)]
+struct Written {
+    epoch: Epoch,
+    /// This node's actors when the file was last rewritten (`RemakeState::own_actors`).
+    actors: usize,
+    heads: Vec<ChangeHash>,
+    /// The full save.
+    base_bytes: u64,
+    /// Appended since.
+    tail_bytes: u64,
+    rewritten_at: Instant,
+}
+
+impl DocumentFile {
+    pub(crate) fn new(
+        path: PathBuf,
+        files: Arc<Files>,
+        append_every: Duration,
+        rewrite_every: Duration,
+    ) -> Self {
+        Self {
+            path,
+            files,
+            append_every,
+            rewrite_every,
+            written: Mutex::new(None),
+        }
+    }
+
+    fn lock(&self) -> MutexGuard<'_, Option<Written>> {
+        self.written
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+}
+
+/// Keeps the document file up to date. Runs for as long as the node does.
+pub(crate) async fn run(state: AppState) {
+    let Some(file) = state.document_file.clone() else {
+        return;
+    };
+    let mut tick = tokio::time::interval(file.append_every);
+    tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+    loop {
+        tick.tick().await;
+        flush(&state).await;
+    }
+}
+
+/// Writes what is due now (see the module docs). Also called at shutdown.
+pub(crate) async fn flush(state: &AppState) {
+    let Some(file) = state.document_file.clone() else {
+        return;
+    };
+    if !state.is_ready().await || state.remake.is_rejoining() {
+        return;
+    }
+    let own = state.remake.own_actors();
+    let (epoch, heads) = {
+        let replica = state.doc.read().await;
+        (replica.epoch, replica.heads().to_vec())
+    };
+    let written = file.lock().clone();
+    let reason = match &written {
+        None => Some("first write"),
+        Some(w) if w.epoch != epoch => Some("new epoch"),
+        Some(w) if w.actors != own.len() => Some("new actor"),
+        Some(w)
+            if w.tail_bytes > w.base_bytes && w.rewritten_at.elapsed() >= file.rewrite_every =>
+        {
+            Some("increments outgrew the full save")
+        }
+        Some(_) => None,
+    };
+    match (reason, written) {
+        (Some(reason), _) => rewrite(state, &file, own, reason).await,
+        (None, Some(written)) if written.heads != heads => append(state, &file, written).await,
+        (None, _) => {}
+    }
+}
+
+async fn rewrite(state: &AppState, file: &Arc<DocumentFile>, own: Vec<ActorId>, reason: &str) {
+    let start = Instant::now();
+    let (epoch, mut copy) = {
+        let replica = state.doc.read().await;
+        (replica.epoch, replica.doc.clone())
+    };
+    let own_count = own.len();
+    let writer = Arc::clone(file);
+    let written = tokio::task::spawn_blocking(move || -> Result<(Vec<ChangeHash>, u64)> {
+        let heads = copy.get_heads();
+        // Only the actors with changes in this document matter for a later rejoin;
+        // the current one in any case.
+        let present: HashSet<ActorId> = copy
+            .get_changes(&[])
+            .iter()
+            .map(|change| change.actor_id().clone())
+            .collect();
+        let current = copy.get_actor().clone();
+        let actors: Vec<ActorId> = own
+            .into_iter()
+            .filter(|actor| *actor == current || present.contains(actor))
+            .collect();
+        let save = copy.save();
+        let bytes = whole(&actors, &save)?;
+        writer
+            .files
+            .write_atomic(FileKind::Document, &writer.path, &bytes)?;
+        Ok((heads, save.len() as u64))
+    })
+    .await;
+    match written {
+        Ok(Ok((heads, base_bytes))) => {
+            debug!(
+                reason,
+                bytes = base_bytes,
+                took_ms = start.elapsed().as_millis() as u64,
+                "Rewrote the document file"
+            );
+            *file.lock() = Some(Written {
+                epoch,
+                actors: own_count,
+                heads,
+                base_bytes,
+                tail_bytes: 0,
+                rewritten_at: Instant::now(),
+            });
+        }
+        Ok(Err(error)) => {
+            error!(path = %file.path.display(), error = %format!("{error:#}"), "Couldn't rewrite the document file");
+        }
+        Err(error) => error!(error = %error, "Document file rewrite task failed"),
+    }
+}
+
+async fn append(state: &AppState, file: &Arc<DocumentFile>, written: Written) {
+    let (heads, increment) = {
+        let mut replica = state.doc.write().await;
+        // Switched since `flush` looked: the next one rewrites the file instead. An
+        // increment of another epoch's document must never be appended.
+        if replica.epoch != written.epoch {
+            return;
+        }
+        let heads = replica.heads().to_vec();
+        (heads, replica.doc.save_after(&written.heads))
+    };
+    let bytes = record(&increment);
+    let writer = Arc::clone(file);
+    let appended = tokio::task::spawn_blocking(move || {
+        writer
+            .files
+            .append(FileKind::Document, &writer.path, &bytes)
+    })
+    .await;
+    match appended {
+        Ok(Ok(())) => {
+            let mut guard = file.lock();
+            if let Some(current) = guard
+                .as_mut()
+                .filter(|current| current.heads == written.heads)
+            {
+                current.heads = heads;
+                current.tail_bytes += increment.len() as u64 + 4;
+                debug!(
+                    bytes = increment.len(),
+                    tail_bytes = current.tail_bytes,
+                    base_bytes = current.base_bytes,
+                    "Appended to the document file"
+                );
+            }
+        }
+        Ok(Err(error)) => {
+            // The file may end in a partial record now; the next rewrite replaces it.
+            error!(path = %file.path.display(), error = %error, "Couldn't append to the document file; rewriting it next");
+            *file.lock() = None;
+        }
+        Err(error) => error!(error = %error, "Document file append task failed"),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use automerge::{ROOT, ReadDoc, transaction::Transactable};
+
+    #[test]
+    fn whole_then_increments() {
+        let mut doc = AutoCommit::new();
+        doc.put(ROOT, "a", 1_i64).unwrap();
+        doc.commit();
+        let actor = doc.get_actor().clone();
+        let heads = doc.get_heads();
+        let mut bytes = whole(std::slice::from_ref(&actor), &doc.save()).unwrap();
+
+        doc.put(ROOT, "b", 2_i64).unwrap();
+        doc.commit();
+        bytes.extend(record(&doc.save_after(&heads)));
+
+        let (loaded, dropped) = parse(&bytes).unwrap();
+        assert_eq!(dropped, 0);
+        assert_eq!(loaded.actors, [actor.clone()]);
+        assert_ne!(loaded.doc.get_actor(), &actor);
+        assert!(loaded.doc.get(ROOT, "b").unwrap().is_some());
+
+        // A record cut short by a crash is dropped.
+        let mut cut = bytes.clone();
+        cut.extend(record(b"whatever")[..6].iter());
+        let (loaded, dropped) = parse(&cut).unwrap();
+        assert_eq!(dropped, 6);
+        assert!(loaded.doc.get(ROOT, "b").unwrap().is_some());
+
+        assert!(parse(b"something else").is_err());
+        assert!(parse(&bytes[..MAGIC.len() + 3]).is_err());
+    }
+}
diff --git a/src/document/mod.rs b/src/document/mod.rs
index ce24815..0da1d67 100644
--- a/src/document/mod.rs
+++ b/src/document/mod.rs
@@ -1,5 +1,6 @@
 //! The replicated document: the replica with its epoch, and the store semantics
 //! (merged-view paths over JSON) the HTTP API and remakes work with.
 
+pub(crate) mod file;
 pub mod replica;
 pub mod store;