armdb 0.7.0

//! Integration tests for Db::clean_shutdown and Fixed clean-recovery path.

#![cfg(all(
    feature = "var-collections",
    feature = "typed-tree",
    feature = "armour",
    feature = "rapira-codec"
))]

use std::path::{Path, PathBuf};

use armdb::armour::Db;
use armdb::fixed::shard::FixedShardInner;
use armdb::{CollectionMeta, Config, FixedConfig, NoHook, RapiraCodec};
use armour_core::GetType;
use rapira::Rapira;
use tempfile::tempdir;

fn shard_id_from_dir(shard_dir: &Path) -> u8 {
    let name = shard_dir.file_name().unwrap().to_str().unwrap();
    name.strip_prefix("shard_")
        .and_then(|s| s.parse().ok())
        .expect("shard dir name")
}

fn fixed_shard_dirs(collection_path: &Path) -> Vec<PathBuf> {
    std::fs::read_dir(collection_path)
        .unwrap()
        .filter_map(|e| {
            let p = e.ok()?.path();
            let name = p.file_name()?.to_str()?;
            if name.starts_with("shard_") {
                Some(p)
            } else {
                None
            }
        })
        .collect()
}

fn all_fixed_shards_clean(
    collection_path: &Path,
    key_len: u16,
    value_len: u16,
    cfg: &FixedConfig,
) -> bool {
    fixed_shard_dirs(collection_path).iter().all(|shard_dir| {
        let shard_id = shard_id_from_dir(shard_dir);
        let shard = FixedShardInner::open(shard_dir, shard_id, key_len, value_len, cfg).unwrap();
        shard.has_clean_shutdown() && shard_dir.join("fixed.versions").exists()
    })
}

fn any_shard_has_clean_flag(
    collection_path: &Path,
    key_len: u16,
    value_len: u16,
    cfg: &FixedConfig,
) -> bool {
    fixed_shard_dirs(collection_path)
        .iter()
        .any(|shard_dir| shard_has_clean_flag(shard_dir, key_len, value_len, cfg))
}

fn shard_has_clean_flag(shard_dir: &Path, key_len: u16, value_len: u16, cfg: &FixedConfig) -> bool {
    let shard_id = shard_id_from_dir(shard_dir);
    FixedShardInner::open(shard_dir, shard_id, key_len, value_len, cfg)
        .unwrap()
        .has_clean_shutdown()
}

fn shard_dir_for_index(collection_path: &Path, shard_idx: usize) -> PathBuf {
    collection_path.join(format!("shard_{shard_idx:03}"))
}

/// Check if any file with the given suffix exists recursively under `root`.
fn walk_has_suffix(root: &Path, suffix: &str) -> bool {
    fn walk(dir: &Path, suffix: &str) -> bool {
        let Ok(entries) = std::fs::read_dir(dir) else {
            return false;
        };
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                if walk(&path, suffix) {
                    return true;
                }
            } else if path.to_string_lossy().ends_with(suffix) {
                return true;
            }
        }
        false
    }
    walk(root, suffix)
}

#[derive(
    Clone,
    Copy,
    Debug,
    PartialEq,
    GetType,
    Rapira,
    zerocopy::FromBytes,
    zerocopy::IntoBytes,
    zerocopy::Immutable,
)]
#[repr(C)]
struct FxCounter {
    n: u32,
}

impl CollectionMeta for FxCounter {
    type SelfId = [u8; 8];
    const NAME: &'static str = "clean_shutdown_counters";
    const VERSION: u16 = 1;
}

#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct User {
    id: u64,
    name: String,
}

impl CollectionMeta for User {
    type SelfId = [u8; 8];
    const NAME: &'static str = "clean_shutdown_users";
    const VERSION: u16 = 1;
}

#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct Blob {
    data: String,
}

impl CollectionMeta for Blob {
    type SelfId = [u8; 8];
    const NAME: &'static str = "clean_shutdown_blobs";
    const VERSION: u16 = 1;
}

#[test]
fn test_db_clean_shutdown_fixed_happy_path() {
    let dir = tempdir().unwrap();
    let fx_cfg = FixedConfig::test();
    let collection_path;

    {
        let db = Db::open_test(dir.path()).unwrap();
        let tree = db
            .open_zero_tree_fixed::<FxCounter, 4, _>(fx_cfg.clone(), NoHook, &[])
            .unwrap();
        for i in 0..20u64 {
            tree.put(&i.to_be_bytes(), &FxCounter { n: i as u32 })
                .unwrap();
        }
        collection_path = db.tree_path(FxCounter::NAME);
        db.clean_shutdown().unwrap();
    }

    assert!(
        all_fixed_shards_clean(&collection_path, 8, 4, &fx_cfg),
        "sidecar + clean flag expected before reopen"
    );

    {
        let db = Db::open_test(dir.path()).unwrap();
        let tree = db
            .open_zero_tree_fixed::<FxCounter, 4, _>(fx_cfg.clone(), NoHook, &[])
            .unwrap();
        assert_eq!(tree.len(), 20);
        for i in 0..20u64 {
            let v = tree.get(&i.to_be_bytes()).unwrap();
            assert_eq!(v.n, i as u32);
        }
        // Recovery clears the flag during open.
        assert!(
            !any_shard_has_clean_flag(&collection_path, 8, 4, &fx_cfg),
            "flag must be cleared after reopen"
        );
    }
}

