yeti_types/backend/

manager.rs

//! Per-table backend registry: `TableInfo`, `BackendManager`, and the
//! `unix_timestamp` / `now_secs` helpers.
//!
//! # Errors
//!
//! Result-returning methods on `BackendManager` return
//! `Result<T, YetiError>` and surface these failure classes:
//!
//! - **Table not registered** — `get_backend_for_table` /
//!   `get_table_info` return `YetiError::Internal` when the named
//!   table has no entry in the registry. Indicates an app-config or
//!   schema mismatch.
//! - **Backend init / open** — `RocksDB` column-family open failure
//!   (corrupt LSM, permission denied, locked DB). Returned as
//!   `YetiError::Internal`.
//! - **Metadata I/O** — encode/decode of stored `TableInfo` records
//!   (under `__yeti_table_backend:` prefix). `YetiError::Internal`
//!   from msgpack or `RocksDB`.
//! - **Clock skew** — `unix_timestamp` returns
//!   `YetiError::Internal("system clock before UNIX epoch")` if the
//!   host clock is set before 1970-01-01.
#![allow(clippy::missing_errors_doc)]

use lru::LruCache;
use std::collections::HashMap;
use std::num::NonZeroUsize;
use std::sync::{Arc, RwLock};

use super::config::BackendType;
use super::config::ConsistencyMode;
use super::traits::KvBackend;
use crate::error::{Result, ResultExt, YetiError};

/// Metadata key prefix for storing backend information in the database.
const METADATA_PREFIX: &str = "__yeti_table_backend:";

/// Default LRU cache capacity for table backend mappings.
const DEFAULT_CACHE_CAPACITY: usize = 10_000;

/// Compile-time `NonZeroUsize::new(1)`. Constructed via const-match so the
/// non-zero invariant is proven at compile time — no runtime `expect`.
const ONE_NONZERO: NonZeroUsize = match NonZeroUsize::new(1) {
    Some(n) => n,
    // SAFETY: the constant `1` is non-zero. Unreachable.
    None => unreachable!(),
};

/// Compile-time `NonZeroUsize::new(DEFAULT_CACHE_CAPACITY)`.
const DEFAULT_CACHE_CAPACITY_NONZERO: NonZeroUsize = match NonZeroUsize::new(DEFAULT_CACHE_CAPACITY)
{
    Some(n) => n,
    // SAFETY: `DEFAULT_CACHE_CAPACITY` is the literal `10_000`. Unreachable.
    None => unreachable!(),
};

// ============================================================================
// TableInfo
// ============================================================================

/// Table information for backend initialization.
#[derive(Debug, Clone)]
pub struct TableInfo {
    /// Owning application id. Stamped onto every transaction-log
    /// entry's `table.app` field so audit / replication observers
    /// can attribute writes to the right app even when multiple apps
    /// share a database (`@table(database: "...")` declares cross-
    /// app sharing). Empty string is the default for test fixtures
    /// that don't care about app attribution; in production every
    /// `TableInfo` is built from an `AppMetadata` with a populated
    /// `id`.
    pub app_id: String,
    /// Table name (used as column family name)
    pub name: String,
    /// Database name (default: "data")
    pub database: String,
    /// Storage backend type
    pub storage: BackendType,
    /// Per-table consistency mode override
    pub consistency: Option<ConsistencyMode>,
    /// Per-field CRDT type declarations
    pub crdt_fields: HashMap<String, String>,
    /// Residency mode
    pub residency: Option<String>,
    /// Per-record TTL expiration from `@table(expiration: N)` — `Some(secs)`
    /// activates the 9-byte record header + `RocksDB` `ExpirationFilter`.
    /// `None` means the column family is written without a header.
    pub expiration_secs: Option<u64>,
    /// True when the schema declared `@audit(capture_state: true)`
    /// on this table. The `LoggingBackend` wrapper uses this flag to
    /// decide whether to read the existing value before each write
    /// and stamp it onto `LogEntry::prev_value` for `AuditEntry.before`.
    /// Costs one extra storage read per write — hence per-table
    /// opt-in.
    pub audit_capture_state: bool,
    /// Per-table durability from `@store(durability:)`. `None` inherits the
    /// database-level default (the WAL setting the `BackendManager` was
    /// opened with). `Some(tier)` overrides this table's `WriteOptions`:
    /// `Lossy` skips the WAL (throughput), `Soft`/`Strong` write through it,
    /// `Strong` additionally fsyncs each write. `RocksDB` applies `disableWAL`
    /// and `sync` per write, so column families in one database can differ.
    pub durability: Option<crate::schema::DurabilityTier>,
    /// The full source [`TableDefinition`](crate::schema::TableDefinition)
    /// this `TableInfo` was projected from. Carried so downstream consumers
    /// (replication's `@distribute`, the router's `@export`/transport flags)
    /// read one source of truth instead of side-channel maps or re-parsing
    /// the schema. `None` for hand-built test fixtures, which only exercise
    /// the flattened backend fields.
    pub def: Option<std::sync::Arc<crate::schema::TableDefinition>>,
}

