path: root/src/compaction/local.rs

//! Per-node, in-memory compaction state: snapshots built ahead of the switch, switch
//! progress for reports, and archived documents of previous epochs.
//!
//! Nothing here is replicated or durable. After a restart the node simply has no
//! archive and rebuilds snapshots when needed.

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;

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

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

#[derive(Default)]
pub(crate) struct Compaction {
    // A std mutex: every access is short and never spans an await.
    inner: Mutex<Inner>,
    /// Set while this node saves its document for a peer; one at a time.
    serving: AtomicBool,
}

#[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 compaction committed.
    first_seen: HashMap<String, Instant>,
    status: HashMap<String, LocalStatus>,
    archive: Vec<Archived>,
    /// The chunk this node is waiting for: (peer, wanted epoch, offset).
    awaiting: Option<(String, Epoch, u64, oneshot::Sender<SnapshotReply>)>,
    catch_up: Option<CatchUp>,
    /// The saved document chunks are served from (one download source at a time).
    served: Option<Arc<Served>>,
}

/// A saved copy of this node's document, served to downloading peers.
pub(crate) struct Served {
    pub(crate) epoch: Epoch,
    pub(crate) bytes: Vec<u8>,
    pub(crate) digest: [u8; 32],
    pub(crate) saved_at: Instant,
}

/// A served copy is dropped this long after it was saved.
const SERVED_KEPT_FOR: Duration = Duration::from_secs(15 * 60);

/// Fetching a newer epoch from a peer.
#[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>,
}

#[derive(Clone, Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct CatchUpReport {
    pub(crate) target: String,
    pub(crate) attempts: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) error: Option<String>,
}

/// What happened to a compaction 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>,
    /// 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 later steps (serving or auditing old epochs); not read yet.
    #[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 Compaction {
    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 compaction 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);
    }

    /// Drops archives past their retention and a served copy nobody needs any more.
    pub(crate) fn prune_archive(&self, current_generation: u64) {
        let mut inner = self.lock();
        prune(&mut inner.archive, current_generation);
        if inner
            .served
            .as_ref()
            .is_some_and(|served| served.saved_at.elapsed() > SERVED_KEPT_FOR)
        {
            inner.served = None;
        }
    }

    /// Registers interest in one chunk from `peer`, replacing any earlier request.
    pub(crate) fn await_reply(
        &self,
        peer: &str,
        wanted: Epoch,
        offset: u64,
    ) -> oneshot::Receiver<SnapshotReply> {
        let (sender, receiver) = oneshot::channel();
        self.lock().awaiting = Some((peer.to_owned(), wanted, offset, sender));
        receiver
    }

    /// Hands a reply to the waiting download. False unless it answers exactly the
    /// pending request (late replies to earlier attempts are dropped).
    pub(crate) fn deliver_reply(&self, peer: &str, reply: SnapshotReply) -> bool {
        let mut inner = self.lock();
        match inner.awaiting.take() {
            Some((expected, wanted, offset, sender))
                if expected == peer && wanted == reply.wanted && offset == reply.offset =>
            {
                sender.send(reply).is_ok()
            }
            other => {
                inner.awaiting = other;
                false
            }
        }
    }

    /// The copy served for `epoch`, if any.
    pub(crate) fn served(&self, epoch: Epoch) -> Option<Arc<Served>> {
        self.lock()
            .served
            .as_ref()
            .filter(|served| served.epoch == epoch)
            .cloned()
    }

    pub(crate) fn set_served(&self, served: Arc<Served>) {
        self.lock().served = Some(served);
    }

    pub(crate) fn stop_awaiting(&self) {
        self.lock().awaiting = None;
    }

    /// Claims the right to serve one snapshot; release with [`Compaction::done_serving`].
    pub(crate) fn try_start_serving(&self) -> bool {
        !self.serving.swap(true, Ordering::AcqRel)
    }

    pub(crate) fn done_serving(&self) {
        self.serving.store(false, Ordering::Release);
    }

    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> {
        self.lock().catch_up.as_ref().map(|catch_up| CatchUpReport {
            target: catch_up.target.to_string(),
            attempts: catch_up.attempts,
            error: catch_up.error.clone(),
        })
    }

    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| {
        let within_epochs = archived.retention.epochs.is_some_and(|epochs| {
            current_generation.saturating_sub(archived.epoch.generation) <= epochs
        });
        let within_time = archived.retention.seconds.is_some_and(|seconds| {
            now.saturating_sub(archived.archived_at) <= seconds.saturating_mul(1000)
        });
        within_epochs || within_time
    });
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn archive_retention() {
        let compaction = Compaction::default();
        let epoch = |generation| Epoch {
            generation,
            id: [generation as u8; 32],
        };

        // Kept by epoch count only.
        compaction.archive(
            epoch(0),
            AutoCommit::new(),
            Retention {
                epochs: Some(1),
                seconds: None,
            },
            1,
        );
        assert_eq!(compaction.archived_epochs().len(), 1);
        compaction.prune_archive(2);
        assert!(compaction.archived_epochs().is_empty());

        // Kept by time while the epoch count has passed.
        compaction.archive(
            epoch(1),
            AutoCommit::new(),
            Retention {
                epochs: Some(0),
                seconds: Some(3600),
            },
            5,
        );
        assert_eq!(compaction.archived_epochs().len(), 1);

        // No limits at all: not kept.
        compaction.archive(
            epoch(2),
            AutoCommit::new(),
            Retention {
                epochs: None,
                seconds: None,
            },
            3,
        );
        assert_eq!(compaction.archived_epochs().len(), 1);
    }
}