lora-database 0.16.1

//! Write-path lock plumbing for [`Database<S>`].
//!
//! The pieces in here form one shared mutating shape and the two
//! published entry points that wrap it:
//!
//! * [`Database::run_with_durable_recorder`] — single canonical write
//!   shape. Acquires the writer mutex, clones the current store into a
//!   staged graph, arms the durable recorder when a WAL is attached,
//!   runs the caller's closure, and only publishes the staged graph
//!   after a successful WAL commit. Used by both the auto-commit OCC
//!   fast path and the admin mutator path so rollback semantics and WAL
//!   ordering live in exactly one place.
//! * [`Database::with_logged_store_mut`] — admin entry. Dispatches
//!   InMemoryGraph traffic to `run_with_durable_recorder`; falls back
//!   to the pessimistic clone-then-publish path for any other backend.
//! * [`Database::with_logged_write_guard`] — pessimistic path. Stages
//!   a clone, runs the closure under the WAL scope, and atomically
//!   publishes on success (or aborts per `WalAbortPolicy` on failure).
//!   Reserved for backends that don't expose the
//!   `set_mutation_recorder` hook.
//! * [`WriteGuard`] — working copy + writer-mutex lease produced by
//!   [`Database::write_store`]. Mutate, then `publish()` to swap the
//!   new state into the live store, or drop to discard.
//! * [`Database::write_store_deadline`] / [`Database::read_store_deadline`]
//!   — deadline-aware variants that participate in the cooperative
//!   query-timeout flow.
//!
//! Failed staged writes are discarded without publishing. The
//! `WalAbortPolicy::PoisonIfMutated` fallback remains for paths that
//! mutate a live graph directly.

use std::any::{Any, TypeId};
use std::sync::{Arc, MutexGuard, TryLockError};
use std::time::Duration;

use web_time::Instant;

use anyhow::{anyhow, Result};
use lora_executor::ExecutorError;
use lora_store::{
    DeletedRecordSink, GraphStorage, GraphStorageMut, InMemoryGraph, MutationEvent,
    MutationRecorder,
};

use crate::changes::{CaptureRecorder, PreImageSink};
use lora_wal::WalRecorder;

use crate::database::Database;
use crate::wal::write_scope::{ensure_wal_query_can_start, WalAbortPolicy, WalWriteScope};

use super::replay::{install_deleted_sink_if_inmemory, install_recorder_if_inmemory};

/// Working copy + writer-mutex lease produced by
/// [`Database::write_store`]. The caller mutates the inner `S`, then
/// calls [`WriteGuard::publish`] to atomically swap the new state into
/// the live store. Dropping without `publish` discards the staged copy
/// (rollback semantics) and releases the writer lock, leaving the
/// authoritative store unchanged.
pub(crate) struct WriteGuard<'db, S> {
    db: &'db Database<S>,
    /// Held for the lifetime of the guard so commits are serialized
    /// (and so the WAL records appear in commit order).
    _writer_lock: MutexGuard<'db, ()>,
    /// `Some` until `publish` consumes the staged graph, `None` after.
    /// Drop on `None` is a no-op; drop on `Some` discards the staged
    /// changes.
    staged: Option<S>,
}

impl<S> WriteGuard<'_, S>
where
    S: Send + Sync + 'static,
{
    /// Atomically replace the live store with the staged graph. After
    /// this returns, subsequent reads see the new state.
    pub(crate) fn publish(mut self) {
        self.publish_in_place();
    }

    /// Publish the staged graph but keep holding the writer lock, so the
    /// caller can finish commit-ordered work (the change feed) first.
    pub(crate) fn publish_in_place(&mut self) {
        if let Some(staged) = self.staged.take() {
            self.db.store.store(Arc::new(staged));
        }
    }

    pub(crate) fn staged_or_error(&self) -> Result<&S> {
        self.staged
            .as_ref()
            .ok_or_else(|| anyhow!("staged graph already published or taken"))
    }

    pub(crate) fn staged_mut_or_error(&mut self) -> Result<&mut S> {
        self.staged
            .as_mut()
            .ok_or_else(|| anyhow!("staged graph already published or taken"))
    }
}

