armdb 0.7.2

use std::mem::size_of;
use std::num::NonZeroU32;
use std::sync::Arc;

#[cfg(feature = "rpc")]
use std::ops::Bound;

#[cfg(feature = "rpc")]
use armour_rpc::protocol::UpsertKey;

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

#[cfg(feature = "rpc")]
use crate::armour::handler::{
    KeyBytes, RpcHandler, ValueBytes, apply_limit, bound_to_key_bound, check_key_len, resolve_key,
};
use crate::armour::handler::{bound_to_key_bound_owned, iteration_direction};
use crate::armour::inspection::{
    InspectionCollection, InspectionEntry, InspectionError, InspectionHandler, InspectionOperation,
    InspectionPage, InspectionSchema, NormalizedBounds, ScanCapability, collect_page,
};
use crate::codec::Codec;
#[cfg(feature = "rpc")]
use crate::error::DbError;
use crate::error::DbResult;
use crate::hook::TypedWriteHook;
use crate::lsm::{LsmTree, map_fjall};
use crate::{CollectionMeta, Key};

/// Handler for [`LsmTree`] — fjall-backed typed values (RPC when `rpc` is enabled).
pub struct LsmTreeHandler<K, T, C, H>
where
    K: Key + Ord,
    T: CollectionMeta<SelfId = K> + Send + Sync,
    C: Codec<T> + Clone,
    H: TypedWriteHook<K, T>,
{
    pub name: String,
    pub typ_hash: u64,
    pub ty: armour_core::Typ,
    pub key_scheme: armour_core::KeyScheme,
    pub version: u16,
    pub tree: Arc<LsmTree<K, T, C, H>>,
    pub codec: Arc<C>,
    pub seq: Arc<crate::armour::seq::SeqGen>,
}

// ─── impl RpcHandler ─────────────────────────────────────────────────────────

