armdb 0.5.3

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::path::{Path, PathBuf};
use std::sync::Arc;

use crate::sync::{self, Mutex};

use crate::error::{DbError, DbResult};
use crate::fixed::config::FixedConfig;
use crate::fixed::shard::FixedShardInner;
use crate::meta::{Backend, DbMeta, META_VERSION};

/// Wrapper that pairs a shard id with a mutex-protected `FixedShardInner`.
#[allow(dead_code)]
pub(crate) struct FixedShard {
    pub id: u8,
    pub inner: Mutex<FixedShardInner>,
}

/// Multi-shard engine for the fixed-slot storage backend.
///
/// Manages the lifecycle of all shards, the `db.meta` file, and
/// open/close coordination.
pub(crate) struct FixedEngine {
    path: PathBuf,
    shards: Arc<Vec<FixedShard>>,
    config: FixedConfig,
}

// ── db.meta I/O ───────────────────────────────────────────────────

fn write_meta(path: &Path, config: &FixedConfig) -> DbResult<()> {
    let meta = DbMeta {
        version: META_VERSION,
        backend: Backend::FixedStore,
        shard_count: config.shard_count as u8,
        shard_prefix_bits: config.shard_prefix_bits as u8,
        flags: 0,
    };

    let file = std::fs::File::create(path)?;
    use std::io::Write;
    (&file).write_all(&meta.encode())?;
    file.sync_all()?;

    if let Some(parent) = path.parent() {
        let dir = std::fs::File::open(parent)?;
        dir.sync_all()?;
    }
    Ok(())
}

fn validate_meta(path: &Path, config: &FixedConfig) -> DbResult<()> {
    let bytes = std::fs::read(path)?;
    let meta = DbMeta::decode(&bytes)?;

    if meta.backend != Backend::FixedStore {
        return Err(DbError::FormatMismatch(format!(
            "db.meta: expected FixedStore backend, found {:?}",
            meta.backend
        )));
    }
    if meta.shard_count as usize != config.shard_count {
        return Err(DbError::FormatMismatch(format!(
            "db.meta: shard_count mismatch: stored {}, config {}",
            meta.shard_count, config.shard_count
        )));
    }
    if meta.shard_prefix_bits as usize != config.shard_prefix_bits {
        return Err(DbError::FormatMismatch(format!(
            "db.meta: shard_prefix_bits mismatch: stored {}, config {}",
            meta.shard_prefix_bits, config.shard_prefix_bits
        )));
    }
    Ok(())
}

// ── FixedEngine ───────────────────────────────────────────────────

impl FixedEngine {
    /// Open or create a fixed-slot database at the given path.
    ///
    /// 1. Validates configuration.
    /// 2. Creates the root directory if needed.
    /// 3. Reads/writes `db.meta`.
    /// 4. Opens every shard (`shard_000/`, `shard_001/`, ...).
    #[tracing::instrument(skip(path,config), fields(path = %path.as_ref().display()))]
    pub fn open(
        path: impl AsRef<Path>,
        config: FixedConfig,
        key_len: u16,
        value_len: u16,
    ) -> DbResult<Self> {
        let path = path.as_ref().to_path_buf();
        config.validate()?;

        tracing::info!(shards = config.shard_count, "opening fixed-slot database");

        std::fs::create_dir_all(&path)?;

        // db.meta: validate or create
        let meta_path = path.join("db.meta");
        if meta_path.exists() {
            validate_meta(&meta_path, &config)?;
        } else {
            write_meta(&meta_path, &config)?;
        }

        // Open each shard
        let mut shards = Vec::with_capacity(config.shard_count);
        for i in 0..config.shard_count {
            let shard_dir = path.join(format!("shard_{i:03}"));
            let inner = FixedShardInner::open(&shard_dir, i as u8, key_len, value_len, &config)?;
            shards.push(FixedShard {
                id: i as u8,
                inner: Mutex::new(inner),
            });
        }

        tracing::info!("fixed-slot database opened");
        Ok(Self {
            path,
            shards: Arc::new(shards),
            config,
        })
    }

    /// Reference to the shard list.
    #[allow(dead_code)]
    pub fn shards(&self) -> &Arc<Vec<FixedShard>> {
        &self.shards
    }

    /// Reference to the engine configuration.
    #[allow(dead_code)]
    pub fn config(&self) -> &FixedConfig {
        &self.config
    }

    /// Root path of the database.
    #[allow(dead_code)]
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Sync all shard data files to disk.
    pub fn flush(&self) -> DbResult<()> {
        tracing::debug!("flushing all fixed-slot shards");
        for shard in self.shards.iter() {
            sync::lock(&shard.inner).sync()?;
        }
        Ok(())
    }

    /// Install SPSC producers into each shard's `replication_tx`.
    /// Called by `FixedReplicationServer` on the first follower accept.
    /// Panics if any shard already has a producer installed or if the
    /// producers slice length mismatches `shard_count`.
    #[cfg(feature = "replication")]
    #[allow(dead_code)]
    pub fn install_replication_producers(
        &self,
        producers: Vec<rtrb::Producer<crate::fixed_replication::FixedReplicationEvent>>,
    ) {
        assert_eq!(
            producers.len(),
            self.shards.len(),
            "producer count must match shard count"
        );
        for (shard, tx) in self.shards.iter().zip(producers) {
            let mut inner = sync::lock(&shard.inner);
            assert!(
                inner.replication_tx.is_none(),
                "producer already installed for shard {}",
                shard.id
            );
            inner.replication_tx = Some(tx);
        }
    }

