armdb 0.14.2

use std::mem::size_of;
use std::ops::Bound;

use armour_core::{CollectionShape, StorageBackend, StorageClass, Typ, ValueModel};

use crate::access::{Access, FallibleRowFn, KeyFn, RowFn, Wire};
use crate::armour::inspection::ScanCapability;
use crate::codec::{Codec, InspectionEncoding};
use crate::error::{DbError, DbResult};
use crate::key::Key;
use crate::lsm::{LsmShard, LsmTree, map_fjall};
use crate::{CollectionMeta, TypedWriteHook};

impl<K, T, C, H> LsmTree<K, T, C, H>
where
    K: Key + Ord + 'static,
    T: CollectionMeta<SelfId = K> + Send + Sync + 'static,
    C: Codec<T> + Clone + Send + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
{
    /// Shared traversal behind `Access::scan_all` and `Access::undecodable_rows`.
    ///
    /// `warn` is off for the counting pass: it reports the same rows as a
    /// `UndecodableRows` finding, and warning about each of them twice per
    /// validation run is noise.
    fn walk_rows(&self, warn: bool, f: FallibleRowFn<'_, K, T>) {
        let codec = self.codec();
        for guard in self.keyspace().iter() {
            let (raw_key, raw_value) = match map_fjall(guard.into_inner()) {
                Ok(pair) => pair,
                Err(error) => {
                    if warn {
                        tracing::warn!(
                            collection = T::NAME,
                            error = %error,
                            "schema scan: fjall iter error in LsmTree, skipping entry — relations not checked"
                        );
                    }
                    if f(Err(error)).is_break() {
                        break;
                    }
                    continue;
                }
            };
            if raw_key.len() != size_of::<K>() {
                if warn {
                    tracing::warn!(
                        collection = T::NAME,
                        len = raw_key.len(),
                        "schema scan: bad key length in LsmTree, skipping entry — relations not checked"
                    );
                }
                if f(Err(DbError::CorruptedEntry { offset: 0 })).is_break() {
                    break;
                }
                continue;
            }
            let key = K::from_bytes(raw_key.as_ref());
            let value = match codec.decode_from(raw_value.as_ref()) {
                Ok(value) => value,
                Err(error) => {
                    if warn {
                        tracing::warn!(
                            collection = T::NAME,
                            "schema scan: decode error in LsmTree, skipping entry — relations not checked"
                        );
                    }
                    if f(Err(error)).is_break() {
                        break;
                    }
                    continue;
                }
            };
            if f(Ok((&key, &value))).is_break() {
                break;
            }
        }
    }
}

