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/2`, 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, the first of them its kind
//! (`Record`): first a full save of the document, then increments (the changes after
//! the heads written before), and the signatures of its changes
//! (`document::signatures`), those the file holds none of yet, as
//! `signatures::pack` packs them. A record cut short by a crash is dropped when
//! reading. A file of `continuum-document/1` has records of the document alone,
//! without their kind: a full save, then increments.
//!
//! Written by a worker (`swarm::document_file`), only while the node is ready and not
//! rejoining (a rejoin keeps the old file until it is done):
//! * every `--document-append`, 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`;
//! * 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 super::{ActorId, ChangeHash, Document};
use anyhow::{Context, Result, anyhow, bail};
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::locks::MutexExt;

const MAGIC: &[u8] = b"continuum-document/2\n";
/// The format before signatures, still read.
const MAGIC_1: &[u8] = b"continuum-document/1\n";

/// What a record holds: its first byte.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Record {
    Save = 0,
    Increment = 1,
    Signatures = 2,
}

#[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: Document,
    /// The actors this node wrote with before.
    pub(crate) actors: Vec<ActorId>,
    /// The signatures of its changes, in the order they were written.
    pub(crate) signatures: Vec<(ChangeHash, [u8; 64])>,
}

/// 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, kinds) = match bytes.strip_prefix(MAGIC) {
        Some(rest) => (rest, true),
        None => (
            bytes
                .strip_prefix(MAGIC_1)
                .ok_or_else(|| anyhow!("not a continuum document file"))?,
            false,
        ),
    };
    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<Document> = None;
    let mut signatures = Vec::new();
    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);
        let (kind, record) = if kinds {
            let (kind, record) = record
                .split_first()
                .ok_or_else(|| anyhow!("an empty record"))?;
            let kind = match kind {
                0 => Record::Save,
                1 => Record::Increment,
                2 => Record::Signatures,
                other => bail!("a record of an unknown kind ({other})"),
            };
            (kind, record)
        } else if doc.is_none() {
            (Record::Save, record)
        } else {
            (Record::Increment, record)
        };
        match (kind, &mut doc) {
            (Record::Save, None) => doc = Some(super::load(record).context("the full save")?),
            (Record::Save, Some(_)) => bail!("a second full save"),
            (Record::Increment, Some(doc)) => {
                doc.load_incremental(record).context("an increment")?;
            }
            (Record::Signatures, Some(_)) => signatures.extend(
                super::signatures::unpack(record)
                    .ok_or_else(|| anyhow!("signatures of {} bytes", record.len()))?,
            ),
            (_, None) => bail!("a record before the full save"),
        }
        records = after;
    }
    let Some(doc) = doc else {
        bail!("no full save");
    };
    let loaded = Loaded {
        doc,
        actors,
        signatures,
    };
    Ok((loaded, records.len()))
}

fn record(kind: Record, bytes: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(bytes.len() + 5);
    out.extend_from_slice(&(bytes.len() as u32 + 1).to_le_bytes());
    out.push(kind as u8);
    out.extend_from_slice(bytes);
    out
}

/// A record of `signatures`; none without any.
fn signatures_record(signatures: &[(ChangeHash, [u8; 64])]) -> Vec<u8> {
    if signatures.is_empty() {
        return Vec::new();
    }
    record(Record::Signatures, &super::signatures::pack(signatures))
}

/// A whole file: header, one full save, and the signatures.
fn whole(
    actors: &[ActorId],
    save: &[u8],
    signatures: &[(ChangeHash, [u8; 64])],
) -> 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(Record::Save, save));
    out.extend(signatures_record(signatures));
    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>,
    /// How many of the document's signatures (`Signatures::count`).
    signatures: usize,
    /// 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_or_recover()
    }
}

/// What the document file needs now (`DocumentFile::due`).
pub(crate) enum Due {
    /// Written whole: for this reason.
    Rewrite(&'static str),
    /// What changed since `since` appended, and the signatures after the first
    /// `signatures`, as long as the document is still of `epoch` (an increment of
    /// another epoch's document must never be appended).
    Append {
        epoch: Epoch,
        since: Vec<ChangeHash>,
        signatures: usize,
    },
    Nothing,
}

impl DocumentFile {
    /// How often what changed is appended (`--document-append`).
    pub(crate) fn append_every(&self) -> Duration {
        self.append_every
    }

    /// What is due for a document of `epoch` at `heads`, written by this node's `actors`
    /// (see the module docs).
    /// How many of this node's actors the file holds; None before it was written.
    pub(crate) fn actors(&self) -> Option<usize> {
        self.lock().as_ref().map(|written| written.actors)
    }

