rindle_replica/

derivation_pool.rs

//! [`DerivationPool`] — the worker-pool half of [`Cluster`](crate::Cluster) for a host
//! that owns its own writer connection and change capture (design 308 §5).
//!
//! `Cluster` bundles the whole write plane: it owns the single writer connection, runs
//! the preupdate capture on it, and coordinates snapshot barriers with writes and
//! [`ClusterWriteTxn::commit`](crate::ClusterWriteTxn::commit). An embedding like
//! `rindle-sqlite-ext` inverts that: the **host application** owns the writer connection
//! and the transaction lifecycle, capture happens on the host's connection, and the only
//! coordination point is a `sqlite3_commit_hook` that fires *before* the host's COMMIT is
//! durable or visible. This facade exposes exactly the machinery such a host needs — the
//! sharded worker pool and the begin-barrier → push → finish → gate handshake — without
//! a second writer connection or a second capture hook.
//!
//! ## The handshake, from a commit hook
//!
//! 1. The host's preupdate capture buffers the transaction's row changes.
//! 2. In the commit hook (pre-durability, pre-visibility): [`begin`](DerivationPool::begin)
//!    — every worker opens + pins its own read snapshot, which sees **T-1** because the
//!    writer's changes are not yet visible — then [`PoolTxn::push`] the captured batch
//!    (bounded chunks, one shared `Arc` per chunk) and [`PoolTxn::finish`] for the
//!    [`PoolGate`].
//! 3. Release the gate — [`commit`](PoolGate::commit) once the host's COMMIT will proceed,
//!    [`abort`](PoolGate::abort) if it will not — and return from the hook. The workers
//!    derive against their pinned T-1 snapshots *concurrently* with the host's commit
//!    fsync. The barrier and bounded capture fan-out can block the writer.
//!
//! The pool trusts the supplied verdict; it cannot observe the host's SQL commit.
//! To give `Progressed` the same durable-commit meaning as `Cluster`, release the gate
//! after a successful COMMIT. A host that releases it inside a pre-commit hook must
//! handle later commit failure itself, invalidate affected subscriptions, and prevent
//! consumers from treating provisional host outcomes as durable confirmation.
//!
//! Derived events flow out the bounded channel returned by [`spawn`](DerivationPool::spawn)
//! as [`ClusterEvent`]s, exactly as under `Cluster`: `Update::Hydrated` at registration,
//! speculative `Update::Changed` per pushed chunk, the per-worker `Progressed` commit
//! marker (the release signal — `Changed(N)` always precedes `Progressed(N)` per worker),
//! and terminal `Faulted`. The same draining contract applies: **the consumer must drain
//! the channel continuously** — a full channel blocks a worker mid-commit, and the next
//! barrier then times out.
//!
//! ## What the caller owns
//!
//! - **Transaction ids**: a monotonic `tx_id` per pushed transaction. `Cluster` persists
//!   its watermark inside the writer transaction; a hook-driven host cannot write from
//!   inside the hook, so its ids are typically in-memory (its consumers recover by
//!   re-subscribe + fresh hydrate, not by replaying a durable cursor).
//! - **The single-writer contract**: exactly one connection writes, and every `begin`
//!   happens while that writer's transaction is still open (the T-1 window). Nothing
//!   here can check that; it is the same contract `Cluster` enforces by construction.
//! - **Barrier-failure recovery**: `begin` returning `None` means at least one worker
//!   missed the barrier and was aborted/reaped — but unlike `Cluster`, the host usually
//!   cannot veto its own commit, so the transaction still lands and **every** worker's
//!   pipelines are now stale. The caller must treat all registered queries as faulted
//!   (deregister + re-register to re-hydrate at head).

use std::path::Path;
use std::sync::mpsc::Receiver;

use rindle::Ast;
use rindle_cdc::Captured;

use crate::parallel::{
    event_channel_bound, CommitGate, StreamBegin, StreamTx, WorkerPool, PUSH_CHUNK_ROWS,
};
use crate::schema::TableSchema;
use crate::{ChangeEvent, ClusterEvent, OpenOptions, QueryId, ReplicaError};

/// The worker-pool seam for an external writer. See the module docs. Single-thread-use,
/// like [`Cluster`](crate::Cluster): the pool handle lives on the writer's thread (the
/// hooks fire there), and only `Send` data crosses to the workers.
pub struct DerivationPool {
    pool: WorkerPool,
    n_workers: usize,
}

