armdb 0.14.1

//! Generic multi-shard held context: the transaction's answer to `atomic()`,
//! which holds exactly one shard (`ctx.rs:135-160`).

use std::collections::BTreeMap;

use crate::ctx::Family;
use crate::error::{DbError, DbResult};
use crate::held_ops::HeldOps;
use crate::txlog::record::{OpKind, VersionTag};
use crate::txn::participant::{DeferredReplay, PlannedOp, TxLocked};
use crate::watch::WatchWakeBatch;

/// Family-specific data accumulated across [`TxFamily::tx_release`] calls.
/// Typed families stash mutation-era `seize` guards here; Const/Var use `()`.
pub(crate) trait TxReleaseExtra: Default {
    fn absorb(&mut self, other: Self);
}

impl TxReleaseExtra for () {
    fn absorb(&mut self, _other: Self) {}
}

/// What a family must expose so a transaction can hold several of its shards.
pub(crate) trait TxFamily: Family + Sized {
    /// The family's per-shard object — the same one `atomic()` hands to its
    /// closure, so it already implements the mutation seam.
    type Shard<'a>: HeldOps
    where
        Self: 'a;
    /// Concrete deferred owner for this family. Used as the inner type inside
    /// [`tx_deferred`](Self::tx_deferred): implementations box an owned or
    /// `Arc`-backed `Self::Deferred<'static>` (or equivalent), not a borrow of
    /// `self`.
    type Deferred<'a>: DeferredReplay
    where
        Self: 'a;
    /// Per-transaction session data threaded through release into deferred.
    type ReleaseExtra: TxReleaseExtra;

    fn tx_shard_count(&self) -> usize;
    fn tx_shard_of(&self, key: &[u8]) -> usize;
    /// Lock `sid` and build the shard object (mirrors the body of `atomic()`).
    fn tx_make_shard(&self, sid: usize) -> Self::Shard<'_>;
    /// Release the shard's guards (LIFO, through `ShardGuards`) and hand back
    /// the events it buffered plus any family-specific release data (e.g.
    /// mutation-era seize guards for the Typed families).
    fn tx_release(
        shard: Self::Shard<'_>,
    ) -> (
        Vec<<Self as Family>::Owned>,
        WatchWakeBatch,
        Self::ReleaseExtra,
    );
    /// Hand back the deferred hook owner for [`TxLocked::finish`].
    ///
    /// Must return `Box<dyn DeferredReplay>` with an owned or `Arc`-backed
    /// deferred struct (e.g. box `Self::Deferred<'static>` built from
    /// `Arc::clone` of the collection). Do **not** return `Self::Deferred<'_>`
    /// borrowing `self` — `Box<dyn DeferredReplay>` requires `'static`.
    ///
    /// `extra` is the merged output of every `tx_release` in the session.
    fn tx_deferred(
        &self,
        events: Vec<<Self as Family>::Owned>,
        extra: Self::ReleaseExtra,
    ) -> DbResult<Box<dyn DeferredReplay>>;
}