impl Default for TableInfo {
    fn default() -> Self {
        Self {
            app_id: String::new(),
            name: String::new(),
            database: "data".to_owned(),
            storage: BackendType::Disk,
            consistency: None,
            crdt_fields: HashMap::new(),
            residency: None,
            expiration_secs: None,
            audit_capture_state: false,
            durability: None,
            def: None,
        }
    }
}

// ============================================================================
// BackendManager
// ============================================================================

/// Manages storage backends with per-table backend selection.
///
/// Each table gets its own backend instance. Multiple applications safely share
/// databases. Uses LRU cache for bounded memory on backend type validation.
pub struct BackendManager {
    /// Map of table name to backend instance
    table_backends: HashMap<String, Arc<dyn KvBackend>>,
    /// Pre-built lowercase → canonical name index for case-insensitive lookup.
    /// Populated at construction time so `get_backend_for_table()` avoids
    /// linear scans with per-call `to_lowercase()` allocations.
    lowercase_index: HashMap<String, String>,
    /// LRU cache tracking which backend type each table uses
    backend_type_cache: RwLock<LruCache<String, BackendType>>,
    /// Metadata key prefix
    metadata_prefix: String,
    /// `PubSub` manager for real-time table change notifications.
    /// Set lazily after construction via `set_pubsub()` — avoids cloning
    /// the `table_backends` `HashMap` just to attach `PubSub`.
    pubsub: std::sync::OnceLock<Arc<crate::pubsub::PubSubManager>>,
    /// Per-table source `TableDefinition`s (by table name). Set once at
    /// startup; the single source of truth for table metadata that the
    /// flattened backend fields don't carry (`@distribute`, `@export`, …).
    table_defs: std::sync::OnceLock<
        Arc<std::collections::HashMap<String, Arc<crate::schema::TableDefinition>>>,
    >,
    /// Per-app vector capabilities (hook, batcher, mappings, cache).
    /// Set once at startup after vector hooks are configured.
    vector: std::sync::OnceLock<Arc<crate::plugins::VectorContext>>,
    /// Per-table secondary-index registry. Populated at app load with one
    /// `Arc<IndexManager>` per table that has any `@indexed` field. Stored
    /// type-erased here because `IndexManager` lives in `yeti-sdk`
    /// (which depends on `yeti-types`); the SDK provides the typed accessor
    /// that downcasts this slot. SDK-side `Table::put_if` consults the
    /// registry to keep hash/range/fulltext indexes in sync with CAS writes
    /// — without this, the `@indexed` directive is a no-op for SDK writes
    /// and worker pools that scan by status pay a full-table scan.
    index_registry: std::sync::OnceLock<Arc<dyn std::any::Any + Send + Sync>>,
    /// Unified transaction log — set at startup by yeti-server when
    /// the deployment opens its transaction-log backend at
    /// `{root_directory}/logs/transactions/{deployment_hash}/`.
    /// Observed by both audit (range-by-time queries) and replication
    /// (tail-by-HLC). `None` when logging is disabled (test fixtures,
    /// dry-run modes).
    transaction_log: std::sync::OnceLock<crate::backend::TransactionLogHandle>,
    /// Shared [`ObserverRegistry`] — every per-table `LoggingBackend`
    /// holds a clone of the same `Arc`, so a single
    /// [`register_observer`](Self::register_observer) call reaches
    /// every commit path. `None` when observer fan-out is disabled.
    ///
    /// [`ObserverRegistry`]: crate::backend::ObserverRegistry
    observer_registry: std::sync::OnceLock<Arc<crate::backend::ObserverRegistry>>,
}

impl std::fmt::Debug for BackendManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BackendManager")
            .field(
                "table_backends",
                &self.table_backends.keys().collect::<Vec<_>>(),
            )
            .field("metadata_prefix", &self.metadata_prefix)
            .field("has_pubsub", &self.pubsub.get().is_some())
            .field("has_vector", &self.vector.get().is_some())
            .field("has_transaction_log", &self.transaction_log.get().is_some())
            .field(
                "has_observer_registry",
                &self.observer_registry.get().is_some(),
            )
            // `lowercase_index`, `backend_type_cache`, dyn-trait values elided
            .finish_non_exhaustive()
    }
}