impl<S> Database<S>
where
    S: GraphStorage + GraphStorageMut + Any + Clone + Send + Sync + 'static,
{
    /// Take the writer lease and clone the current snapshot into a
    /// staged working copy. The caller mutates the staged graph and
    /// either calls `publish()` to install it (atomically replacing the live
    /// store) or drops the guard to discard the changes.
    pub(crate) fn write_store(&self) -> WriteGuard<'_, S> {
        let lock = self
            .writer
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        let snapshot = self.store.load_full();
        let staged: S = (*snapshot).clone();
        WriteGuard {
            db: self,
            _writer_lock: lock,
            staged: Some(staged),
        }
    }

    pub(crate) fn read_store_with_epoch_deadline(
        &self,
        _deadline: Option<Instant>,
    ) -> Result<(Arc<S>, u64)> {
        // Same lock profile as `read_store_deadline`, but includes the
        // coherent epoch needed for cheap plan-cache keys.
        Ok(self.store.load_full_with_epoch())
    }

    pub(crate) fn write_store_deadline(
        &self,
        deadline: Option<Instant>,
    ) -> Result<WriteGuard<'_, S>> {
        let Some(deadline) = deadline else {
            return Ok(self.write_store());
        };

        loop {
            match self.writer.try_lock() {
                Ok(lock) => {
                    let snapshot = self.store.load_full();
                    let staged: S = (*snapshot).clone();
                    return Ok(WriteGuard {
                        db: self,
                        _writer_lock: lock,
                        staged: Some(staged),
                    });
                }
                Err(TryLockError::Poisoned(poisoned)) => {
                    let lock = poisoned.into_inner();
                    let snapshot = self.store.load_full();
                    let staged: S = (*snapshot).clone();
                    return Ok(WriteGuard {
                        db: self,
                        _writer_lock: lock,
                        staged: Some(staged),
                    });
                }
                Err(TryLockError::WouldBlock)
                    if lora_executor::cancel::deadline_reached(deadline) =>
                {
                    return Err(ExecutorError::QueryTimeout.into());
                }
                Err(TryLockError::WouldBlock) => {
                    std::thread::sleep(Duration::from_millis(1));
                }
            }
        }
    }

    pub(crate) fn observe_snapshot_commit_if_needed(
        &self,
        store: &S,
        recorder: &WalRecorder,
    ) -> Result<()> {
        let Some(snapshots) = &self.snapshots else {
            return Ok(());
        };
        let graph = (store as &dyn Any)
            .downcast_ref::<InMemoryGraph>()
            .ok_or_else(|| anyhow!("managed snapshots require InMemoryGraph storage"))?;
        snapshots.observe_commit(graph, recorder)?;
        Ok(())
    }

    /// Admin entry into the live-mutate write path.
    ///
    /// For `InMemoryGraph` (the default backend) this delegates to
    /// [`Self::run_with_durable_recorder`], so the same staged shape
    /// governs both query auto-commit and the direct admin /
    /// `graph_api` mutators. For other backends — which don't expose
    /// `set_mutation_recorder` and therefore can't redirect buffered
    /// events — we fall back to the pessimistic clone-then-publish
    /// path with `WalAbortPolicy::AbortOnly`.
    pub(crate) fn with_logged_store_mut<R>(
        &self,
        f: impl FnOnce(&mut S) -> Result<R>,
    ) -> Result<R> {
        if TypeId::of::<S>() == TypeId::of::<InMemoryGraph>() {
            return self.run_with_durable_recorder(f);
        }
        let guard = self.write_store();
        self.with_logged_write_guard(guard, WalAbortPolicy::AbortOnly, f)
    }

    /// Canonical mutating shape for `InMemoryGraph` writes. Used by
    /// both the auto-commit OCC fast path and the admin mutator path so
    /// the WAL bracketing and publish rules live in exactly one place.
    ///
    /// Sequence:
    ///
    /// 1. Take the writer mutex (serialises commit ordering with
    ///    other writers).
    /// 2. If a WAL is attached, check it isn't poisoned and `arm` the
    ///    durable recorder so subsequent `MutationEvent`s buffer
    ///    inside it.
    /// 3. Clone the current store into a staged graph.
    /// 4. Run `f` against the staged graph. On `Err`: abort the
    ///    recorder, discard the staged graph, and leave the live graph
    ///    unchanged. On `Ok`: commit the recorder, trigger
    ///    managed-snapshot accounting if a commit record was written,
    ///    reinstall the durable recorder, and publish the staged graph.
    pub(crate) fn run_with_durable_recorder<R>(
        &self,
        f: impl FnOnce(&mut S) -> Result<R>,
    ) -> Result<R> {
        let guard = self.write_store();
        self.with_logged_write_guard(guard, WalAbortPolicy::AbortOnly, f)
    }

    /// Low-overhead mutating shape for plans that are proven simple enough to
    /// execute directly against the live graph. This preserves the old
    /// `Arc::make_mut` hot path for CREATE/SET-style writes where staging the
    /// whole graph would dominate the actual mutation cost.
    ///
    /// Complex or fallible plans must use [`Self::run_with_durable_recorder`]
    /// instead so errors discard a staged graph. If an allegedly-safe live
    /// plan still fails after emitting WAL events, poison the recorder: the
    /// live graph may have advanced beyond durable state, so callers should
    /// recover from snapshot + WAL.
    pub(crate) fn run_live_fast_with_durable_recorder<R>(
        &self,
        may_delete: bool,
        f: impl FnOnce(&mut S) -> Result<R>,
    ) -> Result<R> {
        let _commit_lock = self
            .writer
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());

        if let Some(rec) = self.wal.as_ref() {
            ensure_wal_query_can_start(rec)?;
            rec.arm()?;
        }

        let capture = self.changes.is_active();
        let mut handle = self.store.write();
        let capture_rec =
            (capture && self.wal.is_none()).then(|| Arc::new(CaptureRecorder::default()));
        // A captured delete needs the deleted records. The write mutates the
        // live graph in place, so collect them as the store drops them
        // instead of keeping a pre-write copy.
        let deleted = (capture && may_delete).then(|| Arc::new(PreImageSink::default()));

        let result = {
            let live = handle.as_mut();
            if let Some(rec) = self.wal.as_ref() {
                install_recorder_if_inmemory(live, Some(rec.clone() as Arc<dyn MutationRecorder>));
            } else if let Some(cap) = &capture_rec {
                install_recorder_if_inmemory(live, Some(cap.clone() as Arc<dyn MutationRecorder>));
            }
            if let Some(sink) = &deleted {
                install_deleted_sink_if_inmemory(
                    live,
                    Some(sink.clone() as Arc<dyn DeletedRecordSink>),
                );
            }
            let result = f(live);
            if capture_rec.is_some() {
                install_recorder_if_inmemory(live, None);
            }
            if deleted.is_some() {
                install_deleted_sink_if_inmemory(live, None);
            }
            result
        };

        let mut committed: Option<(Option<lora_wal::Lsn>, Vec<MutationEvent>)> = None;
        if let Some(rec) = self.wal.as_ref() {
            match &result {
                Ok(_) => {
                    let wrote = if capture {
                        match rec.commit_capture()? {
                            Some((lsn, events)) => {
                                committed = Some((Some(lsn), events));
                                true
                            }
                            None => false,
                        }
                    } else {
                        rec.commit()?.wrote()
                    };
                    if wrote {
                        let live = handle.snapshot();
                        self.observe_snapshot_commit_if_needed(&*live, rec)?;
                    }
                }
                Err(_) => {
                    let mutated = rec.abort().unwrap_or(false);
                    if mutated {
                        rec.poison(crate::database::QUERY_FAILURE_POISON);
                    }
                }
            }
        } else if let (Some(cap), Ok(_)) = (&capture_rec, &result) {
            committed = Some((None, cap.take()));
        }

        if let Some((lsn, events)) = committed {
            let post = handle.snapshot();
            // Release the store write lock before fan-out; the writer mutex
            // still orders this batch against the next commit.
            drop(handle);
            let pre = deleted.map(|sink| sink.take()).unwrap_or_default();
            self.publish_changes_with(lsn, &events, &pre, &post);
        }

        result
    }

    /// Run `f` against the staged graph inside a WAL transaction. On
    /// success, atomically publishes the staged graph to the live
    /// store; on error, the staged copy is dropped (no observable
    /// state change) and the WAL is aborted per `abort_policy`.
    pub(crate) fn with_logged_write_guard<R>(
        &self,
        mut guard: WriteGuard<'_, S>,
        abort_policy: WalAbortPolicy,
        f: impl FnOnce(&mut S) -> Result<R>,
    ) -> Result<R> {
        let capture = self.changes.is_active();
        let Some(rec) = self.wal.clone() else {
            // No WAL: just run the closure, publish on success.
            let capture_rec = capture.then(|| Arc::new(CaptureRecorder::default()));
            let staged = guard.staged_mut_or_error()?;
            if let Some(cap) = &capture_rec {
                install_recorder_if_inmemory(
                    staged,
                    Some(cap.clone() as Arc<dyn MutationRecorder>),
                );
            }
            let result = f(staged);
            if capture_rec.is_some() {
                install_recorder_if_inmemory(staged, None);
            }
            if result.is_ok() {
                match capture_rec {
                    Some(cap) => {
                        let events = cap.take();
                        let pre = self.store.load_full();
                        guard.publish_in_place();
                        let post = self.store.load_full();
                        self.publish_changes(None, &events, Some(&*pre), &post);
                    }
                    None => guard.publish(),
                }
            }
            return result;
        };

        // Install the durable recorder on the staged graph so the
        // executor's mutations fire into it. `InMemoryGraph::clone`
        // intentionally drops the recorder, so the staged copy starts
        // without one.
        {
            let staged = guard.staged_mut_or_error()?;
            install_recorder_if_inmemory(staged, Some(rec.clone() as Arc<dyn MutationRecorder>));
        }

        let scope = WalWriteScope::arm(&rec, abort_policy)?;
        let result = {
            let staged = guard.staged_mut_or_error()?;
            f(staged)
        };
        let (wrote_commit, captured) = scope.finish(&result, capture)?;
        if wrote_commit {
            let staged = guard.staged_or_error()?;
            self.observe_snapshot_commit_if_needed(staged, &rec)?;
        }

        // Strip the per-mutation recorder before publish — the new live
        // store carries the durable recorder reinstalled below.
        {
            let staged = guard.staged_mut_or_error()?;
            install_recorder_if_inmemory(staged, None);
        }

        if result.is_ok() {
            // Reinstall the durable recorder on staged so the new live
            // graph keeps observing mutations after the swap.
            {
                let staged = guard.staged_mut_or_error()?;
                install_recorder_if_inmemory(
                    staged,
                    Some(rec.clone() as Arc<dyn MutationRecorder>),
                );
            }
            match captured {
                Some((lsn, events)) => {
                    let pre = self.store.load_full();
                    guard.publish_in_place();
                    let post = self.store.load_full();
                    self.publish_changes(Some(lsn), &events, Some(&*pre), &post);
                }
                None => guard.publish(),
            }
        }
        result
    }
}