path: root/src/remake/local.rs

//! Per-node, in-memory remake state: snapshots built ahead of the switch, switch
//! progress for reports, a rejoin in progress, quarantined changes, archived documents
//! of previous epochs, and the content of remakes (`content`).
//!
//! 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, HashSet};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;

use automerge::AutoCommit;
use serde::Serialize;
use tokio::time::Instant;

use super::records::{Retention, unix_millis};
use super::snapshot::Snapshot;
use crate::document::replica::Epoch;

pub(crate) struct RemakeState {
    /// When this node started (Unix milliseconds): votes it cast before may be lost
    /// (`vote`).
    started_at: u64,
    // A std mutex: every access is short and never spans an await.
    inner: Mutex<Inner>,
}

impl Default for RemakeState {
    fn default() -> Self {
        Self {
            started_at: unix_millis(),
            inner: Mutex::default(),
        }
    }
}

#[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>,
    /// While the rejoin downloads the target epoch's snapshot (`catchup`), sync waits.
    fetching_snapshot: Option<Epoch>,
    quarantine: Vec<Quarantined>,
    /// Content of remakes, by BLAKE3 hash, and when it was stored (served to peers).
    contents: HashMap<[u8; 32], (Arc<Vec<u8>>, Instant)>,
    /// Content being downloaded.
    content_fetches: HashSet<[u8; 32]>,
    /// Remakes this node coordinates now (`coordinator`).
    coordinating: HashSet<String>,
}

/// Content no remake refers to any more is kept this long after it was stored, for
/// nodes that switch late.
const CONTENT_GRACE: Duration = Duration::from_secs(3600);

/// 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 {
    pub(crate) fn started_at(&self) -> u64 {
        self.started_at
    }

    /// Whether this node starts coordinating `id` now (false if it already does).
    pub(crate) fn start_coordinating(&self, id: &str) -> bool {
        self.lock().coordinating.insert(id.to_owned())
    }

    pub(crate) fn end_coordinating(&self, id: &str) {
        self.lock().coordinating.remove(id);
    }

    pub(crate) fn is_coordinating(&self, id: &str) -> bool {
        self.lock().coordinating.contains(id)
    }

    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 content(&self, hash: &[u8; 32]) -> Option<Arc<Vec<u8>>> {
        self.lock()
            .contents
            .get(hash)
            .map(|(bytes, _)| Arc::clone(bytes))
    }

    pub(crate) fn put_content(&self, hash: [u8; 32], bytes: Arc<Vec<u8>>) {
        self.lock().contents.insert(hash, (bytes, Instant::now()));
    }

    /// Marks a download of `hash` as started; false if one already is.
    pub(crate) fn start_content_fetch(&self, hash: [u8; 32]) -> bool {
        self.lock().content_fetches.insert(hash)
    }

    pub(crate) fn end_content_fetch(&self, hash: &[u8; 32]) {
        self.lock().content_fetches.remove(hash);
    }

    pub(crate) fn content_fetch_running(&self, hash: &[u8; 32]) -> bool {
        self.lock().content_fetches.contains(hash)
    }

    /// Drops content that no remake in `referenced` needs, once its grace has passed.
    pub(crate) fn prune_contents(&self, referenced: &HashSet<[u8; 32]>) {
        self.lock().contents.retain(|hash, (_, stored)| {
            referenced.contains(hash) || stored.elapsed() < CONTENT_GRACE
        });
    }

    /// Whether sync should wait: a rejoin is downloading its epoch's snapshot.
    pub(crate) fn fetching_snapshot(&self) -> bool {
        self.lock().fetching_snapshot.is_some()
    }

    pub(crate) fn set_fetching_snapshot(&self, target: Option<Epoch>) {
        self.lock().fetching_snapshot = target;
    }

    /// The rejoin of `target`, if still in progress, gets its full time for the sync
    /// that follows the download.
    pub(crate) fn restart_rejoin_clock(&self, target: Epoch) {
        if let Some(rejoin) = self.lock().rejoin.as_mut().filter(|r| r.target == target) {
            rejoin.started = Instant::now();
        }
    }

    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);
    }
}