impl DerivationPool {
    /// Spawn `n_workers` IVM worker threads over the **file-backed** database at `path`,
    /// using the planner, operator-storage, journal, and foreign-key options. The
    /// host owns its writer's `wal_autocheckpoint`; the pool does not configure it.
    /// Returns the pool
    /// and the bounded channel-out of derived [`ClusterEvent`]s (`max(256, n_workers · 16)`
    /// events — FOLLOWER-LAG-SHED §4, rung 0b; see the module docs' draining contract). Asserts
    /// `sqlite3_threadsafe() != 0`.
    ///
    /// [`Cluster::open_with`]: crate::Cluster::open_with
    pub fn spawn(
        path: impl AsRef<Path>,
        n_workers: usize,
        opts: OpenOptions,
    ) -> Result<(DerivationPool, Receiver<ClusterEvent>), ReplicaError> {
        assert!(n_workers > 0, "a derivation pool needs at least one worker");
        if unsafe { rusqlite::ffi::sqlite3_threadsafe() } == 0 {
            return Err(ReplicaError::NotThreadsafe);
        }
        let path = path.as_ref();
        let (out_tx, out_rx) = std::sync::mpsc::sync_channel(event_channel_bound(n_workers));
        let pool = WorkerPool::spawn(
            path,
            n_workers,
            opts.plan_queries,
            opts.operator_storage,
            opts.journal,
            opts.foreign_keys,
            out_tx,
        )?;
        Ok((DerivationPool { pool, n_workers }, out_rx))
    }

    /// The number of worker slots (fixed at spawn; the `query_id % n` shard count).
    pub fn n_workers(&self) -> usize {
        self.n_workers
    }

    /// The worker slot that hosts (or would host) `query_id`. The pool owns the shard
    /// mapping; callers key per-query progress by this index rather than recomputing it.
    pub fn worker_of(&self, query_id: QueryId) -> usize {
        (query_id.0 % self.n_workers as u64) as usize
    }

    /// Build the table's source on **every** worker (any may host a query referencing
    /// it). Idempotence and the PK-UNIQUE-index prerequisite are the caller's concern —
    /// the discovery/ensure half lives with whoever owns a connection (the host), e.g.
    /// [`discover_table_schema`](crate::discover_table_schema) +
    /// [`ensure_unique_pk_index`](crate::ensure_unique_pk_index). Must not be called
    /// while the writer's transaction is open: the workers build sources against
    /// committed state.
    pub fn register_table(&self, table: &str, schema: TableSchema) -> Result<(), ReplicaError> {
        self.pool.register_table_all(table, schema)
    }

    /// Register + hydrate a live query on its shard ([`worker_of`](Self::worker_of));
    /// blocks until the worker built + hydrated it (a `BuildError` surfaces here), then
    /// the worker emits the `Hydrated` baseline — stamped `hydrated_tx`, the caller's
    /// committed watermark — to the channel-out. Returns the hosting worker's index.
    pub fn register_query(
        &self,
        query_id: QueryId,
        ast: Ast,
        hydrated_tx: u64,
    ) -> Result<usize, ReplicaError> {
        let worker = self.worker_of(query_id);
        self.pool
            .register_query_on(worker, query_id, ast, hydrated_tx)?;
        Ok(worker)
    }

    /// Gracefully tear down a query (solicited — no `Faulted` event; after this returns
    /// no further events arrive for `query_id`). Returns `true` if a live query was
    /// found and removed; idempotent.
    pub fn deregister_query(&self, query_id: QueryId) -> bool {
        self.pool
            .deregister_query_on(self.worker_of(query_id), query_id)
            > 0
    }

    /// Re-read the current assembled view of `query_id` from its hosting worker as
    /// hydration `Add`s (the SSR one-shot). FIFO with commits on that worker, so the
    /// returned view reflects every transaction whose gate was released before this
    /// call. A degraded shard or unregistered query reads as an empty snapshot; a raised
    /// read boundary on a live worker is `Err`. Must not be called between
    /// [`begin`](Self::begin) and the gate release.
    pub fn read_snapshot(&self, query_id: QueryId) -> Result<Vec<ChangeEvent>, ReplicaError> {
        self.pool
            .read_snapshot_on(self.worker_of(query_id), query_id)
    }

    /// The begin-barrier for one committing transaction: every live worker opens + pins
    /// its read snapshot at the pre-commit state and acks. `Some` ⇒ all pinned — stream
    /// the capture with [`PoolTxn::push`] and finish for the gate. `None` ⇒ a worker
    /// died or stalled at the barrier: the workers that did pin were rolled back cleanly
    /// (nothing was pushed) and the offender was reaped/respawned (its lost queries got
    /// terminal `Faulted` events) — but every OTHER worker now misses this transaction
    /// too, so if the host's commit proceeds anyway, the caller must fault + re-register
    /// every remaining query (see the module docs).
    pub fn begin(&self, tx_id: u64) -> Option<PoolTxn> {
        // `tx_begin` is a send; `await_acks` is the fence. A hook-driven host has no
        // writer work to overlap them with (it is already at its commit edge), so this
        // facade keeps the blocking begin contract and collects the acks immediately.
        match self.pool.tx_begin(tx_id) {
            StreamBegin::Ready(mut stream) => {
                if stream.await_acks() {
                    return Some(PoolTxn { stream });
                }
                // A worker stalled past the barrier watchdog: nothing pushed yet → a
                // clean snapshot rollback of the workers that pinned, then detect +
                // respawn the offender(s).
                stream.finish().abort();
                self.pool.reap();
                None
            }
            StreamBegin::Failed(stream) => {
                // A worker was dead at send. Same recovery: nothing pushed yet → a clean
                // snapshot rollback of any worker that pinned, then detect + respawn.
                stream.finish().abort();
                self.pool.reap();
                None
            }
        }
    }

