armdb 0.14.4

sharded bitcask key-value storage optimized for NVMe
Documentation
//! A typed value migration between two checked rewrites: interrupted rewrites
//! over the `format.v` it leaves behind, and pending txlog replay across it.
#![cfg(feature = "test-faults")]

use std::process::Command;
use std::sync::atomic::Ordering;

use armdb::armour::{Db, TypedMigration};
use armdb::format_migration::failpoints;
use armdb::{
    CollectionMeta, Config, DbError, FormatMigration, MigrateAction, NoHook, RapiraCodec, TxConfig,
};
use armour_core::GetType;
use rapira::Rapira;
use tempfile::tempdir;

const NAME: &str = "value_tx_rows";

mod v1 {
    use super::*;
    #[derive(Clone, Debug, PartialEq, Rapira, GetType)]
    pub struct Row {
        pub value: u32,
    }
    impl CollectionMeta for Row {
        type SelfId = [u8; 8];
        const NAME: &'static str = NAME;
        const VERSION: u16 = 1;
    }
}

/// v1 → v2: checked rewrite (writes `format.v = 2`).
mod v2 {
    use super::*;
    #[derive(Clone, Debug, PartialEq, Rapira, GetType)]
    pub struct Row {
        pub value: u64,
    }
    pub const WIDEN: FormatMigration = armdb::format_step!(
        1, super::v1::Row => Row, RapiraCodec,
        |old| Ok::<_, DbError>(Some(Row { value: u64::from(old.value) }))
    );
    impl CollectionMeta for Row {
        type SelfId = [u8; 8];
        const NAME: &'static str = NAME;
        const VERSION: u16 = 2;
        const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[WIDEN];
    }
}

/// v2 → v3: typed value migration over the same type; `format.v` stays 2.
mod v3 {
    use super::*;
    #[derive(Clone, Debug, PartialEq, Rapira, GetType)]
    pub struct Row {
        pub value: u64,
    }
    impl CollectionMeta for Row {
        type SelfId = [u8; 8];
        const NAME: &'static str = NAME;
        const VERSION: u16 = 3;
        const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[super::v2::WIDEN];
    }
    pub fn bump(_key: &[u8; 8], old: &Row) -> MigrateAction<Row> {
        MigrateAction::Update(Row {
            value: old.value + 1000,
        })
    }
}

/// v3 → v4: checked rewrite; the 1 → 2 rewrite is retired.
mod v4 {
    use super::*;
    #[derive(Clone, Debug, PartialEq, Rapira, GetType)]
    pub struct Row {
        pub value: u64,
        pub tag: u16,
    }
    pub const EXTEND: FormatMigration = armdb::format_step!(
        3, super::v3::Row => Row, RapiraCodec,
        |old| Ok::<_, DbError>(Some(Row { value: old.value, tag: 7 }))
    );
    impl CollectionMeta for Row {
        type SelfId = [u8; 8];
        const NAME: &'static str = NAME;
        const VERSION: u16 = 4;
        const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[EXTEND];
    }
}

/// v3 → v4 with the 1 → 2 rewrite still declared. v2 and v3 share one
/// fingerprint, so a logged v3 record names both versions.
mod v4_kept {
    use super::*;
    #[derive(Clone, Debug, PartialEq, Rapira, GetType)]
    pub struct Row {
        pub value: u64,
        pub tag: u16,
    }
    impl CollectionMeta for Row {
        type SelfId = [u8; 8];
        const NAME: &'static str = NAME;
        const VERSION: u16 = 4;
        const FORMAT_MIGRATIONS: &'static [FormatMigration] =
            &[super::v2::WIDEN, super::v4::EXTEND];
    }
}

fn key(id: u64) -> [u8; 8] {
    id.to_be_bytes()
}

/// Seed one v1 row, rewrite to v2, value-migrate to v3.
fn fixture() -> tempfile::TempDir {
    assert_eq!(v2::Row::TYPE.h(), v3::Row::TYPE.h());
    let dir = tempdir().unwrap();
    {
        let db = Db::open(dir.path()).unwrap();
        let tree = db
            .open_typed_tree::<v1::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
            .unwrap();
        tree.put(&key(1), v1::Row { value: 1 }).unwrap();
        drop(tree);
        db.close().unwrap();
    }
    {
        let db = Db::open(dir.path()).unwrap();
        db.open_typed_tree::<v2::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
            .unwrap();
        db.close().unwrap();
    }
    {
        let db = Db::open(dir.path()).unwrap();
        let migrations: &[TypedMigration<[u8; 8], v3::Row>] = &[(2, v3::bump)];
        db.open_typed_tree::<v3::Row, RapiraCodec, _>(Config::test(), NoHook, migrations)
            .unwrap();
        db.close().unwrap();
    }
    dir
}

fn child(stage: u8, target: &str, path: &std::path::Path) -> std::process::ExitStatus {
    Command::new(std::env::current_exe().unwrap())
        .arg("--exact")
        .arg("value_crash_child")
        .arg("--nocapture")
        .env("ARMDB_VALUE_CHILD", "1")
        .env("ARMDB_VALUE_DIR", path)
        .env("ARMDB_VALUE_STAGE", stage.to_string())
        .env("ARMDB_VALUE_TARGET", target)
        .status()
        .unwrap()
}