impl<K, T, C, H> Access for LsmTree<K, T, C, H>
where
    K: Key + Ord + 'static,
    T: CollectionMeta<SelfId = K> + Send + Sync + 'static,
    C: Codec<T> + Clone + Send + Sync + 'static,
    H: TypedWriteHook<K, T> + 'static,
{
    type Key = K;
    type Val = T;
    type Shard<'a>
        = LsmShard<'a, K, T, C, H>
    where
        Self: 'a;

    const STORAGE: StorageClass = StorageClass {
        model: ValueModel::Lsm,
        shape: CollectionShape::Tree,
        backend: StorageBackend::Lsm,
    };
    const RAW_WIRE_BYTES: bool = true;

    fn encoding(&self, _ty: &Typ) -> InspectionEncoding {
        C::INSPECTION_ENCODING
    }

    fn approx_len(&self) -> u64 {
        self.len() as u64
    }

    fn rpc_exact_count(&self) -> DbResult<u64> {
        Ok(map_fjall(self.keyspace().len())? as u64)
    }

    fn inspection_exact_count(&self) -> DbResult<u64> {
        Ok(map_fjall(self.keyspace().len())? as u64)
    }

    fn scan_capability(&self) -> ScanCapability {
        ScanCapability::Ordered
    }

    fn reversed(&self) -> bool {
        self.config().reversed
    }

    fn with_value<R>(&self, key: &K, f: impl FnOnce(&T) -> R) -> DbResult<Option<R>> {
        Ok(self.try_get(key)?.map(|value| f(&value)))
    }

    /// Facade policy, deliberately **lenient**: the live RPC handler answers
    /// from `LsmTree::contains`, where a fjall read failure reads as `false`
    /// (`lsm/handler.rs`). The strict check stays on
    /// [`HeldOps::try_contains`] — under a held shard lock an unreadable row
    /// must not be mistaken for an absent one, or a conditional upsert
    /// overwrites live data.
    fn contains(&self, key: &K) -> DbResult<bool> {
        Ok(LsmTree::contains(self, key))
    }

    fn stored_bytes(&self, key: &K) -> DbResult<Option<Vec<u8>>> {
        Ok(self.get_raw(key)?.map(|slice| slice.to_vec()))
    }

    fn entry_len<W: Wire<T>>(&self, key: &K, _wire: &W) -> DbResult<Option<u32>> {
        Ok(self.entry_len(key))
    }

    fn with_first<R>(&self, f: impl FnOnce(&K, &T) -> R) -> DbResult<Option<R>> {
        Ok(self.first().map(|(key, value)| f(&key, &value)))
    }

    fn with_last<R>(&self, f: impl FnOnce(&K, &T) -> R) -> DbResult<Option<R>> {
        Ok(self.last().map(|(key, value)| f(&key, &value)))
    }

    fn scan_range(&self, start: Bound<&K>, end: Bound<&K>, f: RowFn<'_, K, T>) -> DbResult<()> {
        for (key, value) in self.range_bounds(start, end) {
            if f(&key, &value)?.is_break() {
                break;
            }
        }
        Ok(())
    }

    /// Ordered single-structure scan — `batch` is a Map-only knob.
    fn try_scan_range(
        &self,
        start: Bound<&K>,
        end: Bound<&K>,
        _batch: usize,
        f: RowFn<'_, K, T>,
    ) -> DbResult<()> {
        for row in self.try_range_bounds(start, end) {
            let (key, value) = row?;
            if f(&key, &value)?.is_break() {
                break;
            }
        }
        Ok(())
    }

    fn scan_keys(&self, start: Bound<&K>, end: Bound<&K>, f: KeyFn<'_, K>) -> DbResult<()> {
        for (key, _) in self.range_bounds(start, end) {
            if f(&key)?.is_break() {
                break;
            }
        }
        Ok(())
    }

    fn scan_all(&self, f: FallibleRowFn<'_, K, T>) {
        self.walk_rows(true, f);
    }

    fn undecodable_rows(&self) -> u64 {
        let mut n = 0u64;
        self.walk_rows(false, &mut |row| {
            if row.is_err() {
                n += 1;
            }
            std::ops::ControlFlow::Continue(())
        });
        n
    }

    fn put(&self, key: &K, value: T) -> DbResult<()> {
        self.put(key, &value)
    }

    fn put_ref(&self, key: &K, value: &T) -> DbResult<()> {
        self.put(key, value)
    }

    fn delete(&self, key: &K) -> DbResult<()> {
        self.delete(key).map(drop)
    }

    fn atomic_at<R>(
        &self,
        key: &K,
        f: impl FnOnce(&mut Self::Shard<'_>) -> DbResult<R>,
    ) -> DbResult<R> {
        self.atomic(key, f)
    }
}

#[cfg(test)]
mod tests {
    use std::ops::{Bound, ControlFlow};

    use armour_core::{Fuid, GetRefs, GetType};
    use rapira::Rapira;

    use crate::access::{Access, CodecWire};
    use crate::codec::Codec;
    use crate::held_ops::HeldOps;
    use crate::lsm::{LsmConfig, LsmTree};
    use crate::{CollectionMeta, Key, NoHook, RapiraCodec};

    armour_core::const_hasher_or!(
        access_lsm_id,
        "ACCESS_LSM",
        "dev-key-access-lsm",
        brand = "ACCESS_LSM"
    );
    type RowId = Fuid<access_lsm_id::Hasher>;

    #[derive(Clone, Debug, PartialEq, GetType, GetRefs, Rapira)]
    struct Row {
        value: u64,
    }

    impl CollectionMeta for Row {
        type SelfId = RowId;
        const NAME: &'static str = "access_lsm_rows";
    }

    fn key(value: u64) -> RowId {
        RowId::from_be_bytes(value.to_be_bytes())
    }

    fn tree() -> (tempfile::TempDir, LsmTree<RowId, Row, RapiraCodec, NoHook>) {
        let dir = tempfile::tempdir().expect("tmp");
        let db = fjall::Database::builder(dir.path())
            .open()
            .expect("fjall db");
        let tree = LsmTree::open(db, LsmConfig::default(), RapiraCodec, NoHook).expect("lsm tree");
        (dir, tree)
    }

    #[test]
    fn lsm_access_stored_bytes_and_native_entry_len() {
        let (_dir, tree) = tree();
        let first = Row { value: 11 };
        let second = Row { value: 22 };
        tree.put(&key(1), &first).expect("put first");
        tree.put(&key(2), &second).expect("put second");
        let raw = RapiraCodec.encode(&first).expect("encode");

        assert!(<LsmTree<RowId, Row, RapiraCodec, NoHook> as Access>::RAW_WIRE_BYTES);
        assert_eq!(
            Access::stored_bytes(&tree, &key(1)).expect("stored bytes"),
            Some(raw.clone())
        );
        assert_eq!(
            Access::entry_len(&tree, &key(1), &CodecWire(RapiraCodec)).expect("entry len"),
            Some(raw.len() as u32)
        );
        assert_eq!(Access::rpc_exact_count(&tree).expect("RPC count"), 2);
        assert_eq!(
            Access::inspection_exact_count(&tree).expect("inspection count"),
            2
        );
        assert_eq!(Access::approx_len(&tree), 2);

        Access::atomic_at(&tree, &key(3), |shard| {
            HeldOps::put(shard, &key(3), Row { value: 33 })?;
            assert!(HeldOps::try_contains(shard, &key(3))?);
            assert_eq!(
                HeldOps::with_value(shard, &key(3), |value| value.value)?,
                Some(33)
            );
            HeldOps::delete(shard, &key(3))
        })
        .expect("held operations");
        assert!(!tree.contains(&key(3)));
    }

    #[test]
    fn lsm_scan_all_reports_undecodable_rows() {
        let (_dir, tree) = tree();
        tree.put(&key(1), &Row { value: 11 }).expect("put good row");
        tree.keyspace()
            .insert(key(2).as_bytes(), b"corrupt")
            .expect("inject corrupt row");

        let mut good = 0;
        let mut bad = 0;
        Access::scan_all(&tree, &mut |row| {
            match row {
                Ok((_key, value)) => {
                    assert_eq!(value.value, 11);
                    good += 1;
                }
                Err(_) => bad += 1,
            }
            ControlFlow::Continue(())
        });
        assert_eq!((good, bad), (1, 1));

        let mut lenient = 0;
        Access::scan_range(
            &tree,
            Bound::Unbounded,
            Bound::Unbounded,
            &mut |_key, _value| {
                lenient += 1;
                Ok(ControlFlow::Continue(()))
            },
        )
        .expect("lenient range scan");
        assert_eq!(lenient, 1);

        // The counting pass sees the same corrupt row as `scan_all`.
        assert_eq!(Access::undecodable_rows(&tree), 1);
    }

    /// The facade read policy is the RPC handler's (`LsmTree::contains`, where a
    /// fjall read failure reads as absent); the strict check lives on the held
    /// shard, where a conditional write must never mistake unreadable for
    /// absent. Both are exercised here on healthy data — the error paths
    /// themselves need a fjall read failure, which the crate has no fault
    /// injection for (`test_faults` covers the Bitcask engine only).
    #[test]
    fn lsm_facade_contains_is_lenient_while_held_shard_is_strict() {
        let (_dir, tree) = tree();
        tree.put(&key(1), &Row { value: 11 }).expect("put");

        assert!(Access::contains(&tree, &key(1)).expect("facade contains"));
        assert!(!Access::contains(&tree, &key(2)).expect("facade contains"));
        assert_eq!(
            Access::contains(&tree, &key(1)).ok(),
            Some(LsmTree::contains(&tree, &key(1))),
            "facade must answer from the lenient handler method"
        );

        Access::atomic_at(&tree, &key(1), |shard| {
            assert!(HeldOps::try_contains(shard, &key(1))?);
            assert!(!HeldOps::try_contains(shard, &key(2))?);
            Ok(())
        })
        .expect("held shard");
    }
}