    /// Liveness sweep: block until every worker has drained its command queue,
    /// respawning any that died or hung (their lost queries get terminal `Faulted`
    /// events). The quiesce point — a returned sweep implies all prior registrations and
    /// released gates have emitted their events to the channel-out. Mirrors
    /// [`Cluster::sync`](crate::Cluster::sync).
    pub fn reap(&self) {
        self.pool.reap();
    }

    /// Set the per-push derive deadline (FOLLOWER-LAG-SHED §6.6), in ms; `0` disables
    /// it. Applies immediately to every worker, respawns included. Default 10 s.
    pub fn set_push_deadline_ms(&self, ms: u64) {
        self.pool.set_push_deadline_ms(ms);
    }

    /// Override every worker engine's design-306 D4 delta-byte budget (see
    /// [`Db::set_max_delta_bytes`](crate::Db::set_max_delta_bytes)). Workers re-read it
    /// at each commit barrier, so it lands on the next transaction, respawns included.
    /// A worker whose batch delta overflows the budget sheds itself: tear down +
    /// terminal `Faulted` per hosted query.
    pub fn set_max_delta_bytes(&self, bytes: usize) {
        self.pool.set_max_delta_bytes_all(bytes);
    }

    /// Test seam: lower the barrier-ack + liveness-ping timeouts so death/stall
    /// recovery paths run fast. Doc-hidden; not part of the supported API.
    #[doc(hidden)]
    pub fn __test_set_timeouts(&self, ack: std::time::Duration, ping: std::time::Duration) {
        self.pool.set_timeouts(ack, ping);
    }

    /// Test seam: force a derivation-fault recovery on the worker hosting `on_query`
    /// (tears down every query on that worker, emitting terminal `Faulted` for each).
    /// Doc-hidden; not part of the supported API.
    #[doc(hidden)]
    pub fn __test_trigger_fault(&self, on_query: QueryId) {
        self.pool.trigger_fault(self.worker_of(on_query));
    }

    /// Test seam: kill the worker thread hosting `on_query` (a true thread death,
    /// detected + recovered by the next [`reap`](Self::reap)/barrier). Doc-hidden; not
    /// part of the supported API.
    #[doc(hidden)]
    pub fn __test_kill_worker(&self, on_query: QueryId) {
        self.pool.kill_worker(self.worker_of(on_query));
    }
}

/// An in-flight transaction across the pool: every live worker holds a pinned pre-commit
/// snapshot, awaiting pushed chunks and the terminal gate. Dropping it without
/// [`finish`](Self::finish) releases the workers' snapshots as an abort.
#[must_use = "a PoolTxn must be finished (or dropped) to release the workers' held snapshots"]
pub struct PoolTxn {
    stream: StreamTx,
}

impl PoolTxn {
    /// The transaction id this barrier was opened for.
    pub fn tx_id(&self) -> u64 {
        self.stream.tx_id()
    }

    /// Fan captured changes to every worker in bounded chunks (each one shared `Arc` —
    /// a refcount bump per worker, never a row copy). Each worker derives the chunk
    /// against its held snapshot and speculatively emits the deltas, tagged this
    /// transaction; the consumer holds them until the worker's `Progressed` marker. A
    /// worker saturated past its in-flight credit for too long is dropped from the
    /// stream (the liveness sweep owns it from there — its queries fault).
    pub fn push(&mut self, changes: &[Captured]) {
        for chunk in changes.chunks(PUSH_CHUNK_ROWS) {
            self.stream.push(chunk.to_vec().into());
        }
    }

    /// Close the push stream and take the commit gate — the **unsent terminal command**.
    /// Each worker keeps its snapshot until it receives `TxFinish` with the verdict.
    /// Bounded delivery can still block a worker. Release (or drop) the gate
    /// **before** the next [`begin`](DerivationPool::begin) — a begin that overtakes an
    /// unreleased gate makes the affected workers tear down + rehydrate.
    pub fn finish(self) -> PoolGate {
        PoolGate {
            gate: self.stream.finish(),
        }
    }
}

/// The commit-verdict gate, held between "all workers pinned + pushed" and the host
/// COMMIT's outcome. Releasing it sends each worker its terminal `TxFinish` with the
/// verdict inline; dropping it without either releases the workers as an abort (no
/// worker hangs).
#[must_use = "a PoolGate must be committed or aborted, else workers discard the tx as an abort"]
pub struct PoolGate {
    gate: CommitGate,
}

impl PoolGate {
    /// Report a successful host commit. Workers emit `Progressed(N)` after their
    /// speculative `Changed(N)` slices and discard derivation state. The pool trusts
    /// this verdict; it does not check the host connection. Releasing from a pre-commit
    /// hook requires separate host-commit failure handling; see the module docs.
    pub fn commit(self) {
        self.gate.commit();
    }

    /// The host's transaction will not commit: a worker that received pushes has
    /// irreversibly advanced its operator state for a transaction that never happened,
    /// so it tears down + re-hydrates (terminal `Faulted` per hosted query); a worker
    /// that received nothing rolls its snapshot back cleanly.
    pub fn abort(self) {
        self.gate.abort();
    }
}