armdb 0.14.2

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, Shards};

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

#[cfg(test)]
pub(crate) static ENGINE_DURABLE_FLUSHES: AtomicU64 = AtomicU64::new(0);

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<Shards>,
    gsn: Arc<AtomicU64>,
    /// Background hint generation for files this collection's shards retire.
    /// Held here so it outlives the shards that submit to it and is stopped —
    /// draining what is queued — when the collection closes.
    hint_service: Arc<crate::hint_worker::HintService>,
    #[cfg(feature = "encryption")]
    ciphers: Option<Arc<Vec<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);
        if config.enable_fsync {
            static WARNED: std::sync::Once = std::sync::Once::new();
            WARNED.call_once(|| {
                tracing::warn!("`enable_fsync` has no effect on the Bitcask backend");
            });
        }

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

        // Taken before db.meta is read or written and before the F-8B orphan
        // scan, so recovery and migration run under the lease too. It is handed
        // to `Shards` rather than kept here: the replication threads clone that
        // Arc out of the Engine and outlive it, and must keep the lock alive.
        let lock = crate::lock::acquire(&path)?;

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

        #[cfg(feature = "encryption")]
        let encrypted = config.encryption_key.is_some();

        // Salt sidecar (encryption only). Existing DB: validate meta first (clearer
        // error), then load the salt. Fresh DB: generate + write the salt BEFORE
        // db.meta, so a crash between them leaves no meta and the next open starts
        // fresh (regenerating the salt).
        #[cfg(feature = "encryption")]
        let salt: Option<[u8; crate::salt::SALT_LEN]>;

        if meta_existed {
            validate_meta(
                &meta_path,
                &config,
                #[cfg(feature = "encryption")]
                encrypted,
            )?;
            #[cfg(feature = "encryption")]
            {
                salt = if encrypted {
                    Some(crate::salt::read(&path)?)
                } else {
                    None
                };
            }
        } 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(),
                ));
            }

            #[cfg(feature = "encryption")]
            {
                salt = if encrypted {
                    let s = crate::salt::generate()?;
                    crate::salt::write(&path, &s)?;
                    Some(s)
                } else {
                    None
                };
            }

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

        // 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()
                    )));
                }
            }
        }

        // Derive one AES-256-GCM subkey per shard from the master key + salt.
        #[cfg(feature = "encryption")]
        let ciphers: Option<Arc<Vec<Arc<PageCipher>>>> =
            match (config.encryption_key.as_ref(), salt) {
                (Some(key), Some(salt)) => {
                    let mut v = Vec::with_capacity(config.shard_count);
                    for i in 0..config.shard_count {
                        v.push(Arc::new(PageCipher::derive(key, &salt, i as u8)?));
                    }
                    Some(Arc::new(v))
                }
                _ => None,
            };

        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_with_sync_thresholds(
                i as u8,
                &shard_dir,
                config.max_file_size,
                config.write_buffer_size,
                config.hints.unwrap_or(false),
                config.direct_io,
                effective_direct_io,
                config.io_backend,
                ciphers.as_ref().map(|c| c[i].clone()),
                config.sync_batch_size,
                config.sync_interval,
                gsn.clone(),
            )?;
            #[cfg(not(feature = "encryption"))]
            let shard = Shard::open_with_sync_thresholds(
                i as u8,
                &shard_dir,
                config.max_file_size,
                config.write_buffer_size,
                config.hints.unwrap_or(false),
                config.direct_io,
                effective_direct_io,
                config.io_backend,
                config.sync_batch_size,
                config.sync_interval,
                gsn.clone(),
            )?;
            shard.set_compaction_max_files_per_pass(config.compaction_max_files_per_pass);
            shards.push(shard);
        }

        // F8: one background hint generator per collection. The thread itself is
        // started lazily, on the first rotation that would produce a hint, so a
        // collection that never rotates never spawns one.
        let hint_service = Arc::new(crate::hint_worker::HintService::new());
        let shards = Arc::new(Shards::new(shards, lock));
        hint_service.attach_shards(Arc::downgrade(&shards));
        for shard in shards.iter() {
            shard.set_hint_service(hint_service.clone());
        }

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

    /// Clean shutdown: flush every shard, then wait for the hint worker to
    /// finish what it has queued.
    ///
    /// The second half is what restores the guarantee F8 relaxes — hint
    /// generation is asynchronous, so only a clean close can promise that every
    /// retired file has its hint on disk. Distinct from [`Self::flush`], which
    /// the periodic flusher calls and must not stop the worker. Idempotent.
    pub(crate) fn close(&self) -> DbResult<()> {
        self.flush()?;
        self.hint_service.drain_and_stop();
        Ok(())
    }

    /// Stop the hint worker before the shard set — and with it the directory
    /// lease — is released.
    ///
    /// F8's worker upgrades its `Weak<Shards>` for the length of a job, so an
    /// in-flight job keeps the lease alive past the drop. Relying on the Arc
    /// graph to sort itself out does not work: `shards` drops first and, while
    /// the worker holds it up, every `ShardInner`'s `Arc<HintService>` is
    /// still alive, so `HintService::drop` (and its join) never runs during
    /// this drop. An immediate reopen in the same process then fails with
    /// `DbError::Locked`.
    fn stop_hint_worker(&self) {
        self.hint_service.drain_and_stop();
    }

    /// Wait for every queued hint to be written, without closing anything.
    ///
    /// Test-only: `close` is the production way to get this guarantee, but a
    /// test that inspects a retired file's hint — or exercises the replication
    /// guard, which reads one — needs the same barrier mid-flight. The worker
    /// restarts on the next rotation, so this is safe to call repeatedly.
    #[cfg(test)]
    pub(crate) fn drain_hints_for_test(&self) {
        self.hint_service.drain_and_stop();
    }

    /// Hint jobs dropped because the worker's queue was full. Test/observability
    /// hook — see `crate::hint_worker`.
    #[cfg_attr(not(test), allow(dead_code))]
    pub(crate) fn hint_jobs_dropped(&self) -> u64 {
        self.hint_service.dropped()
    }

    #[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<Shards> {
        &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.unwrap_or(false)
    }

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

    /// How many shards currently hold writes no barrier has made durable.
    ///
    /// Sampled without a consistent cut across shards: a concurrent writer or
    /// barrier can move the answer while it is being counted, so treat it as a
    /// gauge of the loss window, not as an exact instant.
    pub fn dirty_shards(&self) -> usize {
        self.shards.iter().filter(|shard| shard.is_dirty()).count()
    }

    /// 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<()> {
        #[cfg(test)]
        ENGINE_DURABLE_FLUSHES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        tracing::debug!("flushing database");
        let mut first_err = None;
        for shard in self.shards.iter() {
            if let Err(e) = shard.flush()
                && first_err.is_none()
            {
                first_err = Some(e);
            }
        }
        tracing::debug!("database flushed");
        match first_err {
            Some(e) => Err(e),
            None => Ok(()),
        }
    }
}

impl Drop for Engine {
    fn drop(&mut self) {
        self.stop_hint_worker();
    }
}

#[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)");

        // Drop the first engine (releases the directory lock) before reopening.
        drop(_db);

        // 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(_)));
    }
}