armdb 0.12.0

//! `TypedTree`/`TypedMap` transaction participants — codec at the byte seam,
//! seize-guard-owning deferred replay for hook/retire work.

use std::sync::Arc;

use crate::codec::Codec;
use crate::ctx::Family;
use crate::error::{DbError, DbResult};
use crate::hook::TypedWriteHook;
use crate::key::Key;
use crate::txlog::record::{OpKind, VersionTag};
use crate::txn::TxTarget;
use crate::txn::held::{TxFamily, TxHeld, TxRaw, TxReleaseExtra};
use crate::txn::participant::{
    ClaimOp, DeferredReplay, PlannedOp, TxBackend, TxCollectionName, TxCounterDelta,
    TxDerivedWrite, TxLocked, TxParticipant,
};
use crate::typed_map::{MutationEvents as MapMutationEvents, TypedMap, TypedMapShard};
use crate::typed_tree::{MutationEvents, TypedShard, TypedTree};
use crate::watch::WatchWakeBatch;

/// Mutation-era seize guards collected from [`tx_release`], kept alive through
/// hook replay — the same discipline as [`TypedTree::atomic`].
#[derive(Default)]
pub(crate) struct TypedTxGuards(Vec<seize::LocalGuard<'static>>);

impl TxReleaseExtra for TypedTxGuards {
    fn absorb(&mut self, mut other: Self) {
        self.0.append(&mut other.0);
    }
}

/// Move a shard's mutation guard into deferred storage (`Box<dyn DeferredReplay>`).
fn retain_mutation_guard(guard: seize::LocalGuard<'_>) -> seize::LocalGuard<'static> {
    // SAFETY: the guard is stored in `TypedTreeDeferred` / `TypedMapDeferred`
    // together with the `Arc<TypedTree>` / `Arc<TypedMap>` that owns the
    // collector it borrows. Struct fields drop in declaration order, so the
    // guards are declared **first** in both structs — dropping the `Arc` first
    // could free the collector while `LocalGuard::drop` still dereferences it.
    unsafe { std::mem::transmute(guard) }
}

/// Arc-backed deferred hook owner. Keeps mutation-era `seize` guards alive
/// through replay so the raw pointers inside buffered `MutationResult` events
/// stay valid.
pub(crate) struct TypedTreeDeferred<
    K: Key + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> {
    /// Entered in `tx_shard` before the first mutation; must not be dropped
    /// until every `replay_one` has run (`typed_tree.rs:504-520`).
    ///
    /// **Declared first on purpose**: fields drop in declaration order, and the
    /// guards borrow a collector owned by the `Arc` below.
    _guards: TypedTxGuards,
    tree: Arc<TypedTree<K, T, C, H>>,
    events: MutationEvents<K, T>,
}

impl<
    K: Key + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> DeferredReplay for TypedTreeDeferred<K, T, C, H>
{
    fn event_count(&self) -> usize {
        self.events.len()
    }
    fn replay_one(&mut self, local: usize) {
        if let Some((key, mr)) = self.events.get(local) {
            self.tree.fire_hook_and_retire(key, mr);
        }
    }
}

pub(crate) struct TypedMapDeferred<
    K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> {
    /// First field for the same reason as `TypedTreeDeferred::_guards`.
    _guards: TypedTxGuards,
    map: Arc<TypedMap<K, T, C, H>>,
    events: MapMutationEvents<K, T>,
}

impl<
    K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> DeferredReplay for TypedMapDeferred<K, T, C, H>
{
    fn event_count(&self) -> usize {
        self.events.len()
    }
    fn replay_one(&mut self, local: usize) {
        if let Some((key, mr)) = self.events.get(local) {
            self.map.fire_hook_and_retire(key, mr);
        }
    }
}

impl<
    K: Key + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> TxFamily for TypedTree<K, T, C, H>
{
    type Shard<'a>
        = TypedShard<'a, K, T, C, H>
    where
        Self: 'a;
    type Deferred<'a>
        = TypedTreeDeferred<K, T, C, H>
    where
        Self: 'a;
    type ReleaseExtra = TypedTxGuards;

    fn tx_shard_count(&self) -> usize {
        self.shard_count()
    }
    fn tx_shard_of(&self, key: &[u8]) -> usize {
        self.shard_for(&K::from_bytes(key))
    }
    fn tx_make_shard(&self, sid: usize) -> Self::Shard<'_> {
        self.tx_shard(sid)
    }
    fn tx_release(
        shard: Self::Shard<'_>,
    ) -> (Vec<<Self as Family>::Owned>, WatchWakeBatch, TypedTxGuards) {
        let (events, wakes, guard) = Self::tx_release_shard(shard);
        (
            events,
            wakes,
            TypedTxGuards(vec![retain_mutation_guard(guard)]),
        )
    }
    fn tx_deferred(
        &self,
        events: Vec<<Self as Family>::Owned>,
        guards: TypedTxGuards,
    ) -> DbResult<Box<dyn DeferredReplay>> {
        Ok(Box::new(TypedTreeDeferred {
            tree: self.tx_arc()?,
            events,
            _guards: guards,
        }))
    }
}

impl<
    K: Key + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> TxRaw for TypedTree<K, T, C, H>
{
    fn raw_read(&self, shard: &mut Self::Shard<'_>, key: &[u8]) -> DbResult<Option<Vec<u8>>> {
        match shard.try_get(&K::from_bytes(key))? {
            Some(value) => {
                let mut out = Vec::new();
                self.codec().encode_to(value, &mut out)?;
                Ok(Some(out))
            }
            None => Ok(None),
        }
    }
    fn raw_contains(&self, shard: &mut Self::Shard<'_>, key: &[u8]) -> DbResult<bool> {
        shard.try_contains(&K::from_bytes(key))
    }
    fn raw_put(&self, shard: &mut Self::Shard<'_>, key: &[u8], value: &[u8]) -> DbResult<()> {
        let decoded = self.codec().decode_from(value)?;
        shard.put(&K::from_bytes(key), decoded).map(drop)
    }
    fn raw_delete(&self, shard: &mut Self::Shard<'_>, key: &[u8]) -> DbResult<()> {
        shard.delete(&K::from_bytes(key)).map(drop)
    }
    fn reserve_in_shard(
        &self,
        _shard: &mut Self::Shard<'_>,
        ops: &[PlannedOp<'_>],
    ) -> DbResult<Vec<VersionTag>> {
        let base = self.tx_reserve_gsns(ops.len());
        Ok((0..ops.len())
            .map(|i| VersionTag::Bitcask {
                gsn: base + i as u64,
            })
            .collect())
    }

    fn arm_tag(shard: &mut Self::Shard<'_>, tag: VersionTag) -> DbResult<()> {
        let VersionTag::Bitcask { gsn } = tag else {
            return Err(DbError::Internal(
                "Fixed version tag on a Bitcask participant",
            ));
        };
        shard.tx_arm_gsn(gsn);
        Ok(())
    }

    fn disarm_tag(shard: &mut Self::Shard<'_>) {
        shard.tx_clear_armed_gsn();
    }

    fn release_reservation(_shard: &mut Self::Shard<'_>, _tag: VersionTag) {
        // A GSN hole is harmless: the counter is monotonic and the replay rule
        // compares against a watermark, not against a dense sequence.
    }
    fn fence_shard(shard: &mut Self::Shard<'_>, tag: VersionTag) -> Option<u64> {
        match tag {
            VersionTag::Bitcask { gsn } => Some(shard.tx_fence(gsn)),
            VersionTag::Fixed { .. } => None,
        }
    }
    fn buffered_events(shard: &Self::Shard<'_>) -> usize {
        shard.tx_buffered_events()
    }

    fn needs_replay(
        &self,
        _shard: &mut Self::Shard<'_>,
        key: &[u8],
        _kind: OpKind,
        tag: VersionTag,
    ) -> DbResult<bool> {
        let VersionTag::Bitcask { gsn } = tag else {
            return Err(DbError::Internal(
                "Fixed version tag on a Bitcask participant",
            ));
        };
        let sid = <Self as TxFamily>::tx_shard_of(self, key);
        Ok(gsn > self.tx_durable_recovered_gsn(sid))
    }
}

impl<
    K: Key + 'static,
    T: armour_core::GetType + TxCollectionName + Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> TxParticipant for TypedTree<K, T, C, H>
{
    fn name(&self) -> &'static str {
        T::tx_collection_name()
    }
    fn backend(&self) -> TxBackend {
        TxBackend::Bitcask
    }
    fn schema_fingerprint(&self) -> u64 {
        T::TYPE.h()
    }
    fn encrypted(&self) -> bool {
        self.tx_encrypted()
    }
    fn max_entry_bytes(&self) -> Option<usize> {
        Some(self.config().write_buffer_size)
    }
    fn shard_of(&self, key: &[u8]) -> DbResult<usize> {
        Ok(<Self as TxFamily>::tx_shard_of(self, key))
    }
    fn lock_shards<'a>(&'a self, shards: &[usize]) -> DbResult<Box<dyn TxLocked + 'a>> {
        Ok(Box::new(TxHeld::lock(self, shards)))
    }
    fn tx_project(
        &self,
        key: &[u8],
        old: Option<&[u8]>,
        new: Option<&[u8]>,
    ) -> DbResult<Vec<TxDerivedWrite>> {
        self.tx_project_writes(key, old, new)
    }
    fn tx_stream_writes(
        &self,
        key: &[u8],
        old: Option<&[u8]>,
        new: Option<&[u8]>,
    ) -> DbResult<crate::txn::participant::TxStreamWrites> {
        self.tx_stream_writes(key, old, new)
    }
    fn tx_counter_deltas(
        &self,
        key: &[u8],
        old: Option<&[u8]>,
        new: Option<&[u8]>,
    ) -> DbResult<Vec<TxCounterDelta>> {
        self.tx_counter_deltas_for(key, old, new)
    }
    fn tx_claims(
        &self,
        key: &[u8],
        old: Option<&[u8]>,
        new: Option<&[u8]>,
    ) -> DbResult<Vec<ClaimOp>> {
        self.tx_claims_for(self.name(), key, old, new)
    }
    fn tx_preplan_projection_targets(&self) -> DbResult<Vec<Arc<dyn TxParticipant>>> {
        TypedTree::tx_preplan_projection_targets(self)
    }
    fn tx_preplan_stream_targets(&self) -> DbResult<Vec<Arc<dyn TxParticipant>>> {
        Ok(TypedTree::tx_preplan_stream_targets(self))
    }
    fn tx_prepare_streams(&self) -> DbResult<Vec<crate::queue::retention::ProducerPermit>> {
        TypedTree::tx_prepare_streams(self)
    }
}

