armdb 0.9.0

sharded bitcask key-value storage optimized for NVMe
Documentation
//! Lock-context seam: one mutation ordering for facade (`Owned`) and
//! `atomic()` (`Held`). See docs/superpowers/specs/26-07-30-collection-contexts-and-rpc-adapter.md.

use crate::DbResult;
use crate::error::DbError;
use crate::sync;

/// Shard guards, released in **reverse acquisition order** (LIFO).
///
/// The Map families take the engine shard mutex first and the per-shard index
/// mutex second, and carry both in a `Locked` tuple. A plain `drop` of that
/// tuple releases `.0` before `.1` — the **outer** lock while the inner one is
/// still held. A waiter then takes the engine shard only to park again on the
/// index, turning one wakeup per operation into two. Measured cost when the
/// seam did that: −52% throughput on `typed_map`/`var_map` at eight writers
/// (`armdb/docs/ctx-seam-bench.md`).
///
/// Releasing through this trait instead of `drop` makes the order a property of
/// the guard shape rather than of each family's teardown code, so a family
/// cannot get it wrong — including families added later.
pub(crate) trait ShardGuards {
    /// Drop the guards, innermost (acquired last) first.
    fn release(self);
}

impl ShardGuards for () {
    #[inline]
    fn release(self) {}
}

impl<T> ShardGuards for sync::MutexGuard<'_, T> {
    #[inline]
    fn release(self) {
        drop(self);
    }
}

/// `.0` is acquired first and `.1` second, so `.1` is released first.
impl<A: ShardGuards, B: ShardGuards> ShardGuards for (A, B) {
    #[inline]
    fn release(self) {
        let (outer, inner) = self;
        inner.release();
        outer.release();
    }
}

pub(crate) trait Family {
    /// Mutex guards held for one shard. Family-specific: one for Tree
    /// families, two (engine shard + per-shard index) for Maps.
    ///
    /// Released via [`ShardGuards::release`], never `drop` — see the trait.
    type Locked<'a>: ShardGuards
    where
        Self: 'a;

    /// Borrowed view of a mutation, produced under the lock. May borrow the
    /// caller's input (`&'v [u8]`), which outlives the locked section.
    type Event<'v>;

    /// Owned form the held context buffers for later replay.
    type Owned;

    fn lock_shard(&self, sid: usize) -> Self::Locked<'_>;

    fn should_sync<'a>(_locked: &Self::Locked<'a>) -> bool
    where
        Self: 'a,
    {
        false
    }
    fn sync(&self, _sid: usize) -> DbResult<()> {
        Ok(())
    }

    fn fire(&self, ev: Self::Event<'_>);
    fn own(ev: Self::Event<'_>) -> Self::Owned;
    fn fire_owned(&self, ev: &Self::Owned);

    fn apply_projection_for_event(&self, _ev: &Self::Event<'_>) -> DbResult<()> {
        Ok(())
    }

    fn apply_projection_for_owned(&self, _ev: &Self::Owned) -> DbResult<()> {
        Ok(())
    }
}

/// `'l` is a trait parameter — do NOT write `f: impl FnOnce(&mut F::Locked<'_>)`
/// (that HRB forces `F: 'static` via the GAT bound; E0309/E0311).
pub(crate) trait Ctx<'l, F: Family + 'l> {
    /// Run `f` under the shard lock, then apply this context's teardown.
    ///
    /// The mutation event leaves through the `&mut Option<Event>` out parameter
    /// rather than the return value. Returning `(R, Option<Event>)` instead made
    /// the closure hand a fat aggregate back across a generic boundary the
    /// optimizer would not collapse, and cost a measurable ~5% on the tightest
    /// benchmark row (`const_tree/4`); the out parameter removes it. See
    /// `armdb/docs/ctx-seam-bench.md`.
    fn with_shard<'v, R>(
        &mut self,
        family: &'l F,
        sid: usize,
        f: impl FnOnce(&mut F::Locked<'l>, &mut Option<F::Event<'v>>) -> DbResult<R>,
    ) -> DbResult<R>;
}

/// Facade context — ZST. Fires borrowed events after unlock; never calls `own`.
#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct Owned;

impl<'l, F: Family + 'l> Ctx<'l, F> for Owned {
    #[inline]
    fn with_shard<'v, R>(
        &mut self,
        family: &'l F,
        sid: usize,
        f: impl FnOnce(&mut F::Locked<'l>, &mut Option<F::Event<'v>>) -> DbResult<R>,
    ) -> DbResult<R> {
        let mut locked = family.lock_shard(sid);
        let mut ev: Option<F::Event<'v>> = None;
        // Not `f(...)?`: an early return would drop the guards through their
        // destructor, which for the Map families' tuple is the inverted order.
        // `should_sync` stays off the error path so the ctx-gates counts hold.
        let r = match f(&mut locked, &mut ev) {
            Ok(r) => r,
            Err(e) => {
                locked.release();
                return Err(e);
            }
        };
        let needs_sync = F::should_sync(&locked);
        locked.release();
        let mut proj_err: Option<DbError> = None;
        if let Some(ref ev) = ev
            && let Err(e) = family.apply_projection_for_event(ev)
        {
            proj_err = Some(e);
        }
        let mut sync_err: Option<DbError> = None;
        if needs_sync && let Err(e) = family.sync(sid) {
            sync_err = Some(e);
        }
        if let Some(ev) = ev {
            family.fire(ev);
        }
        if let (Some(sync_e), Some(proj_e)) = (&sync_err, &proj_err) {
            tracing::error!(
                sync = %sync_e,
                projection = %proj_e,
                "sync error outranked projection error on facade teardown"
            );
        }
        if let Some(e) = sync_err.or(proj_err) {
            return Err(e);
        }
        Ok(r)
    }
}

