armdb 0.2.0

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::sync::Arc;

use crate::{DbResult, FixedConfig, FixedMap};

/// Centralized sequential ID generator backed by a `FixedMap<[u8; 8], 8>`.
///
/// Shared via `Arc<SeqGen>` between [`Db`](super::tree_db::Db) and collections.
pub struct SeqGen {
    map: FixedMap<[u8; 8], 8>,
}

impl SeqGen {
    pub(crate) fn open(path: impl AsRef<std::path::Path>) -> DbResult<Arc<Self>> {
        Ok(Arc::new(Self {
            map: FixedMap::<[u8; 8], 8>::open(
                path,
                FixedConfig {
                    shard_count: 4,
                    ..Default::default()
                },
            )?,
        }))
    }

    /// Generate the next sequential ID for a named collection.
    ///
    /// Uses [`FixedMap::atomic`] to perform the read-modify-write under a single
    /// shard lock, eliminating the TOCTOU race that `get` + `put` would have.
    pub fn next_id(&self, name: &str) -> DbResult<u64> {
        let key = xxhash_rust::xxh3::xxh3_64(name.as_bytes()).to_le_bytes();
        self.map.atomic(&key, |shard| {
            let id = shard.get(&key).map(u64::from_le_bytes).unwrap_or(0);
            let next = id + 1;
            shard.put(&key, &next.to_le_bytes())?;
            Ok(next)
        })
    }

    /// Flush buffers to disk without consuming self.
    pub fn flush(&self) -> DbResult<()> {
        self.map.flush()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashSet;
    use std::sync::Arc;
    use std::thread;

    #[test]
    fn concurrent_next_id_no_duplicates() {
        let tmp = tempfile::tempdir().expect("tmp");
        let seq = SeqGen::open(tmp.path().join("seq")).expect("open");

        const THREADS: usize = 16;
        const PER_THREAD: usize = 500;

        let handles: Vec<_> = (0..THREADS)
            .map(|_| {
                let seq = Arc::clone(&seq);
                thread::spawn(move || {
                    let mut ids = Vec::with_capacity(PER_THREAD);
                    for _ in 0..PER_THREAD {
                        ids.push(seq.next_id("ctr").expect("next_id"));
                    }
                    ids
                })
            })
            .collect();

        let mut all = Vec::with_capacity(THREADS * PER_THREAD);
        for h in handles {
            all.extend(h.join().expect("join"));
        }

        let unique: HashSet<_> = all.iter().copied().collect();
        assert_eq!(unique.len(), all.len(), "duplicate id detected");
        assert_eq!(unique.len(), THREADS * PER_THREAD);
    }
}