#[test]
fn value_crash_child() {
    let Ok(path) = std::env::var("ARMDB_VALUE_DIR") else {
        return;
    };
    if std::env::var("ARMDB_VALUE_CHILD").is_err() {
        return;
    }
    let stage = std::env::var("ARMDB_VALUE_STAGE").unwrap().parse().unwrap();
    failpoints::STAGE.store(stage, Ordering::SeqCst);
    let db = Db::open(path).unwrap();
    match std::env::var("ARMDB_VALUE_TARGET").unwrap().as_str() {
        "v4" => {
            db.open_typed_tree::<v4::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
                .unwrap();
        }
        "pending_v3" => {
            let tree = db
                .open_typed_tree::<v3::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
                .unwrap();
            Db::tx_debug_skip_apply(true);
            db.tx()
                .put(&tree, &key(2), &v3::Row { value: 5 })
                .unwrap()
                .durable()
                .commit()
                .unwrap();
            std::process::exit(9);
        }
        "pending_v4_kept" => {
            let tree = db
                .open_typed_tree::<v4_kept::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
                .unwrap();
            Db::tx_debug_skip_apply(true);
            db.tx()
                .put(&tree, &key(3), &v4_kept::Row { value: 9, tag: 2 })
                .unwrap()
                .durable()
                .commit()
                .unwrap();
            std::process::exit(9);
        }
        other => panic!("unknown target {other}"),
    }
}

#[test]
fn interrupted_rewrite_after_value_migration_converges() {
    let dir = fixture();
    assert_eq!(
        child(failpoints::AFTER_FORMAT_VERSION, "v4", dir.path()).code(),
        Some(9),
        "the first crash must interrupt the rewrite"
    );
    for stage in [
        failpoints::AFTER_SOURCE_RENAME,
        failpoints::BEFORE_INFO_COMMIT,
    ] {
        let status = child(stage, "v4", dir.path());
        assert!(
            matches!(status.code(), Some(9) | Some(0)),
            "stage {stage}: {status:?}"
        );
    }
    let db = Db::open(dir.path()).unwrap();
    let tree = db
        .open_typed_tree::<v4::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    assert_eq!(
        tree.get(&key(1)).map(|row| (*row).clone()),
        Some(v4::Row {
            value: 1001,
            tag: 7
        })
    );
    assert_eq!(db.db_info().collections[NAME].version, 4);
}

#[test]
fn pending_put_logged_after_a_value_migration_replays_through_the_next_rewrite() {
    for kept in [true, false] {
        for discard_pending_on_schema_mismatch in [true, false] {
            let dir = fixture();
            assert_eq!(child(0, "pending_v3", dir.path()).code(), Some(9));
            let db = Db::open_with_tx_config(
                dir.path(),
                TxConfig {
                    discard_pending_on_schema_mismatch,
                    ..TxConfig::default()
                },
            )
            .unwrap();
            assert_eq!(db.pending_transactions().len(), 1);
            let row = if kept {
                db.open_typed_tree::<v4_kept::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
                    .expect("v2 and v3 bytes are one class across the applied value step")
                    .get(&key(2))
                    .map(|row| (row.value, row.tag))
            } else {
                db.open_typed_tree::<v4::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
                    .expect("pending v3 bytes run only the v3 to v4 rewrite")
                    .get(&key(2))
                    .map(|row| (row.value, row.tag))
            };
            assert_eq!(
                row,
                Some((5, 7)),
                "kept: {kept}, discard mode: {discard_pending_on_schema_mismatch}"
            );
        }
    }
}

#[test]
fn pending_put_in_a_collection_created_at_the_current_version_replays() {
    for discard_pending_on_schema_mismatch in [true, false] {
        let dir = tempdir().unwrap();
        // The child creates the collection at v4 and crashes before apply.
        assert_eq!(child(0, "pending_v4_kept", dir.path()).code(), Some(9));
        let db = Db::open_with_tx_config(
            dir.path(),
            TxConfig {
                discard_pending_on_schema_mismatch,
                ..TxConfig::default()
            },
        )
        .unwrap();
        assert_eq!(db.pending_transactions().len(), 1);
        let row = db
            .open_typed_tree::<v4_kept::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
            .expect("declarations below the creation version cannot block the reopen")
            .get(&key(3))
            .map(|row| (row.value, row.tag));
        assert_eq!(
            row,
            Some((9, 2)),
            "discard mode: {discard_pending_on_schema_mismatch}"
        );
    }
}

#[test]
fn pending_put_in_a_collection_created_after_the_value_step_replays_through_the_next_rewrite() {
    for discard_pending_on_schema_mismatch in [true, false] {
        let dir = tempdir().unwrap();
        // Created at v3, whose fingerprint the retained chain also names at v2.
        assert_eq!(child(0, "pending_v3", dir.path()).code(), Some(9));
        let db = Db::open_with_tx_config(
            dir.path(),
            TxConfig {
                discard_pending_on_schema_mismatch,
                ..TxConfig::default()
            },
        )
        .unwrap();
        assert_eq!(db.pending_transactions().len(), 1);
        let row = db
            .open_typed_tree::<v4_kept::Row, RapiraCodec, _>(Config::test(), NoHook, &[])
            .expect("v2 predates the collection, so it shares v3's byte class")
            .get(&key(2))
            .map(|row| (row.value, row.tag));
        assert_eq!(
            row,
            Some((5, 7)),
            "discard mode: {discard_pending_on_schema_mismatch}"
        );
    }
}