gnitz-store 0.1.4

//! The relation registry — what relations this process holds, and the stores
//! behind them.
//!
//! Every relation enters through [`RelationRegistry::register`].

use gnitz_foundation::env::env_num;
use rustc_hash::FxHashMap;

use gnitz_zset::algebra::index_entries;
use gnitz_zset::schema::key::{key_range_between_cuts, KeyCut};
use gnitz_zset::schema::{index_spec_and_schema, KeySpec, SchemaDescriptor};

pub use crate::storage::Cut;
use crate::storage::{RecoverySource, StoreBudgets, Table};
use gnitz_wire::{PkColList, PkKeys, ViewProps};
use gnitz_zset::repr::{Batch, PkSetGather, ReadCursor, StorageError, StoredRow};
use gnitz_zset::schema::{Placement, Slot};

mod build;
mod circuit_state;
mod delta;
mod dirs;
mod disk_usage;
mod ingest;
mod repartition;
mod store;
mod store_lifecycle;
mod unique_pk;

pub use circuit_state::{CircuitState, StateIdx, StateLayout};
pub(crate) use delta::{delta_round, delta_round_prefix};
use dirs::ensure_dir;
pub use dirs::{lock_data_dir, relation_dir, relations_dir, ChildAddr, ChildKind, DirLock};
pub use disk_usage::{disk_usage, DiskUsage};
use store::Store;

// ---------------------------------------------------------------------------
// Secondary index
// ---------------------------------------------------------------------------

/// Who holds a secondary-index circuit; it lives while one remains.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum IndexClaim {
    /// An FK circuit on the owner's column list.
    ForeignKey,
    /// A catalog index.
    Index { id: u64, unique: bool },
}

/// A secondary index on a column list of one relation. Owns this process's store
/// for it — dropping the index drops the store.
pub struct SecondaryIndex {
    store: Store,
    /// The index's declared column list, in order. A 1-element list is the
    /// single-column case. Dedup/lookup is exact ordered-list equality on
    /// `cols.as_slice()`; order is significant (it drives leading-prefix seeks).
    cols: PkColList,
    /// Full-arity span-encode plan, precomputed at registration so the per-push
    /// consumers do no per-call spec rebuild. It survives every column ALTER of
    /// the owner — [`RelationRegistry::swap_schema`] rejects any descriptor that
    /// would not leave it valid.
    key_spec: KeySpec,
    /// Every holder of this circuit.
    claims: Vec<IndexClaim>,
    /// Whether `cols` covers the owner's PK, so the index can never collide.
    covers_pk: bool,
}

impl SecondaryIndex {
    /// By value, not as a slice: `PkColList` is `Copy`, so the list can be keyed
    /// on or held past the borrow of the registry that produced it.
    pub fn cols(&self) -> PkColList {
        self.cols
    }

    pub fn schema(&self) -> SchemaDescriptor {
        *self.store.schema()
    }

    pub fn key_spec(&self) -> KeySpec {
        self.key_spec
    }

    pub fn is_unique(&self) -> bool {
        self.claims
            .iter()
            .any(|c| matches!(c, IndexClaim::Index { unique: true, .. }))
    }

    /// Every holder of this circuit.
    pub fn claims(&self) -> &[IndexClaim] {
        &self.claims
    }

    /// Non-compacting cursor over this index's store, in the index schema.
    pub fn cursor(&self) -> ReadCursor {
        self.store.held().open_cursor()
    }

    /// Every stored entry of `spans`, each this index's leading key bytes.
    pub(crate) fn gather(&self, spans: PkKeys) -> PkSetGather {
        self.store.held().gather(spans, Cut::Now)
    }

    /// A cursor positioned on the key band `r` names under this index's key spec.
    pub(crate) fn cursor_over(&self, r: &gnitz_wire::KeyRange) -> ReadCursor {
        let t = self.store.held();
        t.range_cursor(self.key_spec.range_keys(t.schema().pk_stride(), r))
    }

    /// Project `source` into this index's layout and ingest the result.
    /// `Ok(false)` when the projection was empty: the key spec rejects rows, so
    /// a non-empty `source` can still project to nothing.
    pub(crate) fn project_and_ingest(&mut self, source: &Batch) -> Result<bool, StorageError> {
        let table = self.store.held_mut();
        let projected = index_entries(source, &self.key_spec, table.schema());
        if projected.is_empty() {
            return Ok(false);
        }
        table.ingest_owned_batch(projected).map(|()| true)
    }

