armdb 0.14.2

sharded bitcask key-value storage optimized for NVMe
Documentation
//! RPC handler tests for [`LsmTree`](crate::lsm::LsmTree).
//!
//! The bytes-level bridge itself lives in the generic `Adapter`
//! (`crate::access::Adapter`), which implements both `InspectionHandler` and
//! `RpcHandler` for every collection family, including LSM. This module keeps
//! only the tests that exercise that dispatch path against an `LsmTree`.

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

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

    use crate::access::{Adapter, CodecWire};
    #[cfg(feature = "rpc")]
    use crate::armour::handler::RpcHandler;
    use crate::armour::inspection::{InspectionHandler, NormalizedBounds};
    use crate::armour::seq::SeqGen;
    #[cfg(feature = "rpc")]
    use crate::codec::Codec;
    use crate::error::DbError;
    use crate::lsm::config::LsmConfig;
    use crate::lsm::{LsmTree, map_fjall};
    use crate::tree_meta::TreeMeta;
    use crate::{Key, NoHook, RapiraCodec};
    #[cfg(feature = "rpc")]
    use armour_rpc::protocol::WriteKey;

    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 open_lsm_tree() -> (
        tempfile::TempDir,
        tempfile::TempDir,
        Arc<LsmTree<DocId, Doc, RapiraCodec, NoHook>>,
        Arc<SeqGen>,
    ) {
        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();
        (dir, seq_dir, tree, seq)
    }

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

    /// Builds the same `Adapter` `register_lsm_tree_handler` publishes as the
    /// live Inspection handler (see `armour/db.rs`), so these tests exercise
    /// the production dispatch path rather than the deleted direct impl.
    fn adapter(
        tree: &Arc<LsmTree<DocId, Doc, RapiraCodec, NoHook>>,
        seq: &Arc<SeqGen>,
    ) -> Adapter<LsmTree<DocId, Doc, RapiraCodec, NoHook>, CodecWire<RapiraCodec>> {
        Adapter::new(
            &TreeMeta::of::<Doc>(),
            Arc::clone(tree),
            CodecWire(RapiraCodec),
            Arc::clone(seq),
        )
    }

    #[test]
    fn inspection_lsm_get_propagates_decode_error_and_keeps_row() {
        let (_dir, _seq_dir, tree, seq) = open_lsm_tree();
        let key = k(1);
        tree.keyspace()
            .insert(key.as_bytes(), b"not-a-doc")
            .unwrap();
        let err = adapter(&tree, &seq).get_raw(key.as_bytes()).unwrap_err();
        assert!(matches!(err, DbError::CorruptedEntry { .. }), "{err:?}");
        assert!(
            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, tree, seq) = open_lsm_tree();
        tree.put(&k(1), &Doc { name: "a".into() }).unwrap();
        tree.keyspace()
            .insert(k(2).as_bytes(), b"not-a-doc")
            .unwrap();
        tree.put(&k(3), &Doc { name: "c".into() }).unwrap();
        let err = adapter(&tree, &seq)
            .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, tree, seq) = open_lsm_tree();
        tree.put(&k(1), &Doc { name: "ok".into() }).unwrap();
        tree.keyspace()
            .insert(k(2).as_bytes(), b"not-a-doc")
            .unwrap();
        assert_eq!(adapter(&tree, &seq).exact_count().unwrap(), 2);
        assert_eq!(
            map_fjall(tree.keyspace().len()).unwrap() as u64,
            adapter(&tree, &seq).exact_count().unwrap()
        );
        assert_eq!(tree.iter().count(), 1, "lossy iter skips corrupt row");
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_get_propagates_decode_error() {
        let (_dir, _seq_dir, tree, seq) = open_lsm_tree();
        let key = k(1);
        // Inject a value that cannot decode as `Doc` straight into the keyspace.
        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 mut out = Vec::new();
        let err =
            RpcHandler::get_into(&adapter(&tree, &seq), key.as_bytes(), &mut out).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, tree, seq) = open_lsm_tree();
        let key = k(1);
        let doc = Doc {
            name: "hello".into(),
        };
        let want = RapiraCodec.encode(&doc).unwrap();
        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 mut got = Vec::new();
        assert!(RpcHandler::get_into(&adapter(&tree, &seq), key.as_bytes(), &mut got).unwrap());
        assert_eq!(got, want);
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_take_propagates_decode_error_and_keeps_row() {
        let (_dir, _seq_dir, tree, seq) = open_lsm_tree();
        let key = k(1);
        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 mut out = Vec::new();
        let err = RpcHandler::take_into(&adapter(&tree, &seq), key.as_bytes(), false, &mut out)
            .unwrap_err();
        assert!(matches!(err, DbError::CorruptedEntry { .. }), "{err:?}");
        assert!(
            tree.try_contains(&key).unwrap(),
            "a failed take must leave the row in place"
        );
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_write_insert_only_does_not_clobber_corrupt_row() {
        let (_dir, _seq_dir, tree, seq) = open_lsm_tree();
        let key = k(1);
        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();
        RapiraCodec
            .encode_to(&Doc { name: "new".into() }, &mut buf)
            .unwrap();
        let err = RpcHandler::write(
            &adapter(&tree, &seq),
            WriteKey::Provided(key.as_bytes().to_vec()),
            Some(false),
            buf,
        )
        .unwrap_err();
        assert!(matches!(err, DbError::KeyExists), "{err:?}");
        assert_eq!(
            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_write_update_only_replaces_corrupt_row() {
        let (_dir, _seq_dir, tree, seq) = open_lsm_tree();
        let key = k(1);
        tree.keyspace()
            .insert(key.as_bytes(), b"not-a-doc")
            .unwrap();
        // The mirror image of the insert-only case, and deliberately asymmetric:
        // the precondition is existence-only, so update-only sees a present row
        // and replaces it. It must not surface `CorruptedEntry` — decoding the
        // old value is not part of this contract.
        let doc = Doc { name: "new".into() };
        let mut buf = Vec::new();
        RapiraCodec.encode_to(&doc, &mut buf).unwrap();
        RpcHandler::write(
            &adapter(&tree, &seq),
            WriteKey::Provided(key.as_bytes().to_vec()),
            Some(true),
            buf,
        )
        .expect("update-only over a present (if undecodable) row must succeed");
        assert_eq!(
            tree.get(&key),
            Some(doc),
            "update-only must have replaced the corrupt row"
        );
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_write_update_only_missing_key_does_not_create() {
        let (_dir, _seq_dir, tree, seq) = open_lsm_tree();
        let key = k(1);
        let mut buf = Vec::new();
        RapiraCodec
            .encode_to(&Doc { name: "new".into() }, &mut buf)
            .unwrap();
        let err = RpcHandler::write(
            &adapter(&tree, &seq),
            WriteKey::Provided(key.as_bytes().to_vec()),
            Some(true),
            buf,
        )
        .unwrap_err();
        assert!(matches!(err, DbError::KeyNotFound), "{err:?}");
        assert!(
            !tree.try_contains(&key).unwrap(),
            "a failed update-only must not create the row"
        );
    }

    #[cfg(feature = "rpc")]
    #[test]
    fn lsm_handler_count_exact_includes_corrupt_rows() {
        let (_dir, _seq_dir, tree, seq) = open_lsm_tree();
        tree.put(&k(1), &Doc { name: "ok".into() }).unwrap();
        // Inject a row that cannot decode as `Doc`, straight into the keyspace.
        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(&adapter(&tree, &seq), true).unwrap(), 2);
    }
}