armdb 0.4.1

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

use crate::config::Config;
use crate::error::DbResult;
use crate::meta::{Backend, DbMeta, META_VERSION};
use crate::shard::Shard;

#[cfg(feature = "encryption")]
use crate::crypto::PageCipher;

fn write_meta(
    path: &std::path::Path,
    config: &Config,
    #[cfg(feature = "encryption")] encrypted: bool,
) -> DbResult<()> {
    use std::io::Write;

    #[cfg(feature = "encryption")]
    let flags = if encrypted { 1 } else { 0 };
    #[cfg(not(feature = "encryption"))]
    let flags = 0;

    let meta = DbMeta {
        version: META_VERSION,
        backend: Backend::Bitcask,
        shard_count: config.shard_count as u8,
        shard_prefix_bits: config.shard_prefix_bits as u8,
        flags,
    };

    let tmp = path.with_extension("tmp");
    let mut f = std::fs::File::create(&tmp)?;
    f.write_all(&meta.encode())?;
    f.sync_data()?;
    drop(f);
    std::fs::rename(&tmp, path)?;
    if let Some(parent) = path.parent() {
        let dir = std::fs::File::open(parent)?;
        dir.sync_all()?;
    }
    Ok(())
}

fn validate_meta(
    path: &std::path::Path,
    config: &Config,
    #[cfg(feature = "encryption")] encrypted: bool,
) -> DbResult<()> {
    let bytes = std::fs::read(path)?;
    let meta = DbMeta::decode(&bytes)?;

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

    let was_encrypted = meta.flags & 1 != 0;

    #[cfg(feature = "encryption")]
    {
        if was_encrypted != encrypted {
            return Err(crate::error::DbError::Config(if was_encrypted {
                "database is encrypted but no encryption_key provided"
            } else {
                "database is not encrypted but encryption_key was provided"
            }));
        }
    }

    #[cfg(not(feature = "encryption"))]
    if was_encrypted {
        return Err(crate::error::DbError::FormatMismatch(
            "database is encrypted but binary lacks the encryption feature".into(),
        ));
    }

    Ok(())
}

/// Internal storage engine: owns config, shards, and optional cipher.
/// Each tree type embeds one `Engine` — one tree = one database.
#[allow(dead_code)]
pub(crate) struct Engine {
    path: PathBuf,
    config: Config,
    shards: Arc<Vec<Shard>>,
    gsn: Arc<AtomicU64>,
    #[cfg(feature = "encryption")]
    cipher: Option<Arc<PageCipher>>,
}

/// Whether this build can actually run the io_uring backend.
const URING_AVAILABLE: bool = cfg!(all(target_os = "linux", feature = "io-uring"));

/// Warn once per process when `IoBackend::Uring` was requested but io_uring is
/// not available in this build (the `io-uring` feature is off, or non-Linux).
/// The shard write path transparently falls back to `Pwrite` in that case; this
/// surfaces the silent degradation so it isn't mistaken for real io_uring use.
fn warn_if_uring_unavailable(io_backend: crate::config::IoBackend) {
    if !URING_AVAILABLE && matches!(io_backend, crate::config::IoBackend::Uring { .. }) {
        static WARNED: std::sync::Once = std::sync::Once::new();
        WARNED.call_once(|| {
            tracing::warn!(
                "io_backend=Uring requested but io_uring unavailable \
                 (feature \"io-uring\" disabled or non-Linux); using Pwrite"
            );
        });
    }
}