    /// Whether this process's store of the index came back from a checkpoint
    /// manifest at its open.
    pub fn resumed(&self) -> bool {
        self.store.held().resumed_from_checkpoint()
    }
}

// ---------------------------------------------------------------------------
// Relation kind — what a top-level relation *is*
// ---------------------------------------------------------------------------

/// What a top-level relation *is*.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum RelationKind {
    /// System catalog table: durable, single-partition, never rebuilt from
    /// upstream sources (it has none; recovery is LSN-gated SAL replay).
    SystemCatalog,
    /// User base table: durable, never rebuilt from upstream
    /// sources; owns a DML-enforced PK, so `enforce_unique_pk` runs on every
    /// ingest and its accumulated per-PK weight is always in {0, 1}.
    BaseTable,
    /// Materialised view, with its `WITH (…)` options: ephemeral, resumed from a
    /// generation-stamped manifest when the ephemeral checkpoint round left one.
    ///
    /// Not necessarily *rebuilt* from its sources: a mirror registers its fed
    /// copies under this kind too.
    View(ViewProps),
    /// Ingestion point: storeless, append-only. Pushed rows exist
    /// only as the deltas they produce; nothing is retained, so there is nothing
    /// to recover and no store to read.
    Stream,
}

impl RelationKind {
    /// What a client is told this relation is.
    #[inline]
    pub fn class(self) -> gnitz_wire::RelClass {
        match self {
            RelationKind::Stream => gnitz_wire::RelClass::Stream,
            RelationKind::View(props) => props.into(),
            RelationKind::BaseTable | RelationKind::SystemCatalog => gnitz_wire::RelClass::Table,
        }
    }

    /// What to call this relation in a message to the user.
    #[inline]
    pub fn noun(self) -> &'static str {
        self.class().noun()
    }

    /// True iff this is a user base table.
    #[inline]
    pub fn is_base_table(self) -> bool {
        matches!(self, RelationKind::BaseTable)
    }

    /// True iff this relation's rows arrive from a client push rather than being
    /// derived by a circuit.
    #[inline]
    pub fn is_ingestion_point(self) -> bool {
        matches!(self, RelationKind::BaseTable | RelationKind::Stream)
    }

    /// True iff this is a materialised view.
    #[inline]
    pub fn is_view(self) -> bool {
        matches!(self, RelationKind::View(_))
    }

    /// Whether this kind of relation may carry a secondary index.
    fn admits_index(self) -> bool {
        match self {
            RelationKind::BaseTable | RelationKind::View(ViewProps::Plain | ViewProps::Fed { .. }) => true,
            // A bounded view's sweep drops the payload an entry is projected from.
            RelationKind::View(ViewProps::Bounded { .. }) | RelationKind::Stream | RelationKind::SystemCatalog => false,
        }
    }

    /// Whether this relation keeps a delta feed, whether or not this process holds
    /// its store.
    #[inline]
    pub fn has_delta_feed(self) -> bool {
        matches!(self, RelationKind::View(ViewProps::Fed { .. }))
    }

    /// Whether a capacity bounds this relation's registered shard bytes, so its
    /// sweep may leave skeleton rows behind.
    #[inline]
    pub fn is_bounded(self) -> bool {
        matches!(self, RelationKind::View(ViewProps::Bounded { .. }))
    }
}

// ---------------------------------------------------------------------------
// Residency — what this process is to the stores it registered
// ---------------------------------------------------------------------------

/// What a process is to the stores it registered.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Residency {
    /// The server's master: holds the system families' stores and no user store,
    /// which is what lets it relay and reclaim any rank's children.
    Master,
    /// A standalone host — a mirror, a unit-test engine or fixture: owns every
    /// store it registers, at its slot, and may read any rank's manifests.
    Origin,
    /// A forked worker, at its own slot. Owns its stores; speaks for itself.
    Worker,
}

impl Residency {
    /// True iff this process opens a store for every relation it registers.
    #[inline]
    pub(crate) fn owns_stores(self) -> bool {
        !matches!(self, Residency::Master)
    }
}