impl BackendManager {
    /// Create an empty `BackendManager` for testing.
    #[must_use]
    pub fn empty() -> Self {
        Self {
            table_backends: HashMap::new(),
            lowercase_index: HashMap::new(),
            backend_type_cache: RwLock::new(LruCache::new(ONE_NONZERO)),
            metadata_prefix: METADATA_PREFIX.to_owned(),
            pubsub: std::sync::OnceLock::new(),
            vector: std::sync::OnceLock::new(),
            index_registry: std::sync::OnceLock::new(),
            transaction_log: std::sync::OnceLock::new(),
            observer_registry: std::sync::OnceLock::new(),
            table_defs: std::sync::OnceLock::new(),
        }
    }

    /// Create from pre-built backends.
    pub fn from_backends(table_backends: HashMap<String, Arc<dyn KvBackend>>) -> Result<Self> {
        let cache_capacity = NonZeroUsize::new(DEFAULT_CACHE_CAPACITY).ok_or_else(|| {
            YetiError::Internal("DEFAULT_CACHE_CAPACITY must be non-zero".to_owned())
        })?;

        let lowercase_index: HashMap<String, String> = table_backends
            .keys()
            .map(|k| (k.to_lowercase(), k.clone()))
            .collect();

        Ok(Self {
            table_backends,
            lowercase_index,
            backend_type_cache: RwLock::new(LruCache::new(cache_capacity)),
            metadata_prefix: METADATA_PREFIX.to_owned(),
            pubsub: std::sync::OnceLock::new(),
            vector: std::sync::OnceLock::new(),
            index_registry: std::sync::OnceLock::new(),
            transaction_log: std::sync::OnceLock::new(),
            observer_registry: std::sync::OnceLock::new(),
            table_defs: std::sync::OnceLock::new(),
        })
    }

    /// Set the unified transaction log. Can only be called once
    /// (uses `OnceLock`). Called at startup once the deployment's
    /// transaction-log backend
    /// (`{root_directory}/logs/transactions/{deployment_hash}/`)
    /// has been opened and a [`TransactionLog`] impl built on top.
    ///
    /// [`TransactionLog`]: crate::backend::TransactionLog
    pub fn set_transaction_log(&self, log: crate::backend::TransactionLogHandle) {
        let _ = self.transaction_log.set(log);
    }

    /// Get the unified transaction log, if set. Audit readers,
    /// replication senders, and any future log consumer use this to
    /// reach the durable record of every committed write.
    #[must_use]
    pub fn transaction_log(&self) -> Option<&crate::backend::TransactionLogHandle> {
        self.transaction_log.get()
    }

    /// Set the shared [`ObserverRegistry`]. Can only be called once
    /// (uses `OnceLock`). Called at startup once the same `Arc` has
    /// been threaded into every per-table `LoggingBackend` via
    /// `LogWrapping.observers`, so this is the manager-side handle
    /// to the same instance.
    ///
    /// [`ObserverRegistry`]: crate::backend::ObserverRegistry
    pub fn set_observer_registry(&self, registry: Arc<crate::backend::ObserverRegistry>) {
        let _ = self.observer_registry.set(registry);
    }

    /// Get the shared observer registry, if set. Plugins (replication,
    /// audit, etc.) register their [`TransactionObserver`] here during
    /// startup; every per-table `LoggingBackend` reads the same
    /// snapshot on its commit path.
    ///
    /// [`TransactionObserver`]: crate::backend::TransactionObserver
    #[must_use]
    pub fn observer_registry(&self) -> Option<&Arc<crate::backend::ObserverRegistry>> {
        self.observer_registry.get()
    }

    /// Convenience: register an observer on the shared registry. No-op
    /// when no registry has been attached.
    pub fn register_observer(&self, observer: crate::backend::ObserverHandle) {
        if let Some(registry) = self.observer_registry.get() {
            registry.register(observer);
        }
    }

    /// Set the `PubSub` manager. Can only be called once (uses `OnceLock`).
    /// Called at startup after `BackendManager` is created and wrapped in Arc.
    pub fn set_pubsub(&self, pubsub: Arc<crate::pubsub::PubSubManager>) {
        let _ = self.pubsub.set(pubsub);
    }

    /// Get the `PubSub` manager, if set.
    pub fn pubsub(&self) -> Option<&Arc<crate::pubsub::PubSubManager>> {
        self.pubsub.get()
    }