impl<
    K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> TxFamily for TypedMap<K, T, C, H>
{
    type Shard<'a>
        = TypedMapShard<'a, K, T, C, H>
    where
        Self: 'a;
    type Deferred<'a>
        = TypedMapDeferred<K, T, C, H>
    where
        Self: 'a;
    type ReleaseExtra = TypedTxGuards;

    fn tx_shard_count(&self) -> usize {
        self.shard_count()
    }
    fn tx_shard_of(&self, key: &[u8]) -> usize {
        self.shard_for(&K::from_bytes(key))
    }
    fn tx_make_shard(&self, sid: usize) -> Self::Shard<'_> {
        self.tx_shard(sid)
    }
    fn tx_release(
        shard: Self::Shard<'_>,
    ) -> (Vec<<Self as Family>::Owned>, WatchWakeBatch, TypedTxGuards) {
        let (events, wakes, guard) = Self::tx_release_shard(shard);
        (
            events,
            wakes,
            TypedTxGuards(vec![retain_mutation_guard(guard)]),
        )
    }
    fn tx_deferred(
        &self,
        events: Vec<<Self as Family>::Owned>,
        guards: TypedTxGuards,
    ) -> DbResult<Box<dyn DeferredReplay>> {
        Ok(Box::new(TypedMapDeferred {
            map: self.tx_arc()?,
            events,
            _guards: guards,
        }))
    }
}