// ---------------------------------------------------------------------------
// Relation — one relation in this process
// ---------------------------------------------------------------------------

/// One relation in this process: its identity, its shape, and this process's
/// store for it.
pub struct Relation {
    /// The relation id — the key it is registered under, held here too so a
    /// borrowed `Relation` names itself: nothing that has one needs the id
    /// threaded in beside it to log or to phrase an error.
    id: u64,
    store: Store,
    /// A stream's pushes since its last seal. Nothing reads them before it, so
    /// they are held as they arrived.
    unsealed: Option<Batch>,
    delta: Option<Box<Table>>,
    indexes: Vec<SecondaryIndex>,
    kind: RelationKind,
    placement: Placement,
    pk_repeats: bool,
}

impl Relation {
    /// This relation's delta store, if this process serves its feed.
    pub(crate) fn delta(&self) -> Option<&Table> {
        self.delta.as_deref()
    }

    /// By value, for the reason [`SecondaryIndex::cols`] is: the `tables` borrow
    /// ends with the call, so the descriptor survives a later `&mut` on the
    /// registry.
    pub fn schema(&self) -> SchemaDescriptor {
        *self.store.schema()
    }

    pub fn id(&self) -> u64 {
        self.id
    }

    pub fn kind(&self) -> RelationKind {
        self.kind
    }

    /// Where this relation's rows live.
    pub fn placement(&self) -> Placement {
        self.placement
    }

    /// Whether two of its rows may share a PK.
    pub fn pk_repeats(&self) -> bool {
        self.pk_repeats
    }

    /// The unique secondary indexes a write must still check: one covering the
    /// PK cannot collide, so it is not among them.
    pub fn unique_indexes_to_check(&self) -> impl Iterator<Item = &SecondaryIndex> + '_ {
        self.indexes.iter().filter(|ic| ic.is_unique() && !ic.covers_pk)
    }

    /// Non-compacting cursor over this relation's store.
    pub fn cursor(&self) -> ReadCursor {
        self.table().open_cursor()
    }

    /// Every live row of `keys` — whole PKs, or the same leading columns of one,
    /// each then naming every row it prefixes — among the rows `cut` reads.
    pub fn gather(&self, keys: PkKeys, cut: Cut) -> PkSetGather {
        self.table().gather(keys, cut)
    }

    /// A cursor for probing at the PKs of `keys` — whole PKs, or the same leading
    /// columns of one — positioned on the first, over the rows `cut` reads.
    pub fn cursor_for_keys(&self, keys: &Batch, cut: Cut) -> ReadCursor {
        self.table().cursor_for_keys(keys, cut)
    }

    /// A cursor for probing at the keys in `[first, last]` — whole PKs, or the
    /// same leading bytes of one — positioned on the first, over the rows `cut`
    /// reads.
    pub fn cursor_between(&self, first: &[u8], last: &[u8], cut: Cut) -> ReadCursor {
        self.table().cursor_between(first, last, cut)
    }

    /// Visit every positive-weight row whose OPK key begins with `prefix`.
    pub fn for_each_positive_with_prefix(&self, prefix: &[u8], f: impl FnMut(&ReadCursor)) {
        let table = self.table();
        let band = key_range_between_cuts(
            KeyCut::min_of(prefix),
            KeyCut::above(prefix),
            table.schema().pk_stride(),
        );
        table.range_cursor(band).for_each_positive_while(|_| true, f);
    }

    /// Materialize every row of this relation's store whose net weight is non-zero.
    pub fn full_scan(&self) -> std::rc::Rc<Batch> {
        self.table().full_scan()
    }

    /// The net weight at `key`, and the live row if there is one.
    pub fn live_row_at(&self, key: &[u8]) -> (i64, Option<StoredRow>) {
        self.table().live_row_at(key)
    }

    /// Every secondary index on this relation, in registration order.
    pub fn indexes(&self) -> &[SecondaryIndex] {
        &self.indexes
    }

    /// The secondary index covering exactly `cols`, if one exists. The index
    /// list is deduped by ordered column list, so at most one entry matches.
    pub fn index_on(&self, cols: &[u32]) -> Option<&SecondaryIndex> {
        self.indexes.iter().find(|ix| ix.cols.as_slice() == cols)
    }

    /// Whether rows sit above this relation's cut.
    pub(crate) fn has_pending(&self) -> bool {
        match &self.store {
            Store::Held(t) => t.has_pending(),
            Store::Absent(_) => self.unsealed.is_some(),
        }
    }

    /// This process's store of the relation; panics when it holds none.
    pub(crate) fn table(&self) -> &Table {
        self.store.held()
    }
}