#[cfg(feature = "rpc")]
impl<K, T, C, H> RpcHandler for LsmTreeHandler<K, T, C, H>
where
    K: Key + Ord + Send + Sync,
    T: CollectionMeta<SelfId = K> + Send + Sync,
    C: Codec<T> + Clone,
    H: TypedWriteHook<K, T>,
{
    fn name(&self) -> &str {
        &self.name
    }

    fn info(&self) -> (u64, u16) {
        (self.typ_hash, self.version)
    }

    fn schema(&self) -> armour_rpc::SchemaResponse {
        armour_rpc::SchemaResponse {
            name: self.name.clone(),
            version: self.version,
            typ: self.ty,
            key_scheme: self.key_scheme,
        }
    }

    fn scannable(&self) -> bool {
        true
    }

    fn get(&self, key: &[u8]) -> DbResult<Option<ValueBytes>> {
        let key: K = check_key_len(key)?;
        // `get_raw` propagates a fjall I/O error as `Err`, same as `try_get` —
        // a client must never see a transiently-unreadable key as absent, or it
        // may clobber live data. (Memory-backed Bitcask families rarely see I/O
        // errors here; fjall's disk-backed keyspace does.)
        //
        // Storage and wire share the same codec, so `try_get`'s decode-then-
        // re-encode just reproduces the bytes already on disk. Decode here only
        // to validate — a corrupt row must still surface as `Err` rather than
        // being handed to the client (0de5153/e71f43e) — then return the
        // original raw bytes instead of paying for a redundant encode.
        match self.tree.get_raw(&key)? {
            Some(raw) => {
                self.codec.decode_from(raw.as_ref())?;
                Ok(Some(raw.to_vec()))
            }
            None => Ok(None),
        }
    }

    fn contains(&self, key: &[u8]) -> DbResult<bool> {
        let k = check_key_len::<K>(key)?;
        Ok(self.tree.contains(&k))
    }

    fn entry_len(&self, key: &[u8]) -> DbResult<Option<u32>> {
        let k = check_key_len::<K>(key)?;
        Ok(self.tree.entry_len(&k))
    }

    fn first(&self) -> DbResult<Option<(KeyBytes, ValueBytes)>> {
        match self.tree.first() {
            Some((k, v)) => Ok(Some((k.as_bytes().to_vec(), self.codec.encode(&v)?))),
            None => Ok(None),
        }
    }

    fn last(&self) -> DbResult<Option<(KeyBytes, ValueBytes)>> {
        match self.tree.last() {
            Some((k, v)) => Ok(Some((k.as_bytes().to_vec(), self.codec.encode(&v)?))),
            None => Ok(None),
        }
    }

    fn range(
        &self,
        start: Bound<KeyBytes>,
        end: Bound<KeyBytes>,
        limit: u32,
    ) -> DbResult<Vec<(KeyBytes, ValueBytes)>> {
        let sb = bound_to_key_bound::<K>(&start)?;
        let eb = bound_to_key_bound::<K>(&end)?;
        let iter = self.tree.range_bounds(sb.as_ref(), eb.as_ref());
        let mut result = Vec::new();
        for (k, v) in apply_limit(iter, limit) {
            result.push((k.as_bytes().to_vec(), self.codec.encode(&v)?));
        }
        Ok(result)
    }

    fn range_keys(
        &self,
        start: Bound<KeyBytes>,
        end: Bound<KeyBytes>,
        limit: u32,
    ) -> DbResult<Vec<ValueBytes>> {
        let sb = bound_to_key_bound::<K>(&start)?;
        let eb = bound_to_key_bound::<K>(&end)?;
        let iter = self.tree.range_bounds(sb.as_ref(), eb.as_ref());
        Ok(apply_limit(iter, limit)
            .map(|(k, _)| k.as_bytes().to_vec())
            .collect())
    }

    fn upsert(
        &self,
        key: UpsertKey,
        flag: Option<bool>,
        value: ValueBytes,
    ) -> DbResult<ValueBytes> {
        let (typed_key, key_bytes) = resolve_key::<K>(key, &self.seq, &self.name)?;
        // Pre-decode to validate value bytes before write.
        let decoded = self.codec.decode_from(&value)?;
        match flag {
            None => {
                self.tree.put(&typed_key, &decoded)?;
                Ok(key_bytes)
            }
            Some(false) => self
                .tree
                .atomic(&typed_key, |shard| {
                    // `try_contains`, not `contains`: an I/O error must not read
                    // as "absent" here, or insert-if-absent would overwrite the
                    // live row it failed to see. Matches the sibling handlers.
                    if shard.try_contains(&typed_key)? {
                        return Err(DbError::KeyExists);
                    }
                    shard.put(&typed_key, &decoded)?;
                    Ok(())
                })
                .map(|_| key_bytes),
            Some(true) => self
                .tree
                .atomic(&typed_key, |shard| {
                    if !shard.try_contains(&typed_key)? {
                        return Err(DbError::KeyNotFound);
                    }
                    shard.put(&typed_key, &decoded)?;
                    Ok(())
                })
                .map(|_| key_bytes),
        }
    }

    fn remove(&self, key: &[u8], soft: bool) -> DbResult<()> {
        if soft {
            return Err(DbError::NotImplemented);
        }
        let k = check_key_len::<K>(key)?;
        self.tree.delete(&k).map(|_| ())
    }

    fn take(&self, key: &[u8], soft: bool) -> DbResult<Option<ValueBytes>> {
        if soft {
            return Err(DbError::NotImplemented);
        }
        let k = check_key_len::<K>(key)?;
        let codec = Arc::clone(&self.codec);
        self.tree.atomic(&k, |shard| {
            // Strict read, as in every sibling handler: a corrupt or
            // transiently-unreadable row must not report as `Ok(None)` — that
            // tells the client the key is absent while silently leaving it in
            // place, so a subsequent write clobbers it.
            match shard.try_get(&k)? {
                Some(v) => {
                    let buf = codec.encode(&v)?;
                    shard.delete(&k)?;
                    Ok(Some(buf))
                }
                None => Ok(None),
            }
        })
    }

    fn count(&self, exact: bool) -> DbResult<u64> {
        if exact {
            // Key-only scan via fjall's `Keyspace::len()` — it counts
            // `guard.key()` without decoding values, unlike `self.tree.iter()`
            // which decodes every row. `iter()` also *skips* a row that fails
            // to decode (logged at `warn`), so it under-counts a corrupt
            // collection; this key-only count makes `exact` actually exact.
            Ok(map_fjall(self.tree.keyspace().len())? as u64)
        } else {
            Ok(self.tree.len() as u64)
        }
    }

    fn apply_batch(&self, items: Vec<(KeyBytes, Option<ValueBytes>)>) -> DbResult<()> {
        // Decode every item in a single pass — short-circuits on first error,
        // so no partial writes occur on validation failure.
        let decoded_items: Vec<(K, Option<T>)> = items
            .into_iter()
            .map(|(key, val)| {
                if key.len() != size_of::<K>() {
                    return Err(DbError::Client("invalid key length in batch"));
                }
                let decoded_val = match val {
                    Some(v) => Some(
                        self.codec
                            .decode_from(&v)
                            .map_err(|_| DbError::Client("invalid value bytes in batch"))?,
                    ),
                    None => None,
                };
                Ok((K::from_bytes(&key), decoded_val))
            })
            .collect::<DbResult<_>>()?;
        // Apply from the decoded vec — no second decode.
        for (k, val) in decoded_items {
            match val {
                Some(decoded) => {
                    self.tree.put(&k, &decoded)?;
                }
                None => {
                    self.tree.delete(&k)?;
                }
            }
        }
        Ok(())
    }
}