/// Byte-level access on top of the family's typed shard. Takes `&self` because
/// the Typed families need their codec here: the stored value is a live `T`,
/// while a transaction speaks stored bytes.
pub(crate) trait TxRaw: TxFamily {
    fn raw_read(&self, shard: &mut Self::Shard<'_>, key: &[u8]) -> DbResult<Option<Vec<u8>>>;
    fn raw_contains(&self, shard: &mut Self::Shard<'_>, key: &[u8]) -> DbResult<bool>;
    fn raw_put(&self, shard: &mut Self::Shard<'_>, key: &[u8], value: &[u8]) -> DbResult<()>;
    fn raw_delete(&self, shard: &mut Self::Shard<'_>, key: &[u8]) -> DbResult<()>;
    /// Reserve one version per operation, all of which live on `sid`, in apply
    /// order. Bitcask hands out a contiguous GSN block; FixedStore resolves each
    /// key to `(slot, target_meta)`.
    fn reserve_in_shard(
        &self,
        shard: &mut Self::Shard<'_>,
        ops: &[PlannedOp<'_>],
    ) -> DbResult<Vec<VersionTag>>;
    /// Arm the shard so every mutation on this shard carries `tag`.
    fn arm_tag(shard: &mut Self::Shard<'_>, tag: VersionTag) -> DbResult<()>;
    fn disarm_tag(shard: &mut Self::Shard<'_>);
    /// Does this logged operation still need to be applied?
    ///
    /// Answered under the shard lock: FixedStore reads the slot's live state;
    /// Bitcask compares the reserved GSN with the shard's recovered watermark.
    fn needs_replay(
        &self,
        shard: &mut Self::Shard<'_>,
        key: &[u8],
        kind: OpKind,
        tag: VersionTag,
    ) -> DbResult<bool>;
    /// Undo a reservation that was never applied. Bitcask has nothing to undo
    /// (a GSN hole is harmless); FixedStore returns the slot to the allocator.
    fn release_reservation(shard: &mut Self::Shard<'_>, tag: VersionTag);
    /// Stop this shard from writing past a reservation the transaction logged
    /// but never applied.
    ///
    /// Bitcask must: its replay rule asks whether a reserved GSN is above the
    /// shard's recovered watermark, so any later entry buries the gap and the
    /// committed operation is dropped at the next open. FixedStore need not —
    /// its rule reads the reserved slot's exact state, which no write to another
    /// key can disturb, and a write to the *same* key legitimately supersedes
    /// the operation.
    ///
    /// Returns the effective fence gsn, or `None` when this family does not
    /// fence — the return value is what tells the caller whether the shard
    /// belongs in the reported [`crate::FencedShard`] list.
    fn fence_shard(shard: &mut Self::Shard<'_>, tag: VersionTag) -> Option<u64>;
    /// Hook events this shard has buffered so far. Read before and after every
    /// mutation: the difference is how many events that mutation produced, which
    /// is **not** always one — a write that changed nothing produces none.
    fn buffered_events(shard: &Self::Shard<'_>) -> usize;
}

pub(crate) struct TxHeld<'a, F: TxRaw> {
    family: &'a F,
    /// One shard object per locked shard id, ascending — `BTreeMap` so
    /// `into_iter().rev()` releases in reverse acquisition order.
    ///
    /// `Option` so [`Drop`] can take it: `abort`/`finish` leave `None` behind,
    /// and every other way out of a locked section — a `?` in the replay loop,
    /// a panic in a user closure — lands in `Drop` with the map still present.
    shards: Option<BTreeMap<usize, F::Shard<'a>>>,
    /// Shard ids of successful applies, in application order — used to rebuild
    /// the deferred event list so hooks fire in insertion order across shards.
    apply_order: Vec<usize>,
    /// Reserved-but-not-yet-applied versions, with the shard they belong to.
    /// `abort` (and `Drop`, for a panic) hands each one back before releasing
    /// the guards — but only before the log point; see [`logged`](Self::logged).
    reserved: Vec<(usize, VersionTag)>,
    /// Has this transaction's record reached the log?
    ///
    /// Past that point a reservation belongs to **replay**, not to the
    /// allocator: the durable record names the slot, so handing it back would
    /// let an unrelated key take it and make replay either overwrite that key's
    /// row or fail the open. `finish` is not the only way out past the log point
    /// — a panic between the append and the apply loop unwinds through `Drop` —
    /// so the rule lives on the flag rather than on the exit path.
    logged: bool,
}

/// Release whatever is still held, innermost shard first.
///
/// Never a plain `drop` of the shard objects: a Map family's `Locked` is a
/// guard pair, and dropping the tuple releases the engine shard while the index
/// guard is still held — one wakeup per operation becomes two (`ctx.rs`, the S1
/// regression, −52% throughput).
impl<F: TxRaw> Drop for TxHeld<'_, F> {
    fn drop(&mut self) {
        self.release_reservations();
        if let Some(shards) = self.shards.take() {
            for (_, mut shard) in shards.into_iter().rev() {
                // Disarm before releasing. `apply` disarms on both its own exits,
                // but a panic inside the mutation unwinds between the two — and
                // an arm that outlives the lock is taken by the next unrelated
                // write on this shard, which would then land in a transaction's
                // slot at a transaction's version. Disarming an unarmed shard is
                // a no-op, so this is free on every other path.
                F::disarm_tag(&mut shard);
                let (_, _, extra) = F::tx_release(shard);
                drop(extra);
            }
        }
    }
}

impl<'a, F: TxRaw> TxHeld<'a, F> {
    pub(crate) fn lock(family: &'a F, shards: &[usize]) -> Self {
        // Every real `TxParticipant::lock_shards` funnels here, so this one
        // assert covers the whole rank-order deadlock guarantee: two commits
        // must never take the same pair of shards in opposite orders.
        debug_assert!(
            shards.windows(2).all(|w| w[0] < w[1]),
            "lock_shards requires ascending, deduplicated shards, got {shards:?}"
        );
        let mut map = BTreeMap::new();
        for &sid in shards {
            map.entry(sid).or_insert_with(|| family.tx_make_shard(sid));
        }
        Self {
            family,
            shards: Some(map),
            apply_order: Vec::new(),
            reserved: Vec::new(),
            logged: false,
        }
    }

    fn shard_for(&mut self, key: &[u8]) -> DbResult<&mut F::Shard<'a>> {
        let sid = self.family.tx_shard_of(key);
        self.shards
            .as_mut()
            .and_then(|m| m.get_mut(&sid))
            .ok_or(DbError::ShardMismatch)
    }

    fn release_reservations(&mut self) {
        if self.logged {
            // Past the log point. The record is durable and names these
            // versions; replay will finish the writes into them, so they stay
            // allocated — and the shard is fenced so nothing can write past
            // them meanwhile. This is the `Drop` path — there is nobody left to
            // report the fenced shards to, so they are discarded.
            self.forget_reservations(None);
            return;
        }
        let reserved = std::mem::take(&mut self.reserved);
        let Some(shards) = self.shards.as_mut() else {
            return;
        };
        for (sid, tag) in reserved {
            if let Some(shard) = shards.get_mut(&sid) {
                F::release_reservation(shard, tag);
            }
        }
    }

    /// Give up every reservation that no apply spent, without handing it back.
    ///
    /// Past the log point each one names an operation the apply never
    /// completed — it failed, it was skipped because an earlier one on its
    /// shard failed, or a panic unwound through the locked section. Replay owns
    /// those versions now, so the shard is fenced: a later write there would
    /// take a higher version, and for Bitcask that version becomes the recovered
    /// watermark and makes replay skip the committed operation.
    ///
    /// Before the log point nothing was committed and there is nothing to fence;
    /// the reservations are simply dropped (`abort` and `Drop` hand them back
    /// through [`release_reservations`](Self::release_reservations) instead).
    fn forget_reservations(&mut self, out: Option<&mut Vec<(usize, u64)>>) {
        let reserved = std::mem::take(&mut self.reserved);
        if !self.logged {
            return;
        }
        let Some(shards) = self.shards.as_mut() else {
            return;
        };
        // One transaction can hold several reservations on one shard, and the
        // shard keeps the lowest. Merge here so the report names each shard
        // once, with the value the shard actually holds.
        let mut merged: BTreeMap<usize, u64> = BTreeMap::new();
        for (sid, tag) in reserved {
            if let Some(shard) = shards.get_mut(&sid)
                && let Some(gsn) = F::fence_shard(shard, tag)
            {
                merged
                    .entry(sid)
                    .and_modify(|g| *g = (*g).min(gsn))
                    .or_insert(gsn);
            }
        }
        if let Some(out) = out {
            out.extend(merged);
        }
    }

    /// Arm the reserved version, run one mutation, disarm, and record how many hook
    /// events it produced.
    ///
    /// Both bookkeeping steps are load-bearing:
    ///
    /// * **Disarming unconditionally.** A mutation that changes nothing — a
    ///   delete of an absent key, in every family — appends no entry, so it
    ///   never consumes the armed version. Leaving it armed past the lock release
    ///   would hand a transaction's version to an unrelated later write on this
    ///   shard (a plain `put`, compaction, a TTL tombstone), breaking the
    ///   per-shard ascending-GSN invariant the replay rule depends on — and
    ///   tripping `arm_tag`'s single-shot `debug_assert` in debug builds.
    ///   Arming is single-shot, so disarming after a mutation that *did* append
    ///   is a no-op.
    /// * **Counting events per operation.** One mutation does not always mean
    ///   one buffered event, so the count has to be measured rather than
    ///   assumed. `apply_order` gets one entry per *event*, which is what
    ///   [`TxLocked::finish`] replays.
    fn apply(
        &mut self,
        key: &[u8],
        tag: VersionTag,
        mutate: impl FnOnce(&F, &mut F::Shard<'a>) -> DbResult<()>,
    ) -> DbResult<usize> {
        let sid = self.family.tx_shard_of(key);
        let family = self.family;
        let shard = self.shard_for(key)?;
        let before = F::buffered_events(shard);
        F::arm_tag(shard, tag)?;
        // T10 Task 5: `mutate` reaches the same per-family `Shard::put`/
        // `delete` the facade and `atomic()` use, so it would otherwise trip
        // the non-transactional claim guard meant for writes outside
        // `db.tx()`. Claim admission inside a transaction is
        // `resolve_claims`'s job (`txn::projection_commit`, T10 Task 7), not
        // this guard's — see `unique` module docs for `SkipClaimGuard`. This
        // one call site covers every family's `apply_put`/`apply_delete`,
        // and — since crash-recovery replay reapplies logged operations
        // through this same `TxLocked` entry point — replay too.
        let outcome = {
            let _skip_claim_guard = crate::unique::SkipClaimGuard::enter();
            let _stream_apply_guard = crate::txn_depth::StreamApplyGuard::enter();
            mutate(family, shard)
        };
        // Unconditional: a mutation that changed nothing — a delete of an absent
        // key, in every family — never consumes the arm, and leaving it set would
        // hand a transaction's version to an unrelated later write on this shard.
        F::disarm_tag(shard);
        outcome?;
        let produced = F::buffered_events(shard).saturating_sub(before);
        for _ in 0..produced {
            self.apply_order.push(sid);
        }
        // The write landed (or tried to): the reservation is spent either way.
        // A failed apply leaves the version reserved rather than returning it,
        // because the record is already durable and replay will finish the write
        // into that exact slot.
        if let Some(pos) = self.reserved.iter().position(|(_, t)| *t == tag) {
            self.reserved.swap_remove(pos);
        }
        Ok(produced)
    }
}

impl<'a, F: TxRaw> TxLocked for TxHeld<'a, F> {
    fn read(&mut self, key: &[u8]) -> DbResult<Option<Vec<u8>>> {
        let family = self.family;
        let shard = self.shard_for(key)?;
        family.raw_read(shard, key)
    }

    fn contains(&mut self, key: &[u8]) -> DbResult<bool> {
        let family = self.family;
        let shard = self.shard_for(key)?;
        family.raw_contains(shard, key)
    }

    fn reserve(&mut self, ops: &[PlannedOp<'_>]) -> DbResult<Vec<VersionTag>> {
        // One call per shard, in apply order within that shard: Bitcask needs a
        // contiguous ascending block, FixedStore needs each key's own slot.
        let mut per_shard: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
        for (i, op) in ops.iter().enumerate() {
            per_shard
                .entry(self.family.tx_shard_of(op.key))
                .or_default()
                .push(i);
        }
        let mut tags: Vec<Option<VersionTag>> = vec![None; ops.len()];
        for (sid, indices) in per_shard {
            let planned: Vec<PlannedOp<'_>> = indices
                .iter()
                .map(|&i| PlannedOp {
                    key: ops[i].key,
                    delete: ops[i].delete,
                })
                .collect();
            let family = self.family;
            let shard = self
                .shards
                .as_mut()
                .and_then(|m| m.get_mut(&sid))
                .ok_or(DbError::ShardMismatch)?;
            let reserved = family.reserve_in_shard(shard, &planned)?;
            if reserved.len() != indices.len() {
                return Err(DbError::Internal("reserve_in_shard returned too few tags"));
            }
            for (slot, &i) in reserved.iter().zip(indices.iter()) {
                tags[i] = Some(*slot);
                self.reserved.push((sid, *slot));
            }
        }
        tags.into_iter()
            .map(|t| t.ok_or(DbError::Internal("missing reserved tag")))
            .collect()
    }

    fn apply_put(&mut self, key: &[u8], value: &[u8], tag: VersionTag) -> DbResult<usize> {
        self.apply(key, tag, |family, shard| family.raw_put(shard, key, value))
    }

    fn apply_delete(&mut self, key: &[u8], tag: VersionTag) -> DbResult<usize> {
        self.apply(key, tag, |family, shard| family.raw_delete(shard, key))
    }

    fn needs_replay(&mut self, key: &[u8], kind: OpKind, tag: VersionTag) -> DbResult<bool> {
        let family = self.family;
        let shard = self.shard_for(key)?;
        family.needs_replay(shard, key, kind, tag)
    }

    fn mark_logged(&mut self) {
        self.logged = true;
    }

    #[cfg(test)]
    fn unapplied_reservations(&self) -> usize {
        self.reserved.len()
    }

    #[cfg(test)]
    fn tx_debug_arm(&mut self, key: &[u8], tag: VersionTag) -> DbResult<()> {
        let shard = self.shard_for(key)?;
        F::arm_tag(shard, tag)
    }

    fn abort(self: Box<Self>) {
        // `release_reservations` runs in `Drop` before the guards go.
        drop(self);
    }

    fn finish(
        mut self: Box<Self>,
        fenced: &mut Vec<(usize, u64)>,
        watch_wakes: &mut WatchWakeBatch,
    ) -> DbResult<Box<dyn DeferredReplay>> {
        use std::collections::VecDeque;

        // Any reservation still here belongs to an operation whose apply failed
        // or was skipped. Replay owns those versions; the shard is fenced so no
        // later write can bury the gap it has to fill.
        self.forget_reservations(Some(fenced));

        let family = self.family;
        let apply_order = std::mem::take(&mut self.apply_order);
        // Taking the map here is what tells `Drop` the shards are already
        // released; an early `?` below then leaves nothing double-released.
        let shards = self.shards.take().unwrap_or_default();
        // Release locks LIFO (reverse acquisition), but keep per-shard event
        // queues so we can rebuild the deferred list in application order.
        let mut per_shard: BTreeMap<usize, VecDeque<<F as crate::ctx::Family>::Owned>> =
            BTreeMap::new();
        let mut extra = F::ReleaseExtra::default();
        for (sid, shard) in shards.into_iter().rev() {
            let (ev, wakes, ex) = F::tx_release(shard);
            per_shard.insert(sid, VecDeque::from(ev));
            watch_wakes.extend(wakes);
            extra.absorb(ex);
        }
        // `apply_order` holds one shard id per *event* (see `apply`), so the two
        // counts agree by construction; the check stays as a cheap guard against
        // a family whose `buffered_events` disagrees with what it buffers.
        let mut events = Vec::with_capacity(apply_order.len());
        let buffered: usize = per_shard.values().map(|q| q.len()).sum();
        if buffered != apply_order.len() {
            return Err(DbError::Internal(
                "buffered hook events and apply order disagree",
            ));
        }
        for sid in apply_order {
            let ev =
                per_shard
                    .get_mut(&sid)
                    .and_then(|q| q.pop_front())
                    .ok_or(DbError::Internal(
                        "missing buffered hook event for an applied write",
                    ))?;
            events.push(ev);
        }
        family.tx_deferred(events, extra)
    }
}