// ---------------------------------------------------------------------------
// The relation to register
// ---------------------------------------------------------------------------

/// A relation to [`RelationRegistry::register`].
#[derive(Clone, Copy)]
pub struct RelationSpec {
    pub id: u64,
    pub kind: RelationKind,
    pub schema: SchemaDescriptor,
    pub placement: Placement,
    /// Whether two of its rows may share a PK. `false` is a promise the index
    /// gather relies on, so a host that cannot vouch for it passes `true`.
    pub pk_repeats: bool,
}

// ---------------------------------------------------------------------------
// RelationRegistry
// ---------------------------------------------------------------------------

/// Everything a registry is tuned by. `Default` is the production value of every
/// field; a host overrides them through [`StoreConfig::from_env`].
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct StoreConfig {
    /// The RAM-tier ceiling every store this registry opens is sized to.
    pub ram_tier_bytes: usize,
    /// Rows per `drain_chunk` on every chunked scan this process drives; bounds
    /// peak scan memory at O(chunk × row_width). Never zero.
    pub scan_chunk_rows: usize,
    /// Per-worker distinct-group cap for the ad-hoc aggregate fold; bounds the
    /// accumulator matrix at `cap × aggregates × size_of::<Accumulator>()`.
    pub adhoc_group_cap: usize,
    /// Bytes a key-span sort holds in RAM before it spills a run.
    pub key_spans_spill_bytes: usize,
}

impl Default for StoreConfig {
    fn default() -> Self {
        StoreConfig {
            ram_tier_bytes: crate::storage::DEFAULT_RAM_TIER_BYTES,
            scan_chunk_rows: 65_536,
            adhoc_group_cap: 65_536,
            key_spans_spill_bytes: 128 << 20,
        }
    }
}

impl StoreConfig {
    /// Every field from `<prefix>RAM_TIER_BYTES`, `<prefix>SCAN_CHUNK_ROWS`,
    /// `<prefix>ADHOC_GROUP_CAP` and `<prefix>KEY_SPANS_SPILL_BYTES`, each falling
    /// back to [`Default`].
    pub fn from_env(prefix: &str) -> Self {
        let d = StoreConfig::default();
        StoreConfig {
            ram_tier_bytes: env_num(&format!("{prefix}RAM_TIER_BYTES"), d.ram_tier_bytes),
            scan_chunk_rows: env_num(&format!("{prefix}SCAN_CHUNK_ROWS"), d.scan_chunk_rows),
            adhoc_group_cap: env_num(&format!("{prefix}ADHOC_GROUP_CAP"), d.adhoc_group_cap),
            key_spans_spill_bytes: env_num(&format!("{prefix}KEY_SPANS_SPILL_BYTES"), d.key_spans_spill_bytes),
        }
    }
}

/// Which relations this process holds, the stores behind them, and what this
/// process is to those stores.
///
/// Nothing here compiles a circuit or runs an epoch — those are
/// `gnitz-server`'s, and the crate graph is what says so.
///
/// **Thread contract.** `!Send + !Sync` by auto-trait: every store holds its
/// runs behind `Rc`. Nothing beneath a registry is thread-affine, so a host may
/// move one to another thread while it has exclusive access to it and no
/// `Rc`-bearing value it handed out is still alive outside it: those refcounts
/// are non-atomic.
pub struct RelationRegistry {
    pub(crate) tables: FxHashMap<u64, Relation>,
    /// The data directory every relation's [`relation_dir`] sits under.
    pub(crate) base_dir: String,
    /// Which worker this process is, of how many: names the `w{k}of{n}` child
    /// its stores open under.
    pub(crate) slot: Slot,
    pub(crate) config: StoreConfig,
    /// What this process is to the stores it registered.
    pub(crate) residency: Residency,
}