// ─── InspectionHandler ───────────────────────────────────────────────────────

impl<K, T, C, H> InspectionHandler for LsmTreeHandler<K, T, C, H>
where
    K: Key + Ord + Send + Sync,
    T: CollectionMeta<SelfId = K> + Send + Sync,
    C: Codec<T> + Clone,
    H: TypedWriteHook<K, T>,
{
    fn metadata(&self) -> Result<InspectionCollection, InspectionError> {
        Ok(InspectionCollection {
            name: self.name.clone(),
            schema: InspectionSchema {
                version: self.version,
                typ: self.ty,
                key_scheme: self.key_scheme,
                key_size: size_of::<K>() as u32,
                encoding: C::INSPECTION_ENCODING,
            },
            storage: StorageClass {
                model: ValueModel::Lsm,
                shape: CollectionShape::Tree,
                backend: StorageBackend::Lsm,
            },
            count: Some(self.tree.len() as u64),
            scan: ScanCapability::Ordered,
            direction: Some(iteration_direction(self.tree.config().reversed)),
        })
    }

    fn get_raw(&self, key: &[u8]) -> DbResult<Option<Vec<u8>>> {
        let key = K::from_bytes(key);
        match self.tree.get_raw(&key)? {
            Some(raw) => {
                self.codec.decode_from(raw.as_ref())?;
                Ok(Some(raw.to_vec()))
            }
            None => Ok(None),
        }
    }

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

    fn ordered_page(
        &self,
        bounds: &NormalizedBounds,
        limit: NonZeroU32,
    ) -> DbResult<InspectionPage> {
        if bounds.empty {
            return Ok(InspectionPage {
                entries: vec![],
                next_key: None,
            });
        }
        let start = bound_to_key_bound_owned::<K>(&bounds.start)?;
        let end = bound_to_key_bound_owned::<K>(&bounds.end)?;
        let codec = Arc::clone(&self.codec);
        collect_page(
            self.tree.try_range_bounds(start.as_ref(), end.as_ref()),
            limit,
            move |k, v| {
                let mut bytes = Vec::new();
                codec.encode_to(&v, &mut bytes)?;
                Ok(InspectionEntry {
                    key: k.as_bytes().to_vec(),
                    value: bytes,
                })
            },
        )
    }

    fn unordered_snapshot(
        &self,
        _max_entries: NonZeroU32,
    ) -> Result<Vec<InspectionEntry>, InspectionError> {
        Err(InspectionError::UnsupportedOperation {
            name: self.name.clone(),
            operation: InspectionOperation::SnapshotUnordered,
        })
    }
}

