moraine 0.4.3

//! Opening a moraine store.
//!
//! [`StoreBuilder`] opens a store on an object store as a read-write [`Db`]
//! or a read-only [`DbReader`], carrying the shared open configuration.
//! Every store is created — and must thereafter be opened — with the
//! tag-byte segment extractor; SlateDB persists the extractor identity and
//! refuses a mismatched open.

use std::{path::PathBuf, sync::Arc, time::Duration};

use object_store::ObjectStore;
use slatedb::{
    Db, DbReader,
    admin::AdminBuilder,
    config::{
        CheckpointOptions, CheckpointScope, DbReaderOptions, ObjectStoreCacheOptions, Settings,
    },
};
use uuid::Uuid;

use crate::{
    error::{Error, Result},
    store::segment::TagSegmentExtractor,
};

/// The default WAL flush cadence when none is configured.
const DEFAULT_FLUSH_INTERVAL: Duration = Duration::from_millis(100);

/// How often a read-only handle polls for new state when none is
/// configured — SlateDB's own default, kept so an unconfigured reader
/// costs exactly what it always has.
const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(10);

/// Creates a checkpoint of everything `db` has committed, expiring after
/// `lifetime` (never, if `None`), and reports its id. The scope is every
/// write the handle has taken, not only the already-durable ones, so a
/// commit that has returned is inside the checkpoint it precedes.
pub(crate) async fn create_checkpoint(db: &Db, lifetime: Option<Duration>) -> Result<Uuid> {
    let created = db
        .create_checkpoint(
            CheckpointScope::All,
            &CheckpointOptions {
                lifetime,
                ..CheckpointOptions::default()
            },
        )
        .await
        .map_err(Error::from)?;
    Ok(created.id)
}

/// Opens a moraine store on `object_store` — a read-write [`Db`] via
/// [`open_writer`](Self::open_writer) or a read-only [`DbReader`] via
/// [`open_reader`](Self::open_reader) — carrying the shared open
/// configuration: the WAL flush cadence (writer only) and the on-disk cache.
pub(crate) struct StoreBuilder<'a> {
    path: &'a str,
    object_store: Arc<dyn ObjectStore>,
    flush_interval: Duration,
    poll_interval: Duration,
    cache_dir: Option<PathBuf>,
    checkpoint: Option<Uuid>,
}

impl<'a> StoreBuilder<'a> {
    /// A builder for the store at `path` on `object_store`, with the default
    /// flush cadence and no on-disk cache.
    pub(crate) fn new(path: &'a str, object_store: Arc<dyn ObjectStore>) -> Self {
        Self {
            path,
            object_store,
            flush_interval: DEFAULT_FLUSH_INTERVAL,
            poll_interval: DEFAULT_POLL_INTERVAL,
            cache_dir: None,
            checkpoint: None,
        }
    }

    /// Sets how often a read-only handle polls for new state. Reader only —
    /// a writer reads its own writes.
    pub(crate) fn poll_interval(mut self, poll_interval: Duration) -> Self {
        self.poll_interval = poll_interval;
        self
    }

    /// Sets the WAL flush cadence. Durable commits wait for the next flush,
    /// so this bounds per-commit latency; smaller values mean more frequent
    /// (on S3, costlier) object-store PUTs. Zero flushes continuously (no
    /// timer), so a durable commit waits only on the object-store PUT — the
    /// lowest latency, at the cost of a busy flush loop. Writer only — a
    /// reader never flushes.
    pub(crate) fn flush_interval(mut self, flush_interval: Duration) -> Self {
        self.flush_interval = flush_interval;
        self
    }

    /// Sets the local directory backing SlateDB's on-disk block cache. When
    /// set, warm reads skip repeat object-store GETs and survive process
    /// restarts — worthwhile for remote (`s3://`) stores, redundant for local
    /// ones. `None` (the default) uses only SlateDB's in-memory cache.
    pub(crate) fn cache_dir(mut self, cache_dir: Option<PathBuf>) -> Self {
        self.cache_dir = cache_dir;
        self
    }

    /// Pins the reader to an existing checkpoint instead of the latest
    /// manifest. Reader only, and the only truly zero-write open: a
    /// latest-mode reader creates and refreshes a checkpoint of its own,
    /// which is a manifest write.
    pub(crate) fn checkpoint(mut self, checkpoint: Option<Uuid>) -> Self {
        self.checkpoint = checkpoint;
        self
    }