impl RelationRegistry {
    /// An empty [`Residency::Origin`] registry over the data directory `base_dir`.
    pub fn new(base_dir: &str, slot: Slot, config: StoreConfig) -> Self {
        RelationRegistry {
            tables: FxHashMap::default(),
            base_dir: base_dir.to_string(),
            slot,
            config: StoreConfig {
                scan_chunk_rows: config.scan_chunk_rows.max(1),
                ..config
            },
            residency: Residency::Origin,
        }
    }

    /// An empty [`Residency::Master`] registry over `base_dir` for `worker_count` workers.
    pub fn master(base_dir: &str, worker_count: u32, config: StoreConfig) -> Self {
        RelationRegistry {
            residency: Residency::Master,
            ..Self::new(base_dir, Slot::new(0, worker_count), config)
        }
    }

    // ── Relation registry ───────────────────────────────────────────────

    /// Drop `id`'s entry, and with it its owned stores and fds.
    pub fn unregister(&mut self, id: u64) {
        self.tables.remove(&id);
    }

    /// Enter a secondary index on `owner` over `cols`, filled from this process's
    /// slice of the owner, or add a claim to the one already on `cols`. On
    /// `Err` nothing is entered. A claim's `unique` is trusted: a duplicate can
    /// straddle two workers' slices, so only the caller can check it.
    pub fn add_index(&mut self, owner: u64, claim: IndexClaim, cols: PkColList) -> Result<(), String> {
        self.enter_index(owner, claim, cols, None)
    }

    /// [`Self::add_index`] beside an owner [reopened](Self::reopen_view) at
    /// `generation`: the index's store resumes from its manifest at that
    /// generation, and is filled from the owner when it carries none.
    pub fn reopen_index(
        &mut self,
        owner: u64,
        claim: IndexClaim,
        cols: PkColList,
        generation: u64,
    ) -> Result<(), String> {
        self.enter_index(owner, claim, cols, Some(generation))
    }

    fn enter_index(
        &mut self,
        owner: u64,
        claim: IndexClaim,
        cols: PkColList,
        resume_at: Option<u64>,
    ) -> Result<(), String> {
        let unique = matches!(claim, IndexClaim::Index { unique: true, .. });
        let owner_schema = self.index_owner(owner, unique)?.schema();
        let entry = self.tables.get_mut(&owner).expect("resolved above");
        if let Some(ix) = entry.indexes.iter_mut().find(|ix| ix.cols == cols) {
            debug_assert!(!ix.claims.contains(&claim), "{claim:?} claimed twice");
            ix.claims.push(claim);
            return Ok(());
        }
        let (key_spec, index_schema) = index_spec_and_schema(cols.as_slice(), &owner_schema)?;
        let mut ix = SecondaryIndex {
            cols,
            store: Store::Absent(Box::new(index_schema)),
            key_spec,
            claims: vec![claim],
            covers_pk: owner_schema.covers_pk(cols.as_slice()),
        };
        if self.residency.owns_stores() {
            ix.store = Store::Held(Box::new(self.open_child(
                owner,
                ChildKind::Index(ix.cols),
                index_schema,
                RecoverySource::Rederive { resume_at },
                self.store_budgets(),
            )?));
            if !ix.resumed() {
                let owner_store = &self.tables[&owner].store;
                ingest::fill_indexes(owner_store, self.config.scan_chunk_rows, owner, &mut [&mut ix])?;
            }
        }
        self.tables.get_mut(&owner).expect("resolved above").indexes.push(ix);
        Ok(())
    }

    /// Remove catalog index `index_id`'s claim from whichever of `owner`'s
    /// circuits holds it, dropping a circuit left with none. A no-op when the
    /// owner or claim is absent.
    pub fn release_index(&mut self, owner: u64, index_id: u64) {
        let Some(entry) = self.tables.get_mut(&owner) else {
            return;
        };
        for ix in &mut entry.indexes {
            ix.claims
                .retain(|c| !matches!(*c, IndexClaim::Index { id, .. } if id == index_id));
        }
        entry.indexes.retain(|ix| !ix.claims.is_empty());
    }

