path: root/src/document/file.rs

//! 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());
    }
}