    /// Opens (or creates) the store as a read-write [`Db`].
    pub(crate) async fn open_writer(&self) -> Result<Db> {
        let settings = self.settings();
        Db::builder(self.path, Arc::clone(&self.object_store))
            .with_settings(settings)
            .with_segment_extractor(Arc::new(TagSegmentExtractor))
            .build()
            .await
            .map_err(Error::from)
    }

    /// Opens the store read-only as a [`DbReader`]. A `DbReader` never opens
    /// the writer `Db`, so it never fences a live writer. The flush cadence,
    /// if set, is ignored.
    ///
    /// Without a checkpoint the reader follows the latest manifest, which
    /// costs a manifest write on open and a refresh for the reader's
    /// lifetime. Pinned to a [`checkpoint`](Self::checkpoint) it writes
    /// nothing at all, and reads the fixed cut that checkpoint names.
    pub(crate) async fn open_reader(&self) -> Result<DbReader> {
        let options = DbReaderOptions {
            object_store_cache_options: self.cache_options(),
            manifest_poll_interval: self.poll_interval,
            ..Default::default()
        };
        let mut builder = DbReader::builder(self.path, Arc::clone(&self.object_store))
            .with_segment_extractor(Arc::new(TagSegmentExtractor))
            .with_options(options);
        if let Some(checkpoint) = self.checkpoint {
            builder = builder.with_checkpoint_id(checkpoint);
        }
        builder.build().await.map_err(Error::from)
    }

    /// Deletes the checkpoint `checkpoint`, unpinning whatever it held
    /// against SlateDB's garbage collection. Readers still open against it
    /// keep serving from the objects they have already resolved.
    pub(crate) async fn delete_checkpoint(&self, checkpoint: Uuid) -> Result<()> {
        AdminBuilder::new(self.path, Arc::clone(&self.object_store))
            .build()
            .delete_checkpoint(checkpoint)
            .await
            .map_err(Error::from)
    }

    /// Every checkpoint the store's manifest currently carries, oldest
    /// first — reader-established ones included.
    pub(crate) async fn list_checkpoints(&self) -> Result<Vec<Uuid>> {
        let checkpoints = AdminBuilder::new(self.path, Arc::clone(&self.object_store))
            .build()
            .list_checkpoints(None)
            .await
            .map_err(Error::from)?;
        Ok(checkpoints.into_iter().map(|c| c.id).collect())
    }

    /// SlateDB settings for a writer. A zero flush interval flushes
    /// continuously rather than on a timer: durable commits then wait only
    /// on the object-store PUT, at the cost of a busy flush loop.
    fn settings(&self) -> Settings {
        Settings {
            flush_interval: Some(self.flush_interval),
            object_store_cache_options: self.cache_options(),
            ..Default::default()
        }
    }

    /// SlateDB's on-disk block cache options: a disk-backed cache under
    /// `cache_dir` when set, otherwise none (in-memory cache only). Every
    /// field but `root_folder` stays at SlateDB's defaults (16 GiB, 4 MiB
    /// parts).
    fn cache_options(&self) -> ObjectStoreCacheOptions {
        ObjectStoreCacheOptions {
            root_folder: self.cache_dir.clone(),
            ..Default::default()
        }
    }
}

#[cfg(test)]
mod tests {
    use object_store::memory::InMemory;
    use slatedb::{IsolationLevel, config::WriteOptions};

    use super::*;
    use crate::store::key::{self, Key, SysKey};

    fn memory_store() -> Arc<dyn ObjectStore> {
        Arc::new(InMemory::new())
    }