    /// Attach the per-table source [`TableDefinition`](crate::schema::TableDefinition)s
    /// (keyed by table name). Set once at construction so downstream
    /// consumers — replication's `@distribute` filter, the router's
    /// `@export` flags — read this single source of truth rather than a
    /// side-channel map. Can only be set once (uses `OnceLock`).
    pub fn set_table_defs(
        &self,
        defs: std::collections::HashMap<String, Arc<crate::schema::TableDefinition>>,
    ) {
        let _ = self.table_defs.set(Arc::new(defs));
    }

    /// The source `TableDefinition` for `table_name`, if known. `None` when
    /// defs weren't attached (test fixtures) or the table is unknown.
    #[must_use]
    pub fn table_def(&self, table_name: &str) -> Option<Arc<crate::schema::TableDefinition>> {
        self.table_defs.get()?.get(table_name).map(Arc::clone)
    }

    /// Set the per-app vector context. Can only be called once (uses `OnceLock`).
    /// Called at startup after vector hooks and batcher are configured.
    pub fn set_vector(&self, vector: Arc<crate::plugins::VectorContext>) {
        let _ = self.vector.set(vector);
    }

    /// Get the per-app vector context, if set.
    pub fn vector(&self) -> Option<&Arc<crate::plugins::VectorContext>> {
        self.vector.get()
    }