#[cfg(test)]
mod tests {
    use std::ops::Bound;
    use std::sync::Arc;

    use super::*;
    use crate::armour::inspection::NormalizedBounds;
    use crate::armour::seq::SeqGen;
    use crate::error::DbError;
    use crate::lsm::config::LsmConfig;
    use crate::{NoHook, RapiraCodec};
    use armour_core::{Fuid, GetType};
    use rapira::Rapira;

    armour_core::const_hasher_or!(doc_id, "DOC", "dev-key-doc", brand = "DOC");
    type DocId = Fuid<doc_id::Hasher>;

    #[derive(Clone, Debug, PartialEq, GetType, Rapira)]
    struct Doc {
        name: String,
    }
    impl crate::CollectionMeta for Doc {
        type SelfId = DocId;
        const NAME: &'static str = "docs";
    }

    fn k(n: u64) -> DocId {
        DocId::from_be_bytes(n.to_be_bytes())
    }

    fn handler() -> (
        tempfile::TempDir,
        tempfile::TempDir,
        LsmTreeHandler<DocId, Doc, RapiraCodec, NoHook>,
    ) {
        let dir = tempfile::tempdir().unwrap();
        let db = fjall::Database::builder(dir.path()).open().unwrap();
        let tree = Arc::new(LsmTree::open(db, LsmConfig::default(), RapiraCodec, NoHook).unwrap());
        // Separate root: a `FixedMap` directory inside the fjall root would be a
        // foreign entry in a keyspace directory (a recovery hazard).
        let seq_dir = tempfile::tempdir().unwrap();
        let seq = SeqGen::open(seq_dir.path().join("seq")).unwrap();
        let h = LsmTreeHandler {
            name: "docs".to_string(),
            typ_hash: Doc::TYPE.h(),
            ty: Doc::TYPE,
            key_scheme: DocId::KEY_SCHEME,
            version: 0,
            tree,
            codec: Arc::new(RapiraCodec),
            seq,
        };
        (dir, seq_dir, h)
    }

    fn full_bounds() -> NormalizedBounds {
        NormalizedBounds {
            start: Bound::Unbounded,
            end: Bound::Unbounded,
            empty: false,
        }
    }

    #[test]
    fn inspection_lsm_get_propagates_decode_error_and_keeps_row() {
        let (_dir, _seq_dir, h) = handler();
        let key = k(1);
        h.tree
            .keyspace()
            .insert(key.as_bytes(), b"not-a-doc")
            .unwrap();
        let err = h.get_raw(key.as_bytes()).unwrap_err();
        assert!(matches!(err, DbError::CorruptedEntry { .. }), "{err:?}");
        assert!(
            h.tree.keyspace().get(key.as_bytes()).unwrap().is_some(),
            "strict inspection get must leave the row untouched"
        );
    }

    #[test]
    fn inspection_lsm_page_fails_on_poison_without_partial_page() {
        let (_dir, _seq_dir, h) = handler();
        h.tree.put(&k(1), &Doc { name: "a".into() }).unwrap();
        h.tree
            .keyspace()
            .insert(k(2).as_bytes(), b"not-a-doc")
            .unwrap();
        h.tree.put(&k(3), &Doc { name: "c".into() }).unwrap();
        let err = h
            .ordered_page(&full_bounds(), NonZeroU32::new(10).unwrap())
            .unwrap_err();
        assert!(matches!(err, DbError::CorruptedEntry { .. }), "{err:?}");
    }