impl Engine {
    /// Open or create a database at the given path.
    /// Creates shard directories, validates encryption meta, initializes shards.
    #[tracing::instrument(skip(path, config), fields(path = %path.as_ref().display()))]
    pub fn open(path: impl AsRef<Path>, config: Config) -> DbResult<Self> {
        let path = path.as_ref().to_path_buf();
        config.validate()?;
        tracing::info!(shards = config.shard_count, "opening database");
        warn_if_uring_unavailable(config.io_backend);

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

        #[cfg(feature = "encryption")]
        let cipher = config
            .encryption_key
            .as_ref()
            .map(|key| PageCipher::new(key).map(Arc::new))
            .transpose()?;

        // db.meta: validate immutable parameters (shard_count, shard_prefix_bits, encryption)
        let meta_path = path.join("db.meta");
        let meta_existed = meta_path.exists();

        if meta_existed {
            validate_meta(
                &meta_path,
                &config,
                #[cfg(feature = "encryption")]
                cipher.is_some(),
            )?;
        } else {
            // F-8B: refuse to create fresh meta if orphaned shard dirs exist.
            let has_shard_dirs = std::fs::read_dir(&path)?.filter_map(|e| e.ok()).any(|e| {
                e.file_type().is_ok_and(|ft| ft.is_dir())
                    && e.file_name()
                        .to_str()
                        .is_some_and(|n| n.starts_with("shard_"))
            });
            if has_shard_dirs {
                return Err(crate::error::DbError::FormatMismatch(
                    "db.meta missing but shard directories present; \
                     refusing to recreate metadata"
                        .into(),
                ));
            }

            write_meta(
                &meta_path,
                &config,
                #[cfg(feature = "encryption")]
                cipher.is_some(),
            )?;
        }

        // F-5: existing DB must have all shard dirs intact.
        if meta_existed {
            for i in 0..config.shard_count {
                let shard_dir = path.join(format!("shard_{i:03}"));
                if !shard_dir.is_dir() {
                    return Err(crate::error::DbError::FormatMismatch(format!(
                        "shard directory {} missing; \
                         refusing to recreate (possible data loss)",
                        shard_dir.display()
                    )));
                }
            }
        }

        let effective_direct_io = if config.direct_io {
            let ok = crate::io::direct::probe_direct_io(&path);
            if !ok {
                tracing::warn!(
                    "direct_io requested but the filesystem does not support O_DIRECT; using buffered I/O"
                );
            }
            ok
        } else {
            false
        };

        let gsn = Arc::new(AtomicU64::new(1));

        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}"));
            #[cfg(feature = "encryption")]
            let shard = Shard::open_encrypted(
                i as u8,
                &shard_dir,
                config.max_file_size,
                config.write_buffer_size,
                config.hints,
                config.direct_io,
                effective_direct_io,
                config.io_backend,
                cipher.clone(),
                gsn.clone(),
            )?;
            #[cfg(not(feature = "encryption"))]
            let shard = Shard::open(
                i as u8,
                &shard_dir,
                config.max_file_size,
                config.write_buffer_size,
                config.hints,
                config.direct_io,
                effective_direct_io,
                config.io_backend,
                gsn.clone(),
            )?;
            shards.push(shard);
        }

        tracing::info!("database opened");
        Ok(Self {
            path,
            config,
            shards: Arc::new(shards),
            gsn,
            #[cfg(feature = "encryption")]
            cipher,
        })
    }

    #[allow(dead_code)]
    pub fn path(&self) -> &Path {
        &self.path
    }

    #[allow(dead_code)]
    pub fn config(&self) -> &Config {
        &self.config
    }

    pub fn gsn(&self) -> &AtomicU64 {
        &self.gsn
    }

    pub fn shards(&self) -> &Arc<Vec<Shard>> {
        &self.shards
    }

    pub fn shard_dirs(&self) -> Vec<std::path::PathBuf> {
        self.shards.iter().map(|s| s.dir().to_path_buf()).collect()
    }

    pub fn shard_dir_refs(dirs: &[std::path::PathBuf]) -> Vec<&Path> {
        dirs.iter().map(|p| p.as_path()).collect()
    }

    pub fn shard_ids(&self) -> Vec<u8> {
        self.shards.iter().map(|s| s.id).collect()
    }

    pub fn hints(&self) -> bool {
        self.config.hints
    }

    #[cfg(feature = "encryption")]
    pub fn cipher(&self) -> Option<Arc<PageCipher>> {
        self.cipher.clone()
    }

    /// Flush all shard write buffers to disk (without fsync).
    pub fn flush_buffers(&self) -> DbResult<()> {
        tracing::debug!("flushing all buffers");
        for shard in self.shards.iter() {
            shard.flush_buf()?;
        }
        Ok(())
    }

    /// Flush write buffers + fsync all shard active files.
    pub fn flush(&self) -> DbResult<()> {
        tracing::info!("flushing database");
        for shard in self.shards.iter() {
            shard.flush()?;
        }
        tracing::info!("database flushed");
        Ok(())
    }
}

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

    #[test]
    fn bitcask_meta_is_versioned_armd() {
        let dir = tempfile::tempdir().unwrap();
        let config = Config {
            shard_count: 2,
            shard_prefix_bits: 1,
            ..Config::test()
        };
        let _db = Engine::open(dir.path(), config.clone()).unwrap();

        let meta = std::fs::read(dir.path().join("db.meta")).unwrap();
        assert_eq!(
            meta.len(),
            16,
            "Bitcask db.meta must be the 16-byte ARMD format"
        );
        assert_eq!(&meta[0..4], b"ARMD");
        assert_eq!(meta[5], 0, "backend byte must be 0 (Bitcask)");

        // Reopen must validate cleanly.
        let _db2 = Engine::open(dir.path(), config).unwrap();
    }

    #[test]
    fn bitcask_rejects_fixedstore_meta() {
        let dir = tempfile::tempdir().unwrap();
        // Hand-write a FixedStore meta (backend = 1).
        let fixed = [b'A', b'R', b'M', b'D', 2, 1, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0];
        std::fs::write(dir.path().join("db.meta"), fixed).unwrap();
        let config = Config {
            shard_count: 2,
            shard_prefix_bits: 1,
            ..Config::test()
        };
        let err = Engine::open(dir.path(), config).err().unwrap();
        assert!(matches!(err, crate::error::DbError::FormatMismatch(_)));
    }
}