#[test]
fn test_db_clean_shutdown_fixed_downgrade_on_write() {
    let dir = tempdir().unwrap();
    let fx_cfg = FixedConfig::test();

    let db = Db::open_test(dir.path()).unwrap();
    let tree = db
        .open_zero_tree_fixed::<FxCounter, 4, _>(fx_cfg.clone(), NoHook, &[])
        .unwrap();
    tree.put(&1u64.to_be_bytes(), &FxCounter { n: 1 }).unwrap();
    let collection_path = db.tree_path(FxCounter::NAME);
    db.clean_shutdown().unwrap();
    assert!(all_fixed_shards_clean(&collection_path, 8, 4, &fx_cfg));

    // Straggler write after clean_shutdown.
    tree.put(&2u64.to_be_bytes(), &FxCounter { n: 2 }).unwrap();
    let affected = shard_dir_for_index(&collection_path, tree.shard_for(&2u64.to_be_bytes()));
    assert!(
        !shard_has_clean_flag(&affected, 8, 4, &fx_cfg),
        "post-shutdown write must clear clean flag on the affected shard"
    );

    // The reopen below builds a second engine over these shard files: drop the
    // first handle explicitly, or its lease (and its index) would still be live.
    // `let tree = ...` below only shadows the binding — it does not drop it.
    drop(tree);
    drop(db);
    let db = Db::open_test(dir.path()).unwrap();
    let tree = db
        .open_zero_tree_fixed::<FxCounter, 4, _>(fx_cfg, NoHook, &[])
        .unwrap();
    assert_eq!(tree.len(), 2);
    assert_eq!(tree.get(&1u64.to_be_bytes()).unwrap().n, 1);
    assert_eq!(tree.get(&2u64.to_be_bytes()).unwrap().n, 2);
}

#[test]
fn test_db_clean_shutdown_idempotent() {
    let dir = tempdir().unwrap();
    let fx_cfg = FixedConfig::test();
    let db = Db::open_test(dir.path()).unwrap();
    let _tree = db
        .open_zero_tree_fixed::<FxCounter, 4, _>(fx_cfg.clone(), NoHook, &[])
        .unwrap();
    let collection_path = db.tree_path(FxCounter::NAME);
    db.clean_shutdown().unwrap();
    db.clean_shutdown().unwrap();
    assert!(all_fixed_shards_clean(&collection_path, 8, 4, &fx_cfg));
}

#[test]
fn test_db_clean_shutdown_survives_drop() {
    let dir = tempdir().unwrap();
    let fx_cfg = FixedConfig::test();
    let collection_path;

    {
        let db = Db::open_test(dir.path()).unwrap();
        let _tree = db
            .open_zero_tree_fixed::<FxCounter, 4, _>(fx_cfg.clone(), NoHook, &[])
            .unwrap();
        collection_path = db.tree_path(FxCounter::NAME);
        db.clean_shutdown().unwrap();
        // Drop → shutdown → close (flush only, no Fixed sidecar).
    }

    assert!(
        all_fixed_shards_clean(&collection_path, 8, 4, &fx_cfg),
        "clean flag must survive Db drop"
    );
}

#[test]
fn test_db_clean_shutdown_bitcask_hints() {
    let dir = tempdir().unwrap();
    let cfg = Config::balanced().shard_count(2).hints(true).build();

    {
        let db = Db::open_test(dir.path()).unwrap();
        let users = db
            .open_typed_tree::<User, RapiraCodec, _>(cfg, NoHook, &[])
            .unwrap();
        users
            .put(
                &1u64.to_be_bytes(),
                User {
                    id: 1,
                    name: "ada".into(),
                },
            )
            .unwrap();
        db.clean_shutdown().unwrap();
    }

    assert!(
        walk_has_suffix(dir.path(), ".hint"),
        "clean_shutdown must write Bitcask hint files when hints enabled"
    );
}

#[test]
fn test_db_clean_shutdown_var_typed_tree() {
    let dir = tempdir().unwrap();
    let cfg = Config::balanced().shard_count(2).hints(true).build();
    let db = Db::open_test(dir.path()).unwrap();
    let tree = db
        .open_var_typed_tree::<Blob, RapiraCodec, _>(cfg, NoHook, &[])
        .unwrap();
    let blob = Blob {
        data: "hello".into(),
    };
    tree.put(&1u64.to_be_bytes(), &blob).unwrap();
    db.clean_shutdown().unwrap();
}

#[test]
fn test_db_clean_shutdown_mixed_backends() {
    let dir = tempdir().unwrap();
    let bitcask_cfg = Config::balanced().shard_count(2).hints(true).build();
    let fx_cfg = FixedConfig::test();

    let db = Db::open_test(dir.path()).unwrap();
    let users = db
        .open_typed_tree::<User, RapiraCodec, _>(bitcask_cfg, NoHook, &[])
        .unwrap();
    users
        .put(
            &1u64.to_be_bytes(),
            User {
                id: 1,
                name: "bob".into(),
            },
        )
        .unwrap();
    let counters = db
        .open_zero_tree_fixed::<FxCounter, 4, _>(fx_cfg.clone(), NoHook, &[])
        .unwrap();
    counters
        .put(&1u64.to_be_bytes(), &FxCounter { n: 42 })
        .unwrap();

    let fx_path = db.tree_path(FxCounter::NAME);
    db.clean_shutdown().unwrap();

    assert!(walk_has_suffix(dir.path(), ".hint"));
    assert!(all_fixed_shards_clean(&fx_path, 8, 4, &fx_cfg));
}