    #[test]
    fn inspection_lsm_count_includes_undecodable_row() {
        let (_dir, _seq_dir, h) = handler();
        h.tree.put(&k(1), &Doc { name: "ok".into() }).unwrap();
        h.tree
            .keyspace()
            .insert(k(2).as_bytes(), b"not-a-doc")
            .unwrap();
        assert_eq!(h.exact_count().unwrap(), 2);
        assert_eq!(
            map_fjall(h.tree.keyspace().len()).unwrap() as u64,
            h.exact_count().unwrap()
        );
        assert_eq!(h.tree.iter().count(), 1, "lossy iter skips corrupt row");
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_get_propagates_decode_error() {
        let (_dir, _seq_dir, h) = handler();
        let key = k(1);
        // Inject a value that cannot decode as `Doc` straight into the keyspace.
        h.tree
            .keyspace()
            .insert(key.as_bytes(), b"not-a-doc")
            .unwrap();
        // `get` must surface the corrupt value as an error, not mask it as
        // `Ok(None)` (which would let a client treat the key as absent and
        // clobber live data on a transient decode/IO failure).
        let err = RpcHandler::get(&h, key.as_bytes()).unwrap_err();
        assert!(matches!(err, DbError::CorruptedEntry { .. }), "{err:?}");
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_get_returns_bytes_identical_to_stored_encoding() {
        let (_dir, _seq_dir, h) = handler();
        let key = k(1);
        let doc = Doc {
            name: "hello".into(),
        };
        let want = h.codec.encode(&doc).unwrap();
        h.tree.put(&key, &doc).unwrap();
        // `get` now returns the raw stored bytes instead of decode-then-
        // re-encode; the wire contract (bytes identical to the on-disk
        // encoding) must be unchanged.
        let got = RpcHandler::get(&h, key.as_bytes()).unwrap().unwrap();
        assert_eq!(got, want);
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_take_propagates_decode_error_and_keeps_row() {
        let (_dir, _seq_dir, h) = handler();
        let key = k(1);
        h.tree
            .keyspace()
            .insert(key.as_bytes(), b"not-a-doc")
            .unwrap();
        // `take` must not report a corrupt row as `Ok(None)`: that tells the
        // client the key is absent while silently leaving it on disk, so the
        // client's follow-up write clobbers it.
        let err = RpcHandler::take(&h, key.as_bytes(), false).unwrap_err();
        assert!(matches!(err, DbError::CorruptedEntry { .. }), "{err:?}");
        assert!(
            h.tree.try_contains(&key).unwrap(),
            "a failed take must leave the row in place"
        );
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_upsert_insert_only_does_not_clobber_corrupt_row() {
        let (_dir, _seq_dir, h) = handler();
        let key = k(1);
        h.tree
            .keyspace()
            .insert(key.as_bytes(), b"not-a-doc")
            .unwrap();
        // insert-if-absent over an existing (if undecodable) key must report
        // KeyExists from the strict `try_contains`, never silently overwrite.
        let mut buf = Vec::new();
        h.codec
            .encode_to(&Doc { name: "new".into() }, &mut buf)
            .unwrap();
        let err = RpcHandler::upsert(
            &h,
            UpsertKey::Provided(key.as_bytes().to_vec()),
            Some(false),
            buf,
        )
        .unwrap_err();
        assert!(matches!(err, DbError::KeyExists), "{err:?}");
        assert_eq!(
            h.tree.keyspace().get(key.as_bytes()).unwrap().as_deref(),
            Some(&b"not-a-doc"[..]),
            "insert-only must not have overwritten the existing row"
        );
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_count_exact_includes_corrupt_rows() {
        let (_dir, _seq_dir, h) = handler();
        h.tree.put(&k(1), &Doc { name: "ok".into() }).unwrap();
        // Inject a row that cannot decode as `Doc`, straight into the keyspace.
        h.tree
            .keyspace()
            .insert(k(2).as_bytes(), b"not-a-doc")
            .unwrap();
        // `count(true)` must count every row, including one that fails to
        // decode — `self.tree.iter()` silently skips it, which would
        // under-count a corrupt collection.
        assert_eq!(RpcHandler::count(&h, true).unwrap(), 2);
    }
}