    /// With `signatures`, how many the document has (`Signatures::count`); None for a
    /// document whose signatures aren't kept (the one before a rejoin).
    pub(crate) fn due(
        &self,
        epoch: Epoch,
        actors: usize,
        heads: &[ChangeHash],
        signatures: Option<usize>,
    ) -> Due {
        let written = self.lock().clone();
        let reason = match &written {
            None => Some("first write"),
            Some(w) if w.epoch != epoch => Some("new epoch"),
            Some(w) if w.actors != actors => Some("new actor"),
            Some(w)
                if w.tail_bytes > w.base_bytes
                    && w.rewritten_at.elapsed() >= self.rewrite_every =>
            {
                Some("increments outgrew the full save")
            }
            Some(_) => None,
        };
        match (reason, written) {
            (Some(reason), _) => Due::Rewrite(reason),
            (None, Some(written))
                if written.heads != heads
                    || signatures.is_some_and(|count| count != written.signatures) =>
            {
                Due::Append {
                    epoch: written.epoch,
                    since: written.heads,
                    signatures: written.signatures,
                }
            }
            (None, _) => Due::Nothing,
        }
    }

    /// Writes the whole file, atomically: `copy` (a copy of the document, of `epoch`),
    /// its `signatures`, and this node's actors `own` that wrote in it.
    pub(crate) async fn rewrite(
        self: &Arc<Self>,
        epoch: Epoch,
        copy: Document,
        signatures: Vec<(ChangeHash, [u8; 64])>,
        own: Vec<ActorId>,
        reason: &str,
    ) {
        let signature_count = signatures.len();
        let start = Instant::now();
        let own_count = own.len();
        let writer = Arc::clone(self);
        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 = super::save(&copy);
            let bytes = whole(&actors, &save, &signatures)?;
            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"
                );
                *self.lock() = Some(Written {
                    epoch,
                    actors: own_count,
                    heads,
                    signatures: signature_count,
                    base_bytes,
                    tail_bytes: 0,
                    rewritten_at: Instant::now(),
                });
            }
            Ok(Err(error)) => {
                error!(path = %self.path.display(), error = %format!("{error:#}"), "Couldn't rewrite the document file");
            }
            Err(error) => error!(error = %error, "Document file rewrite task failed"),
        }
    }

    /// Appends `increment`: what changed from `since` (the heads written last) to
    /// `heads`; and `signatures`, those of the document the file holds none of yet
    /// (`Due::Append`).
    pub(crate) async fn append(
        self: &Arc<Self>,
        since: Vec<ChangeHash>,
        heads: Vec<ChangeHash>,
        increment: Vec<u8>,
        signatures: Vec<(ChangeHash, [u8; 64])>,
    ) {
        let mut bytes = if heads == since {
            Vec::new()
        } else {
            record(Record::Increment, &increment)
        };
        bytes.extend(signatures_record(&signatures));
        let length = bytes.len() as u64;
        let writer = Arc::clone(self);
        let appended = tokio::task::spawn_blocking(move || {
            writer
                .files
                .append(FileKind::Document, &writer.path, &bytes)
        })
        .await;
        match appended {
            Ok(Ok(())) => {
                let mut guard = self.lock();
                if let Some(current) = guard.as_mut().filter(|current| current.heads == since) {
                    current.heads = heads;
                    current.signatures += signatures.len();
                    current.tail_bytes += length;
                    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 = %self.path.display(), error = %error, "Couldn't append to the document file; rewriting it next");
                *self.lock() = None;
            }
            Err(error) => error!(error = %error, "Document file append task failed"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::document::testing::change;
    use automerge::{ROOT, ReadDoc, transaction::Transactable};

    #[test]
    fn whole_then_increments() {
        let mut doc = Document::new();
        change(&mut doc, |doc| doc.put(ROOT, "a", 1_i64).unwrap());
        let actor = doc.get_actor().clone();
        let heads = doc.get_heads();
        let first = (heads[0], [1; 64]);
        let mut bytes = whole(std::slice::from_ref(&actor), &doc.save(), &[first]).unwrap();

        change(&mut doc, |doc| doc.put(ROOT, "b", 2_i64).unwrap());
        bytes.extend(record(Record::Increment, &doc.save_after(&heads)));
        let second = (doc.get_heads()[0], [2; 64]);
        bytes.extend(signatures_record(&[second]));

        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());
        assert_eq!(loaded.signatures, [first, second]);

        // A record cut short by a crash is dropped.
        let mut cut = bytes.clone();
        cut.extend(record(Record::Increment, 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());
    }

    /// A file of the format before signatures: records of the document alone.
    #[test]
    fn the_format_before_signatures() {
        let mut doc = Document::new();
        change(&mut doc, |doc| doc.put(ROOT, "a", 1_i64).unwrap());
        let heads = doc.get_heads();
        let untagged = |bytes: &[u8]| {
            let mut out = (bytes.len() as u32).to_le_bytes().to_vec();
            out.extend_from_slice(bytes);
            out
        };
        let mut bytes = MAGIC_1.to_vec();
        bytes.extend(b"{\"actors\":[]}\n");
        bytes.extend(untagged(&doc.save()));
        change(&mut doc, |doc| doc.put(ROOT, "b", 2_i64).unwrap());
        bytes.extend(untagged(&doc.save_after(&heads)));
        let (loaded, dropped) = parse(&bytes).unwrap();
        assert_eq!(dropped, 0);
        assert!(loaded.doc.get(ROOT, "b").unwrap().is_some());
        assert!(loaded.signatures.is_empty());
    }
}