    /// Create a new manager with this manager's tables plus additional tables.
    #[must_use]
    pub fn with_merged_tables(&self, other: &Self) -> Self {
        let mut table_backends = self.table_backends.clone();
        for (name, backend) in &other.table_backends {
            table_backends
                .entry(name.clone())
                .or_insert_with(|| Arc::clone(backend));
        }
        let lowercase_index: HashMap<String, String> = table_backends
            .keys()
            .map(|k| (k.to_lowercase(), k.clone()))
            .collect();
        Self {
            table_backends,
            lowercase_index,
            backend_type_cache: RwLock::new(LruCache::new(DEFAULT_CACHE_CAPACITY_NONZERO)),
            metadata_prefix: self.metadata_prefix.clone(),
            pubsub: self
                .pubsub
                .get()
                .cloned()
                .map_or_else(std::sync::OnceLock::new, |ps| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(ps);
                    lock
                }),
            vector: self
                .vector
                .get()
                .cloned()
                .map_or_else(std::sync::OnceLock::new, |vc| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(vc);
                    lock
                }),
            index_registry: self.index_registry.get().cloned().map_or_else(
                std::sync::OnceLock::new,
                |r| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(r);
                    lock
                },
            ),
            transaction_log: self.transaction_log.get().cloned().map_or_else(
                std::sync::OnceLock::new,
                |l| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(l);
                    lock
                },
            ),
            observer_registry: self.observer_registry.get().cloned().map_or_else(
                std::sync::OnceLock::new,
                |r| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(r);
                    lock
                },
            ),
            table_defs: self
                .table_defs
                .get()
                .cloned()
                .map_or_else(std::sync::OnceLock::new, |d| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(d);
                    lock
                }),
        }
    }

    /// Merge another manager's tables with database-qualified keys.
    #[must_use]
    pub fn with_qualified_merge(&self, other: &Self, database: &str) -> Self {
        let mut table_backends = self.table_backends.clone();
        for (name, backend) in &other.table_backends {
            let qualified = format!("{database}.{name}");
            table_backends
                .entry(qualified)
                .or_insert_with(|| Arc::clone(backend));
            table_backends
                .entry(name.clone())
                .or_insert_with(|| Arc::clone(backend));
        }
        let lowercase_index: HashMap<String, String> = table_backends
            .keys()
            .map(|k| (k.to_lowercase(), k.clone()))
            .collect();
        Self {
            table_backends,
            lowercase_index,
            backend_type_cache: RwLock::new(LruCache::new(DEFAULT_CACHE_CAPACITY_NONZERO)),
            metadata_prefix: self.metadata_prefix.clone(),
            pubsub: self
                .pubsub
                .get()
                .cloned()
                .map_or_else(std::sync::OnceLock::new, |ps| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(ps);
                    lock
                }),
            vector: self
                .vector
                .get()
                .cloned()
                .map_or_else(std::sync::OnceLock::new, |vc| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(vc);
                    lock
                }),
            index_registry: self.index_registry.get().cloned().map_or_else(
                std::sync::OnceLock::new,
                |r| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(r);
                    lock
                },
            ),
            transaction_log: self.transaction_log.get().cloned().map_or_else(
                std::sync::OnceLock::new,
                |l| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(l);
                    lock
                },
            ),
            observer_registry: self.observer_registry.get().cloned().map_or_else(
                std::sync::OnceLock::new,
                |r| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(r);
                    lock
                },
            ),
            table_defs: self
                .table_defs
                .get()
                .cloned()
                .map_or_else(std::sync::OnceLock::new, |d| {
                    let lock = std::sync::OnceLock::new();
                    let _ = lock.set(d);
                    lock
                }),
        }
    }

    /// Set the per-app secondary-index registry. Type-erased so this crate
    /// stays free of yeti-sdk imports — yeti-sdk's `Tables::get` does the
    /// downcast. Can only be called once (uses `OnceLock`). Called at
    /// startup once `IndexManager`s are built from the loaded schema.
    pub fn set_index_registry(&self, registry: Arc<dyn std::any::Any + Send + Sync>) {
        let _ = self.index_registry.set(registry);
    }

    /// Get the per-app secondary-index registry, if set. Callers (yeti-sdk)
    /// downcast to their concrete `HashMap<String, Arc<IndexManager>>`.
    pub fn index_registry(&self) -> Option<&Arc<dyn std::any::Any + Send + Sync>> {
        self.index_registry.get()
    }

    /// List all table names managed by this `BackendManager`.
    pub fn table_names(&self) -> Vec<String> {
        self.table_backends.keys().cloned().collect()
    }

    /// Get the storage backend for a specific table by name.
    pub fn get_backend_for_table(&self, table_name: &str) -> Result<Arc<dyn KvBackend>> {
        // Exact match first (no allocation)
        if let Some(backend) = self.table_backends.get(table_name) {
            return Ok(Arc::clone(backend));
        }
        // Case-insensitive via pre-built index (single to_lowercase, no linear scan)
        let lower = table_name.to_lowercase();
        if let Some(canonical) = self.lowercase_index.get(&lower)
            && let Some(backend) = self.table_backends.get(canonical)
        {
            return Ok(Arc::clone(backend));
        }
        Err(YetiError::NotFound {
            resource_type: "backend".to_owned(),
            id: table_name.to_owned(),
        })
    }

    /// Probe storage health by performing lightweight reads across all backends.
    pub async fn health_probe(&self) -> Result<()> {
        if self.table_backends.is_empty() {
            return Err(YetiError::Internal(
                "No table backends registered".to_owned(),
            ));
        }
        for backend in self.table_backends.values() {
            backend.get(b"__health_probe__").await?;
        }
        Ok(())
    }

    /// Validate that a table is using the correct backend.
    pub async fn validate_table_backend(
        &self,
        table_name: &str,
        key_prefix: &str,
        expected_backend: BackendType,
    ) -> Result<()> {
        let metadata_key = format!("{}{}:{}", self.metadata_prefix, key_prefix, table_name);

        {
            let mut cache = self
                .backend_type_cache
                .write()
                .with_context(|| "table backends cache lock".to_owned())?;
            if let Some(&existing_backend) = cache.get(table_name) {
                if existing_backend != expected_backend {
                    return Err(YetiError::Validation(format!(
                        "Table '{table_name}' already exists with {existing_backend} backend but schema specifies {expected_backend} backend. \
                         Storage backend migration is not yet supported."
                    )));
                }
                return Ok(());
            }
        }

        let backend = self.get_backend_for_table(table_name)?;

        if let Some(stored_value) = backend.get(metadata_key.as_bytes()).await? {
            let stored_backend_str = String::from_utf8_lossy(&stored_value);
            let stored_backend: BackendType = stored_backend_str.parse()?;

            if stored_backend != expected_backend {
                return Err(YetiError::Validation(format!(
                    "Table '{table_name}' has stored backend metadata indicating {stored_backend}, \
                     but schema specifies {expected_backend}. Storage backend migration is not yet supported."
                )));
            }

            self.backend_type_cache
                .write()
                .with_context(|| "table backends cache lock".to_owned())?
                .put(table_name.to_owned(), stored_backend);
            return Ok(());
        }

        backend
            .put(
                metadata_key.as_bytes(),
                expected_backend.to_string().as_bytes(),
            )
            .await?;

        self.backend_type_cache
            .write()
            .with_context(|| "table backends cache lock".to_owned())?
            .put(table_name.to_owned(), expected_backend);

        tracing::debug!(
            "  Initialized new table '{}' with {} backend",
            table_name,
            expected_backend
        );

        Ok(())
    }
}

/// Get current Unix timestamp in seconds.
pub fn unix_timestamp() -> Result<u64> {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .map_err(|_| {
            YetiError::Internal("System time is before UNIX epoch - check system clock".to_owned())
        })
}

/// Get current Unix timestamp in seconds with safe fallback.
#[inline]
#[must_use]
pub fn now_secs() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_or(0, |d| d.as_secs())
}