    /// Perform a clean shutdown of every shard.
    pub fn close(&self) -> DbResult<()> {
        tracing::info!("closing fixed-slot database");
        for shard in self.shards.iter() {
            sync::lock(&shard.inner).clean_shutdown()?;
        }
        tracing::info!("fixed-slot database closed");
        Ok(())
    }
}

#[cfg(feature = "replication")]
impl crate::fixed_replication::FixedEngineAccess for FixedEngine {
    fn shard_count(&self) -> usize {
        self.shards.len()
    }

    fn key_len(&self) -> usize {
        sync::lock(&self.shards[0].inner).key_len() as usize
    }

    fn value_len(&self) -> usize {
        sync::lock(&self.shards[0].inner).value_len() as usize
    }

    fn slot_size(&self) -> u16 {
        sync::lock(&self.shards[0].inner).slot_size
    }

    fn current_slot_count(&self, shard_id: usize) -> u32 {
        sync::lock(&self.shards[shard_id].inner).slot_count()
    }

    fn shard_prefix_bits(&self) -> u8 {
        self.config.shard_prefix_bits as u8
    }

    fn read_shard_chunk(
        &self,
        shard_id: usize,
        start_slot: u32,
        count: usize,
    ) -> crate::error::DbResult<Vec<u8>> {
        let inner = sync::lock(&self.shards[shard_id].inner);
        let slot_size = inner.slot_size as usize;
        let mut buf = vec![0u8; count * slot_size];
        let offset = crate::fixed::shard::HEADER_SIZE + start_slot as u64 * slot_size as u64;
        inner.read_chunk_at(offset, &mut buf)?;
        Ok(buf)
    }

    fn install_replication_producers(
        &self,
        producers: Vec<rtrb::Producer<crate::fixed_replication::FixedReplicationEvent>>,
    ) {
        FixedEngine::install_replication_producers(self, producers);
    }

    fn update_min_replicated_version(&self, _shard_id: usize, _version: u32) {
        // Metric-only hook; no-op for v1. A future version can track per-shard
        // atomic gauges for dashboard visibility.
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    fn test_config() -> FixedConfig {
        FixedConfig {
            shard_count: 2,
            grow_step: 64,
            ..FixedConfig::test()
        }
    }

    #[test]
    fn test_engine_open_creates_shards() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("testdb");
        let config = test_config();

        let engine = FixedEngine::open(&db_path, config.clone(), 8, 32).unwrap();

        // db.meta exists
        assert!(db_path.join("db.meta").exists());

        // Shard directories exist
        for i in 0..config.shard_count {
            assert!(
                db_path.join(format!("shard_{i:03}")).exists(),
                "shard_{i:03} directory should exist"
            );
        }

        // Correct number of shards
        assert_eq!(engine.shards().len(), config.shard_count);

        // Each shard has the right id
        for (i, shard) in engine.shards().iter().enumerate() {
            assert_eq!(shard.id, i as u8);
        }
    }

    #[test]
    fn test_engine_reopen() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("testdb");
        let config = test_config();

        // Open and close
        {
            let engine = FixedEngine::open(&db_path, config.clone(), 8, 32).unwrap();
            engine.close().unwrap();
        }

        // Reopen — must succeed with same config
        {
            let engine = FixedEngine::open(&db_path, config.clone(), 8, 32).unwrap();
            assert_eq!(engine.shards().len(), config.shard_count);
            engine.close().unwrap();
        }
    }

    #[test]
    fn fixedstore_rejects_bitcask_meta() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("testdb");
        let config = test_config();

        // Create the database directory and write a valid Bitcask db.meta
        // (backend byte = 0). shard_count/shard_prefix_bits match test_config()
        // so the backend check fires before any shard mismatch check.
        std::fs::create_dir_all(&db_path).unwrap();
        let meta_path = db_path.join("db.meta");
        let bitcask_meta: [u8; 16] = [
            b'A',
            b'R',
            b'M',
            b'D', // magic
            2,    // version
            0,    // backend: 0 = Bitcask
            config.shard_count as u8,
            config.shard_prefix_bits as u8,
            0, // flags
            0,
            0,
            0,
            0,
            0,
            0,
            0, // reserved
        ];
        std::fs::write(&meta_path, bitcask_meta).unwrap();

        let result = FixedEngine::open(&db_path, config, 8, 32);
        assert!(result.is_err());
        let err = result.err().unwrap();
        assert!(
            matches!(err, DbError::FormatMismatch(_)),
            "expected FormatMismatch, got: {err}"
        );
    }

    #[test]
    fn test_engine_shard_count_mismatch() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("testdb");

        // Open with 2 shards
        {
            let config = FixedConfig {
                shard_count: 2,
                grow_step: 64,
                ..FixedConfig::test()
            };
            let engine = FixedEngine::open(&db_path, config, 8, 32).unwrap();
            engine.close().unwrap();
        }

        // Reopen with 4 shards — should fail
        {
            let config = FixedConfig {
                shard_count: 4,
                grow_step: 64,
                ..FixedConfig::test()
            };
            let result = FixedEngine::open(&db_path, config, 8, 32);
            assert!(result.is_err());
            let msg = result.err().unwrap().to_string();
            assert!(
                msg.contains("shard_count mismatch"),
                "expected shard_count mismatch error, got: {msg}"
            );
        }
    }
}