persistence.rs

use std::{
    io,
    num::{NonZeroU64, NonZeroUsize},
    sync::Arc,
};

use async_trait::async_trait;
use spacetimedb_commitlog::SizeOnDisk;
use spacetimedb_durability::{DurabilityExited, TxOffset};
use spacetimedb_paths::server::ServerDataDir;
use spacetimedb_snapshot::DynSnapshotRepo;

use crate::{messages::control_db::Database, util::asyncify};
use spacetimedb_runtime::Handle;

use super::{
    relational_db::{self, Txdata},
    snapshot::{self, SnapshotDatabaseState, SnapshotWorker},
};

/// Local durability configuration exposed through server config.
#[derive(Clone, Copy, Debug, Default, serde::Deserialize, PartialEq, Eq)]
#[serde(rename_all = "kebab-case")]
pub struct DurabilityConfig {
    #[serde(default)]
    pub commitlog: CommitlogConfig,
}

/// Commitlog configuration exposed through server config.
#[derive(Clone, Copy, Debug, Default, serde::Deserialize, PartialEq, Eq)]
#[serde(rename_all = "kebab-case")]
pub struct CommitlogConfig {
    pub log_format_version: Option<u8>,
    pub max_segment_size: Option<NonZeroU64>,
    #[serde(alias = "offset-interval-bytes")]
    pub offset_index_interval_bytes: Option<NonZeroU64>,
    #[serde(alias = "offset-index-require-fsync")]
    pub offset_index_require_segment_fsync: Option<bool>,
    pub preallocate_segments: Option<bool>,
    pub write_buffer_size: Option<NonZeroUsize>,
}

impl DurabilityConfig {
    fn into_options(self) -> spacetimedb_durability::local::Options {
        let mut opts = spacetimedb_durability::local::Options::default();
        self.commitlog.apply_to(&mut opts.commitlog);
        opts
    }
}

impl CommitlogConfig {
    fn apply_to(self, opts: &mut spacetimedb_commitlog::Options) {
        if let Some(log_format_version) = self.log_format_version {
            opts.log_format_version = log_format_version;
        }
        if let Some(max_segment_size) = self.max_segment_size {
            opts.max_segment_size = max_segment_size.get();
        }
        if let Some(offset_index_interval_bytes) = self.offset_index_interval_bytes {
            opts.offset_index_interval_bytes = offset_index_interval_bytes;
        }
        if let Some(offset_index_require_segment_fsync) = self.offset_index_require_segment_fsync {
            opts.offset_index_require_segment_fsync = offset_index_require_segment_fsync;
        }
        if let Some(preallocate_segments) = self.preallocate_segments {
            opts.preallocate_segments = preallocate_segments;
        }
        if let Some(write_buffer_size) = self.write_buffer_size {
            opts.write_buffer_size = write_buffer_size.get();
        }
    }
}

/// [spacetimedb_durability::Durability] impls with a [`Txdata`] transaction
/// payload, suitable for use in the [`relational_db::RelationalDB`].
pub type Durability = dyn spacetimedb_durability::Durability<TxData = Txdata>;

/// A function to determine the size on disk of the durable state of the
/// local database instance. This is used for metrics and energy accounting
/// purposes.
///
/// It is not part of the [`Durability`] trait because it must report disk
/// usage of the local instance only, even if exclusively remote durability is
/// configured or the database is in follower state.
pub type DiskSizeFn = Arc<dyn Fn() -> io::Result<SizeOnDisk> + Send + Sync>;

/// Persistence services for a database.
pub struct Persistence {
    /// The [Durability] to use, for persisting transactions.
    pub durability: Arc<Durability>,
    /// The [DiskSizeFn].
    ///
    /// Currently the expectation is that the reported size is the commitlog
    /// size only.
    pub disk_size: DiskSizeFn,
    /// An optional [SnapshotWorker].
    ///
    /// The current expectation is that snapshots are only enabled for
    /// persistent (as opposed to in-memory) databases. This is enforced by
    /// this type.
    pub snapshots: Option<SnapshotWorker>,
    /// Runtime onto which durability-related tasks shall be spawned.
    pub runtime: Handle,
}

impl Persistence {
    /// Convenience constructor of a [Persistence] that handles boxing.
    pub fn new(
        durability: impl spacetimedb_durability::Durability<TxData = Txdata> + 'static,
        disk_size: impl Fn() -> io::Result<SizeOnDisk> + Send + Sync + 'static,
        snapshots: Option<SnapshotWorker>,
        runtime: tokio::runtime::Handle,
    ) -> Self {
        Self::new_with_runtime(durability, disk_size, snapshots, Handle::tokio(runtime))
    }