    /// Publish a new column schema for a registered base table in place (any
    /// column ALTER): update the registry copy and push the same value down into
    /// the owned `Table`, which rebinds its shards if the region count grew.
    pub fn swap_schema(&mut self, id: u64, schema: SchemaDescriptor) -> Result<(), String> {
        let entry = self.relation_mut_or_err(id)?;
        // Checked, not asserted: what a stale `key_spec` produces is a silently
        // wrong index projection, which release codegen would not guard at all.
        if !schema.is_trailing_append_of(entry.store.schema()) {
            return Err(format!(
                "ALTER on table {id}: the new descriptor is not a trailing append of the stored one"
            ));
        }
        entry
            .store
            .swap_schema(schema)
            .map_err(|e| format!("ALTER on table {id}: rebinding shards: {e}"))
    }

    // ── Registry reads ──────────────────────────────────────────────────

    pub fn has_id(&self, id: u64) -> bool {
        self.tables.contains_key(&id)
    }

    /// The data directory every relation's [`relation_dir`] sits under.
    pub fn base_dir(&self) -> &str {
        &self.base_dir
    }

    pub fn slot(&self) -> Slot {
        self.slot
    }

    /// This process's `kind` child directory of relation `id`.
    pub fn child_dir(&self, id: u64, kind: ChildKind<'_>) -> String {
        ChildAddr { kind, slot: self.slot }.dir(&relation_dir(&self.base_dir, id))
    }

    /// See [`StoreConfig::scan_chunk_rows`].
    pub fn scan_chunk_rows(&self) -> usize {
        self.config.scan_chunk_rows
    }

    /// Set the chunk size, clamped to at least one row. Per registry, not
    /// process-wide. Production sets it once through [`StoreConfig`].
    pub fn set_scan_chunk_rows(&mut self, rows: usize) {
        self.config.scan_chunk_rows = rows.max(1);
    }

    /// The relation `id` names, or `None` for an unknown id.
    pub fn relation(&self, id: u64) -> Option<&Relation> {
        self.tables.get(&id)
    }

    /// [`Self::relation`] plus the one "not registered" sentence — spelled the
    /// same by the mutating paths.
    pub fn relation_or_err(&self, id: u64) -> Result<&Relation, String> {
        self.relation(id).ok_or_else(|| Self::unregistered(id))
    }

    /// `owner`, if it may carry a secondary index, `unique` or not. Every path
    /// that enters an index asks here.
    pub fn index_owner(&self, owner: u64, unique: bool) -> Result<&Relation, String> {
        let e = self.relation_or_err(owner)?;
        if !e.kind.admits_index() {
            return Err(format!(
                "Index: owner {owner} is a {}; only a base table or a view without a capacity can be indexed",
                e.kind.noun()
            ));
        }
        if unique && e.kind.is_view() {
            return Err(format!(
                "Index: view {owner} cannot carry a UNIQUE index: its circuit cannot refuse a duplicate"
            ));
        }
        // An index walk gathers every row under an in-range entry's source PK, so
        // a second row under that PK would come back whatever its own key.
        if e.pk_repeats {
            return Err(format!(
                "Index: owner {owner} repeats its primary key, so an index entry does not name one row"
            ));
        }
        Ok(e)
    }

    /// [`Self::relation_or_err`] as `&mut`.
    pub(crate) fn relation_mut_or_err(&mut self, id: u64) -> Result<&mut Relation, String> {
        self.tables.get_mut(&id).ok_or_else(|| Self::unregistered(id))
    }

    fn unregistered(id: u64) -> String {
        format!("relation {id} is not registered")
    }

    /// True iff at least one registered relation carries a delta feed.
    pub fn any_delta_feed(&self) -> bool {
        self.tables.values().any(|r| r.kind.has_delta_feed())
    }

    /// Every registered view id.
    pub fn view_ids(&self) -> impl Iterator<Item = u64> + '_ {
        self.tables.iter().filter(|(_, e)| e.kind.is_view()).map(|(&id, _)| id)
    }

    /// What every store this registry opens starts from; the two bounded kinds
    /// narrow it.
    pub(crate) fn store_budgets(&self) -> StoreBudgets {
        StoreBudgets::new(self.config.ram_tier_bytes)
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
#[path = "tests/relation.rs"]
mod tests;