armdb 0.4.0

sharded bitcask key-value storage optimized for NVMe
Documentation
#![cfg(feature = "typed-tree")]

use armdb::codec::Codec;
use armdb::{Config, DbError, TypedMap};
use std::sync::atomic::{AtomicUsize, Ordering};
use tempfile::tempdir;

struct U64Codec;

impl Codec<u64> for U64Codec {
    fn encode_to(&self, value: &u64, buf: &mut Vec<u8>) -> armdb::DbResult<()> {
        buf.clear();
        buf.extend_from_slice(&value.to_be_bytes());
        Ok(())
    }

    fn decode_from(&self, bytes: &[u8]) -> armdb::DbResult<u64> {
        let arr: [u8; 8] = bytes
            .try_into()
            .map_err(|_| DbError::CorruptedEntry { offset: 0 })?;
        Ok(u64::from_be_bytes(arr))
    }
}

type Map = TypedMap<[u8; 8], u64, U64Codec>;

fn open(dir: &std::path::Path) -> Map {
    TypedMap::<[u8; 8], u64, U64Codec>::open(dir, Config::test(), U64Codec).unwrap()
}

struct DeleteHook {
    writes: std::sync::Arc<AtomicUsize>,
    last_old: std::sync::Arc<std::sync::Mutex<Option<u64>>>,
    last_new: std::sync::Arc<std::sync::Mutex<Option<u64>>>,
}

impl Default for DeleteHook {
    fn default() -> Self {
        Self {
            writes: std::sync::Arc::new(AtomicUsize::new(0)),
            last_old: std::sync::Arc::new(std::sync::Mutex::new(None)),
            last_new: std::sync::Arc::new(std::sync::Mutex::new(None)),
        }
    }
}

impl armdb::TypedWriteHook<[u8; 8], u64> for DeleteHook {
    fn on_write(&self, _key: &[u8; 8], old: Option<&u64>, new: Option<&u64>) {
        self.writes.fetch_add(1, Ordering::Relaxed);
        *self.last_old.lock().unwrap() = old.copied();
        *self.last_new.lock().unwrap() = new.copied();
    }
}

#[test]
fn compare_delete_match_mismatch_absent() {
    let dir = tempdir().unwrap();
    let map = open(dir.path());

    let k = 1u64.to_be_bytes();
    map.put(&k, 42).unwrap();

    let err = map.compare_delete(&k, &99).unwrap_err();
    assert!(matches!(err, DbError::CasMismatch));
    assert_eq!(map.get(&k).map(|r| *r), Some(42));

    map.compare_delete(&k, &42).unwrap();
    assert!(map.get(&k).is_none());

    let err = map.compare_delete(&k, &42).unwrap_err();
    assert!(matches!(err, DbError::KeyNotFound));
}

#[test]
fn compare_delete_fires_hook_and_retires() {
    let dir = tempdir().unwrap();
    let hook = DeleteHook::default();
    let writes = std::sync::Arc::clone(&hook.writes);
    let last_old = std::sync::Arc::clone(&hook.last_old);
    let last_new = std::sync::Arc::clone(&hook.last_new);
    let map = TypedMap::<[u8; 8], u64, U64Codec, DeleteHook>::open_hooked(
        dir.path(),
        Config::test(),
        U64Codec,
        hook,
    )
    .unwrap();

    let k = 5u64.to_be_bytes();
    map.put(&k, 100).unwrap();

    let writes_before = writes.load(Ordering::Relaxed);
    map.compare_delete(&k, &100).unwrap();

    // exactly one on_write(Some(old), None) for the delete
    assert_eq!(writes.load(Ordering::Relaxed) - writes_before, 1);
    assert_eq!(*last_old.lock().unwrap(), Some(100));
    assert_eq!(*last_new.lock().unwrap(), None);

    // retire path (MutationResult.retire = Some) must be sound:
    map.put(&k, 7).unwrap();
    assert_eq!(map.get(&k).map(|r| *r), Some(7));
}