impl<
    K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> TxRaw for TypedMap<K, T, C, H>
{
    fn raw_read(&self, shard: &mut Self::Shard<'_>, key: &[u8]) -> DbResult<Option<Vec<u8>>> {
        match shard.try_get(&K::from_bytes(key))? {
            Some(value) => {
                let mut out = Vec::new();
                self.codec().encode_to(value, &mut out)?;
                Ok(Some(out))
            }
            None => Ok(None),
        }
    }
    fn raw_contains(&self, shard: &mut Self::Shard<'_>, key: &[u8]) -> DbResult<bool> {
        shard.try_contains(&K::from_bytes(key))
    }
    fn raw_put(&self, shard: &mut Self::Shard<'_>, key: &[u8], value: &[u8]) -> DbResult<()> {
        let decoded = self.codec().decode_from(value)?;
        shard.put(&K::from_bytes(key), decoded).map(drop)
    }
    fn raw_delete(&self, shard: &mut Self::Shard<'_>, key: &[u8]) -> DbResult<()> {
        shard.delete(&K::from_bytes(key)).map(drop)
    }
    fn reserve_in_shard(
        &self,
        _shard: &mut Self::Shard<'_>,
        ops: &[PlannedOp<'_>],
    ) -> DbResult<Vec<VersionTag>> {
        let base = self.tx_reserve_gsns(ops.len());
        Ok((0..ops.len())
            .map(|i| VersionTag::Bitcask {
                gsn: base + i as u64,
            })
            .collect())
    }

    fn arm_tag(shard: &mut Self::Shard<'_>, tag: VersionTag) -> DbResult<()> {
        let VersionTag::Bitcask { gsn } = tag else {
            return Err(DbError::Internal(
                "Fixed version tag on a Bitcask participant",
            ));
        };
        shard.tx_arm_gsn(gsn);
        Ok(())
    }

    fn disarm_tag(shard: &mut Self::Shard<'_>) {
        shard.tx_clear_armed_gsn();
    }

    fn release_reservation(_shard: &mut Self::Shard<'_>, _tag: VersionTag) {
        // A GSN hole is harmless: the counter is monotonic and the replay rule
        // compares against a watermark, not against a dense sequence.
    }
    fn fence_shard(shard: &mut Self::Shard<'_>, tag: VersionTag) -> Option<u64> {
        match tag {
            VersionTag::Bitcask { gsn } => Some(shard.tx_fence(gsn)),
            VersionTag::Fixed { .. } => None,
        }
    }
    fn buffered_events(shard: &Self::Shard<'_>) -> usize {
        shard.tx_buffered_events()
    }

    fn needs_replay(
        &self,
        _shard: &mut Self::Shard<'_>,
        key: &[u8],
        _kind: OpKind,
        tag: VersionTag,
    ) -> DbResult<bool> {
        let VersionTag::Bitcask { gsn } = tag else {
            return Err(DbError::Internal(
                "Fixed version tag on a Bitcask participant",
            ));
        };
        let sid = <Self as TxFamily>::tx_shard_of(self, key);
        Ok(gsn > self.tx_durable_recovered_gsn(sid))
    }
}

impl<
    K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
    T: armour_core::GetType + TxCollectionName + Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> TxParticipant for TypedMap<K, T, C, H>
{
    fn name(&self) -> &'static str {
        T::tx_collection_name()
    }
    fn backend(&self) -> TxBackend {
        TxBackend::Bitcask
    }
    fn schema_fingerprint(&self) -> u64 {
        T::TYPE.h()
    }
    fn encrypted(&self) -> bool {
        self.tx_encrypted()
    }
    fn max_entry_bytes(&self) -> Option<usize> {
        Some(self.config().write_buffer_size)
    }
    fn shard_of(&self, key: &[u8]) -> DbResult<usize> {
        Ok(<Self as TxFamily>::tx_shard_of(self, key))
    }
    fn lock_shards<'a>(&'a self, shards: &[usize]) -> DbResult<Box<dyn TxLocked + 'a>> {
        Ok(Box::new(TxHeld::lock(self, shards)))
    }
    fn tx_project(
        &self,
        key: &[u8],
        old: Option<&[u8]>,
        new: Option<&[u8]>,
    ) -> DbResult<Vec<TxDerivedWrite>> {
        self.tx_project_writes(key, old, new)
    }
    fn tx_stream_writes(
        &self,
        key: &[u8],
        old: Option<&[u8]>,
        new: Option<&[u8]>,
    ) -> DbResult<crate::txn::participant::TxStreamWrites> {
        self.tx_stream_writes(key, old, new)
    }
    fn tx_counter_deltas(
        &self,
        key: &[u8],
        old: Option<&[u8]>,
        new: Option<&[u8]>,
    ) -> DbResult<Vec<TxCounterDelta>> {
        self.tx_counter_deltas_for(key, old, new)
    }
    fn tx_claims(
        &self,
        key: &[u8],
        old: Option<&[u8]>,
        new: Option<&[u8]>,
    ) -> DbResult<Vec<ClaimOp>> {
        self.tx_claims_for(self.name(), key, old, new)
    }
    fn tx_preplan_projection_targets(&self) -> DbResult<Vec<Arc<dyn TxParticipant>>> {
        TypedMap::tx_preplan_projection_targets(self)
    }
    fn tx_preplan_stream_targets(&self) -> DbResult<Vec<Arc<dyn TxParticipant>>> {
        Ok(TypedMap::tx_preplan_stream_targets(self))
    }
    fn tx_prepare_streams(&self) -> DbResult<Vec<crate::queue::retention::ProducerPermit>> {
        TypedMap::tx_prepare_streams(self)
    }
}