    pub fn new_with_runtime(
        durability: impl spacetimedb_durability::Durability<TxData = Txdata> + 'static,
        disk_size: impl Fn() -> io::Result<SizeOnDisk> + Send + Sync + 'static,
        snapshots: Option<SnapshotWorker>,
        runtime: Handle,
    ) -> Self {
        Self {
            durability: Arc::new(durability),
            disk_size: Arc::new(disk_size),
            snapshots,
            runtime,
        }
    }

    /// If snapshots are enabled, get the [SnapshotRepo] they are stored in.
    pub fn snapshot_repo(&self) -> Option<Arc<DynSnapshotRepo>> {
        self.snapshots.as_ref().map(|worker| worker.snapshot_repo())
    }

    /// Get the [TxOffset] reported as durable by the [Durability] impl.
    ///
    /// Returns `Ok(None)` if no offset is durable yet, and `Err(DurabilityExited)`
    /// if the [Durability] has shut down already.
    pub fn durable_tx_offset(&self) -> Result<Option<TxOffset>, DurabilityExited> {
        self.durability.durable_tx_offset().get()
    }

    /// Initialize the [SnapshotWorker], no-op if snapshots are not enabled.
    pub(super) fn set_snapshot_state(&self, state: SnapshotDatabaseState) {
        if let Some(worker) = &self.snapshots {
            worker.set_state(state)
        }
    }

    /// Convenience to deconstruct an [Option<Self>] into parts.
    ///
    /// Returns `(Some(durability), Some(disk_size), Option<SnapshotWorker>, Some(runtime))`
    /// if `this` is `Some`, and `(None, None, None, None)` if `this` is `None`.
    pub(super) fn unzip(
        this: Option<Self>,
    ) -> (
        Option<Arc<Durability>>,
        Option<DiskSizeFn>,
        Option<SnapshotWorker>,
        Option<Handle>,
    ) {
        this.map(
            |Self {
                 durability,
                 disk_size,
                 snapshots,
                 runtime,
             }| (Some(durability), Some(disk_size), snapshots, Some(runtime)),
        )
        .unwrap_or_default()
    }
}

/// A persistence provider is a "factory" of sorts that can produce [Persistence]
/// services for a given replica.
///
/// The [crate::host::HostController] uses this to obtain [Persistence]s from
/// an external source, and construct [relational_db::RelationalDB]s with it.
///
/// This is an `async_trait` to allow it to be used as a trait object.
#[async_trait]
pub trait PersistenceProvider: Send + Sync {
    async fn persistence(&self, database: &Database, replica_id: u64) -> anyhow::Result<Persistence>;
}

/// The standard [PersistenceProvider] for non-replicated databases.
///
/// [Persistence] services are provided for the local [ServerDataDir].
///
/// Note that its [PersistenceProvider::persistence] impl will spawn a
/// background task that [compresses] older commitlog segments whenever a
/// snapshot is taken.
///
/// [compresses]: relational_db::snapshot_watching_commitlog_compressor
pub struct LocalPersistenceProvider {
    data_dir: Arc<ServerDataDir>,
    durability: DurabilityConfig,
}

impl LocalPersistenceProvider {
    pub fn new(data_dir: impl Into<Arc<ServerDataDir>>) -> Self {
        Self {
            data_dir: data_dir.into(),
            durability: DurabilityConfig::default(),
        }
    }

    pub fn with_durability_config(mut self, durability: DurabilityConfig) -> Self {
        self.durability = durability;
        self
    }
}

#[async_trait]
impl PersistenceProvider for LocalPersistenceProvider {
    async fn persistence(&self, database: &Database, replica_id: u64) -> anyhow::Result<Persistence> {
        let replica_dir = self.data_dir.replica(replica_id);
        let snapshot_dir = replica_dir.snapshots();
        let runtime = Handle::tokio_current();

        let database_identity = database.database_identity;
        let snapshot_worker =
            asyncify(move || relational_db::open_snapshot_repo(snapshot_dir, database_identity, replica_id))
                .await
                .map(|repo| SnapshotWorker::new(repo, snapshot::Compression::Enabled, runtime.clone()))?;
        let (durability, disk_size) = relational_db::local_durability_with_options(
            replica_dir,
            runtime.clone(),
            Some(&snapshot_worker),
            self.durability.into_options(),
        )
        .await?;

        tokio::spawn(relational_db::snapshot_watching_commitlog_compressor(
            snapshot_worker.subscribe(),
            None,
            None,
            durability.clone(),
        ));

        Ok(Persistence {
            durability,
            disk_size,
            snapshots: Some(snapshot_worker),
            runtime,
        })
    }
}