/// Held context inside `atomic()` — lock already acquired for `sid`.
pub(crate) struct Held<'a, F: Family + 'a> {
    pub locked: F::Locked<'a>,
    pub sid: usize,
    pub events: Vec<F::Owned>,
}

impl<'a, F: Family + 'a> Ctx<'a, F> for Held<'a, F> {
    #[inline]
    fn with_shard<'v, R>(
        &mut self,
        family: &'a F,
        sid: usize,
        f: impl FnOnce(&mut F::Locked<'a>, &mut Option<F::Event<'v>>) -> DbResult<R>,
    ) -> DbResult<R> {
        if sid != self.sid {
            // Callers should check_shard before this; belt-and-suspenders.
            return Err(crate::DbError::ShardMismatch);
        }
        let mut ev: Option<F::Event<'v>> = None;
        let r = f(&mut self.locked, &mut ev)?;
        if let Some(ev) = ev {
            self.events.push(F::own(ev));
        }
        let _ = family; // sync deferred to atomic() teardown
        Ok(r)
    }
}

#[cfg(test)]
mod release_order {
    use std::cell::RefCell;
    use std::rc::Rc;

    use super::*;

    /// Records its label when released, so a `ShardGuards` release order is
    /// observable without depending on timing.
    struct Recorder(&'static str, Rc<RefCell<Vec<&'static str>>>);

    impl ShardGuards for Recorder {
        fn release(self) {
            self.1.borrow_mut().push(self.0);
        }
    }

    /// A two-guard `Locked` must release the **inner** guard (acquired last)
    /// before the outer one. Releasing the outer first lets a waiter take it and
    /// immediately park on the still-held inner guard — two wakeups per
    /// operation instead of one, measured at −52% throughput on the Map
    /// families at eight writers.
    #[test]
    fn pair_releases_inner_guard_first() {
        let log = Rc::new(RefCell::new(Vec::new()));
        let pair = (
            Recorder("outer/engine_shard", log.clone()),
            Recorder("inner/index", log.clone()),
        );

        pair.release();

        assert_eq!(*log.borrow(), ["inner/index", "outer/engine_shard"]);
    }

    /// Nested pairs stay LIFO all the way down.
    #[test]
    fn nested_pairs_release_innermost_first() {
        let log = Rc::new(RefCell::new(Vec::new()));
        let nested = (
            Recorder("a", log.clone()),
            (Recorder("b", log.clone()), Recorder("c", log.clone())),
        );

        nested.release();

        assert_eq!(*log.borrow(), ["c", "b", "a"]);
    }

    /// The real guard pair unlocks both mutexes, in LIFO order: after
    /// `release`, both are free.
    #[test]
    fn real_guard_pair_unlocks_both() {
        let engine_shard = sync::Mutex::new(1u32);
        let index = sync::Mutex::new(2u32);

        let guards = (sync::lock(&engine_shard), sync::lock(&index));
        guards.release();

        assert!(engine_shard.try_lock().is_some());
        assert!(index.try_lock().is_some());
    }
}

#[cfg(test)]
mod compile_gate {
    use super::*;

    struct StubFamily;

    impl Family for StubFamily {
        type Locked<'a> = ();
        type Event<'v> = &'v [u8];
        type Owned = Vec<u8>;

        fn lock_shard(&self, _sid: usize) -> Self::Locked<'_> {}

        fn fire(&self, _ev: Self::Event<'_>) {}

        fn own(ev: Self::Event<'_>) -> Self::Owned {
            ev.to_vec()
        }

        fn fire_owned(&self, _ev: &Self::Owned) {}
    }

    #[test]
    fn owned_and_held_seam_compiles() {
        let family = StubFamily;
        let mut owned = Owned;
        owned.with_shard(&family, 0, |_locked, _ev| Ok(())).unwrap();

        let locked = family.lock_shard(0);
        let mut held = Held {
            locked,
            sid: 0,
            events: Vec::new(),
        };
        held.with_shard(&family, 0, |_locked, ev| {
            *ev = Some(b"x" as &[u8]);
            Ok(())
        })
        .unwrap();
        assert_eq!(held.events.len(), 1);
    }

    struct ProjSyncFamily {
        sync_err: bool,
        proj_err: bool,
    }

    impl Family for ProjSyncFamily {
        type Locked<'a> = ();
        type Event<'v> = ();
        type Owned = ();

        fn lock_shard(&self, _sid: usize) -> Self::Locked<'_> {}

        fn should_sync<'a>(_locked: &Self::Locked<'a>) -> bool
        where
            Self: 'a,
        {
            true
        }

        fn sync(&self, _sid: usize) -> DbResult<()> {
            if self.sync_err {
                Err(DbError::Internal("sync failed"))
            } else {
                Ok(())
            }
        }

        fn apply_projection_for_event(&self, _ev: &Self::Event<'_>) -> DbResult<()> {
            if self.proj_err {
                Err(DbError::Internal("projection failed"))
            } else {
                Ok(())
            }
        }

        fn fire(&self, _ev: Self::Event<'_>) {}

        fn own(_ev: Self::Event<'_>) -> Self::Owned {}

        fn fire_owned(&self, _ev: &Self::Owned) {}
    }

    #[test]
    fn sync_error_outranks_projection_error() {
        let family = ProjSyncFamily {
            sync_err: true,
            proj_err: true,
        };
        let mut owned = Owned;
        let err = owned
            .with_shard(&family, 0, |_locked, ev| {
                *ev = Some(());
                Ok(())
            })
            .unwrap_err();
        assert!(matches!(err, DbError::Internal("sync failed")));
    }
}