impl<
    K: Key + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> TxTarget for TypedTree<K, T, C, H>
{
    type Key = K;
    type Val = T;

    fn tx_encode(&self, value: &T, out: &mut Vec<u8>) -> DbResult<()> {
        self.codec().encode_to(value, out)
    }

    fn tx_decode(&self, bytes: &[u8]) -> DbResult<T> {
        self.codec().decode_from(bytes)
    }
}

impl<
    K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
    T: Send + Sync + 'static,
    C: Codec<T> + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
> TxTarget for TypedMap<K, T, C, H>
{
    type Key = K;
    type Val = T;

    fn tx_encode(&self, value: &T, out: &mut Vec<u8>) -> DbResult<()> {
        self.codec().encode_to(value, out)
    }

    fn tx_decode(&self, bytes: &[u8]) -> DbResult<T> {
        self.codec().decode_from(bytes)
    }
}

#[cfg(test)]
#[cfg(feature = "rapira-codec")]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};

    use crate::txn::participant::{PlannedOp, TxParticipant};
    use crate::{Config, RapiraCodec, TypedTree, TypedWriteHook};

    #[derive(Clone, Debug, PartialEq, armour_core::GetType, rapira::Rapira)]
    struct Row {
        n: u32,
    }

    impl crate::txn::participant::TxCollectionName for Row {}

    #[derive(Clone, Default)]
    struct CountingHook(Arc<AtomicUsize>);
    impl TypedWriteHook<[u8; 8], Row> for CountingHook {
        const NEEDS_WRITE: bool = true;
        fn on_write(&self, _k: &[u8; 8], _old: Option<&Row>, _new: Option<&Row>) {
            self.0.fetch_add(1, Ordering::SeqCst);
        }
    }

    fn open_tree(
        dir: &std::path::Path,
        hook: CountingHook,
    ) -> Arc<TypedTree<[u8; 8], Row, RapiraCodec, CountingHook>> {
        let tree =
            Arc::new(TypedTree::open_hooked(dir, Config::test(), RapiraCodec, hook).unwrap());
        TypedTree::bind_tx_self(&tree);
        tree
    }

    #[test]
    fn tx_typed_round_trips_through_the_codec_and_defers_hooks() {
        let dir = tempfile::tempdir().unwrap();
        let fired = Arc::new(AtomicUsize::new(0));
        let tree = open_tree(dir.path(), CountingHook(fired.clone()));
        tree.put(&[1u8; 8], Row { n: 1 }).unwrap();
        fired.store(0, Ordering::SeqCst);

        let sid = TxParticipant::shard_of(&*tree, &[1u8; 8]).unwrap();
        let mut locked = TxParticipant::lock_shards(&*tree, &[sid]).unwrap();

        let stored = locked.read(&[1u8; 8]).unwrap().expect("row is present");
        assert_eq!(rapira::deserialize::<Row>(&stored).unwrap(), Row { n: 1 });

        let new_bytes = rapira::serialize(&Row { n: 2 });
        let tags = locked
            .reserve(&[PlannedOp {
                key: &[1u8; 8],
                delete: false,
            }])
            .unwrap();
        locked.apply_put(&[1u8; 8], &new_bytes, tags[0]).unwrap();
        assert_eq!(fired.load(Ordering::SeqCst), 0, "no hook under the lock");

        let mut deferred = locked
            .finish(
                &mut Vec::new(),
                &mut crate::watch::WatchWakeBatch::default(),
            )
            .unwrap();
        assert_eq!(deferred.event_count(), 1);
        deferred.replay_one(0);
        assert_eq!(fired.load(Ordering::SeqCst), 1);
        drop(deferred);

        assert_eq!(tree.get(&[1u8; 8]).map(|r| r.n), Some(2));
    }

    #[test]
    fn tx_typed_rejects_a_value_the_codec_cannot_decode() {
        let dir = tempfile::tempdir().unwrap();
        let tree = open_tree(dir.path(), CountingHook::default());
        let sid = TxParticipant::shard_of(&*tree, &[5u8; 8]).unwrap();
        let mut locked = TxParticipant::lock_shards(&*tree, &[sid]).unwrap();
        let tags = locked
            .reserve(&[PlannedOp {
                key: &[5u8; 8],
                delete: false,
            }])
            .unwrap();
        let err = locked
            .apply_put(&[5u8; 8], &[0xFF; 3], tags[0])
            .unwrap_err();
        assert!(
            matches!(
                err,
                crate::DbError::CorruptedEntry { .. } | crate::DbError::Client(_)
            ),
            "{err}"
        );
        locked.abort();
        assert!(tree.get(&[5u8; 8]).is_none());
    }
}