    /// A commit-shaped transaction spanning several subspaces lands
    /// atomically, and per-subspace prefix scans see exactly their own
    /// segment's keys (multi-segment batches satisfy the antichain rule).
    #[tokio::test]
    async fn multi_subspace_transaction_and_prefix_scans() {
        let db = StoreBuilder::new("test/store", memory_store())
            .open_writer()
            .await
            .unwrap();

        let head = Key::Sys(SysKey::Head).encode();
        let snapshot = Key::Snapshot { snapshot_id: 1 }.encode();
        let table = Key::current(key::EntityKey::Table { table_id: 7 }).encode();

        let tx = db.begin(IsolationLevel::Snapshot).await.unwrap();
        tx.put(&head, b"head").unwrap();
        tx.put(&snapshot, b"snap").unwrap();
        tx.put(&table, b"table").unwrap();
        tx.commit_with_options(&WriteOptions {
            await_durable: true,
            ..Default::default()
        })
        .await
        .unwrap();

        assert_eq!(db.get(&head).await.unwrap().unwrap().as_ref(), b"head");

        // Each subspace scan returns exactly its own keys.
        let mut iter = db
            .scan_prefix(key::subspace_prefix(key::Subspace::Current), ..)
            .await
            .unwrap();
        let entry = iter.next().await.unwrap().unwrap();
        assert_eq!(entry.key.as_ref(), table.as_slice());
        assert!(iter.next().await.unwrap().is_none());

        let mut iter = db
            .scan_prefix(key::subspace_prefix(key::Subspace::Snapshot), ..)
            .await
            .unwrap();
        let entry = iter.next().await.unwrap().unwrap();
        assert_eq!(entry.key.as_ref(), snapshot.as_slice());
        assert!(iter.next().await.unwrap().is_none());

        db.close().await.unwrap();
    }

    /// A zero flush interval is allowed: SlateDB flushes continuously
    /// rather than on a timer, so the store opens and a durable write lands.
    #[tokio::test]
    async fn zero_flush_interval_opens_a_working_store() {
        let db = StoreBuilder::new("test/store", memory_store())
            .flush_interval(Duration::ZERO)
            .open_writer()
            .await
            .unwrap();

        let head = Key::Sys(SysKey::Head).encode();
        db.put(&head, b"head").await.unwrap();
        assert_eq!(db.get(&head).await.unwrap().unwrap().as_ref(), b"head");

        db.close().await.unwrap();
    }

    /// An explicit flush interval reaches the SlateDB builder: the store
    /// opens, and a durable commit still lands.
    #[tokio::test]
    async fn explicit_flush_interval_opens_a_working_store() {
        let db = StoreBuilder::new("test/store", memory_store())
            .flush_interval(Duration::from_millis(1))
            .open_writer()
            .await
            .unwrap();

        let head = Key::Sys(SysKey::Head).encode();
        db.put(&head, b"head").await.unwrap();
        assert_eq!(db.get(&head).await.unwrap().unwrap().as_ref(), b"head");

        db.close().await.unwrap();
    }

    /// SlateDB persists the extractor identity: reopening the store
    /// without it is refused rather than silently mis-segmented.
    #[tokio::test]
    async fn reopen_without_extractor_is_refused() {
        let object_store = memory_store();
        let db = StoreBuilder::new("test/store", object_store.clone())
            .open_writer()
            .await
            .unwrap();
        db.put(&Key::Sys(SysKey::Head).encode(), b"head")
            .await
            .unwrap();
        db.close().await.unwrap();

        let bare = Db::builder("test/store", object_store).build().await;
        assert!(bare.is_err(), "unsegmented reopen must be refused");
    }

    /// A configured `cache_dir` reaches SlateDB: a fresh `DbReader`, whose
    /// in-memory cache has never seen the store's blocks, serves committed
    /// data and in doing so populates the on-disk cache directory.
    #[tokio::test]
    async fn cache_dir_backs_the_reader_with_an_on_disk_cache() {
        let object_store = memory_store();
        let cache = std::env::temp_dir().join(format!("moraine-cache-test-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&cache);

        let head = Key::Sys(SysKey::Head).encode();
        let db = StoreBuilder::new("s", object_store.clone())
            .flush_interval(Duration::from_millis(1))
            .cache_dir(Some(cache.clone()))
            .open_writer()
            .await
            .unwrap();
        db.put(&head, b"head").await.unwrap();
        db.close().await.unwrap();

        // The reader is cold: it must GET the store's blocks, which the disk
        // cache records under `cache`.
        let reader = StoreBuilder::new("s", object_store)
            .cache_dir(Some(cache.clone()))
            .open_reader()
            .await
            .unwrap();
        assert_eq!(reader.get(head).await.unwrap().unwrap().as_ref(), b"head");
        reader.close().await.unwrap();

        let populated = std::fs::read_dir(&cache).is_ok_and(|mut entries| entries.next().is_some());
        assert!(populated, "expected an on-disk cache under {cache:?}");
        let _ = std::fs::remove_dir_all(&cache);
    }
}