armdb 0.14.2

//! Public transaction-scoped bulk deletion.
#![cfg(all(
    feature = "armour",
    feature = "typed-tree",
    feature = "rapira-codec",
    feature = "var-collections"
))]

use std::process::Command;
use std::sync::Arc;

use armdb::armour::Db;
use armdb::{
    CollectionMeta, Config, CounterTargets, DbError, FixedConfig, IndexTargets, NoHook,
    OnMissingTarget, Progress, Projected, RapiraCodec, Selection, TtlHook, TtlPolicy, TxConfig,
    TxScope, Upsert,
};
use armour_core::{GetType, KeyScheme, KeyType};
use rapira::Rapira;
use zerocopy::{FromBytes, Immutable, IntoBytes};

#[cfg(feature = "test-faults")]
use armdb::txn::failpoints;

#[derive(
    Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, FromBytes, Immutable, IntoBytes,
)]
#[repr(C)]
struct Id {
    prefix: u64,
    row: u64,
}

armdb::impl_key_zerocopy!(Id, KeyScheme::Typed(&[KeyType::U64, KeyType::U64]));

fn key(prefix: u64, row: u64) -> Id {
    Id { prefix, row }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, GetType, Rapira)]
struct Row(u64);

impl CollectionMeta for Row {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_rows";
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, GetType, Rapira)]
struct TtlRow {
    expires_at_ms: u64,
}

impl CollectionMeta for TtlRow {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_ttl_rows";
    const TTL: Option<TtlPolicy<Self>> = Some(TtlPolicy {
        revision: 1,
        deadline: |row, _| Some(row.expires_at_ms),
    });
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, GetType, Rapira)]
struct MapRow(u64);
impl CollectionMeta for MapRow {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_map_rows";
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, GetType, Rapira)]
struct VarTreeRow(u64);
impl CollectionMeta for VarTreeRow {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_var_tree_rows";
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, GetType, Rapira)]
struct VarMapRow(u64);
impl CollectionMeta for VarMapRow {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_var_map_rows";
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, GetType, FromBytes, Immutable, IntoBytes)]
#[repr(C)]
struct ZeroTreeRow(u64);
impl CollectionMeta for ZeroTreeRow {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_zero_tree_rows";
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, GetType, FromBytes, Immutable, IntoBytes)]
#[repr(C)]
struct ZeroMapRow(u64);
impl CollectionMeta for ZeroMapRow {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_zero_map_rows";
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, GetType, FromBytes, Immutable, IntoBytes)]
#[repr(C)]
struct FixedZeroRow(u64);
impl CollectionMeta for FixedZeroRow {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_fixed_zero_rows";
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, GetType, Rapira)]
struct DerivedRow {
    group: [u8; 8],
}

impl CollectionMeta for DerivedRow {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_derived_rows";
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, GetType, FromBytes, Immutable, IntoBytes)]
#[repr(C)]
struct DerivedIndexRow(u64);

impl CollectionMeta for DerivedIndexRow {
    type SelfId = Id;
    const NAME: &'static str = "bulk_tx_derived_index";
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, GetType, Rapira)]
struct DerivedCount {
    rows: u64,
}

impl CollectionMeta for DerivedCount {
    type SelfId = [u8; 8];
    const NAME: &'static str = "bulk_tx_derived_counts";
}

type DerivedHook = Projected<Id, DerivedRow>;
type DerivedSource = armdb::typed_tree!(DerivedRow, RapiraCodec, DerivedHook);

fn derived_hook(
    index: &Arc<armdb::zero_tree!(DerivedIndexRow)>,
    counts: &Arc<armdb::typed_tree!(DerivedCount, RapiraCodec)>,
) -> DerivedHook {
    let index_for_projection = Arc::clone(index);
    let counts_for_projection = Arc::clone(counts);
    Projected::<Id, DerivedRow, NoHook>::new(
        IndexTargets::new().with_index(index),
        move |key, _row, sink| sink.row(&index_for_projection, *key, DerivedIndexRow(1)),
    )
    .and_counters(
        CounterTargets::new().with_counter(
            counts,
            "rows",
            OnMissingTarget::Error,
            |count: &DerivedCount| count.rows,
            |_current, rows| Upsert::Put(DerivedCount { rows }),
        ),
        move |_key, row, sink| sink.add(&counts_for_projection, row.group, 1),
    )
}

fn open_derived_source(
    db: &Arc<Db>,
) -> (
    Arc<DerivedSource>,
    Arc<armdb::zero_tree!(DerivedIndexRow)>,
    Arc<armdb::typed_tree!(DerivedCount, RapiraCodec)>,
) {
    let index = db
        .open_zero_tree::<DerivedIndexRow, { std::mem::size_of::<DerivedIndexRow>() }, _>(
            Config::test(),
            NoHook,
            &[],
        )
        .unwrap();
    let counts = db
        .open_typed_tree::<DerivedCount, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    let group = [7; 8];
    counts.insert(&group, DerivedCount { rows: 0 }).unwrap();
    let source = db
        .open_typed_tree::<DerivedRow, RapiraCodec, _>(
            Config::test(),
            derived_hook(&index, &counts),
            &[],
        )
        .unwrap();
    (source, index, counts)
}

fn seed_derived(source: &DerivedSource, rows: u64) {
    for row in 0..rows {
        source
            .insert(&key(7, row), DerivedRow { group: [7; 8] })
            .unwrap();
    }
}

fn open_db() -> (tempfile::TempDir, Arc<Db>) {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    (dir, db)
}

fn seed_typed(rows: &armdb::TypedTree<Id, Row, RapiraCodec, TtlHook<Row, NoHook>>) {
    for row in 0..4 {
        rows.insert(&key(7, row), Row(row)).unwrap();
    }
}

/// The break this catches: one of the promised bulk public types stops being
/// nameable by a downstream crate because it is no longer publicly re-exported
/// or `TxScope`'s generic shape changes incompatibly.
#[test]
fn bulk_public_surface_is_nameable() {
    let _progress = Progress {
        deleted: 0,
        done: false,
    };
    let _selection = Selection {
        queued: 0,
        done: false,
    };

    fn accept_scope<S, K: armdb::Key, const N: usize>(_scope: &TxScope<S, K, N>) {}

    let _ = accept_scope::<(), [u8; 8], 1>;
}

/// The break this catches: a scoped delete writes eagerly, ignores its limit,
/// or incorrectly reports a range truncated by the requested page size as done.
#[test]
fn bulk_tx_scope_queues_a_bounded_typed_range() {
    let (_dir, db) = open_db();
    let rows = db
        .open_typed_tree::<Row, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    seed_typed(&rows);

    let scope = rows.tx_scope::<1>(&key(7, 0));
    let (_tx, selection) = db.tx().delete_scope(&scope, 0).unwrap();
    assert_eq!(
        selection,
        Selection {
            queued: 0,
            done: false
        }
    );
    assert_eq!(rows.scope::<1>(&key(7, 0)).count().unwrap(), 4);

    let (tx, selection) = db.tx().delete_scope(&scope, 2).unwrap();
    assert_eq!(
        selection,
        Selection {
            queued: 2,
            done: false
        }
    );
    assert_eq!(rows.scope::<1>(&key(7, 0)).count().unwrap(), 4);

    tx.durable().commit().unwrap();
    assert_eq!(rows.scope::<1>(&key(7, 0)).count().unwrap(), 2);

    let (tx, selection) = db.tx().delete_scope(&scope, 5).unwrap();
    assert_eq!(
        selection,
        Selection {
            queued: 2,
            done: true
        }
    );
    tx.durable().commit().unwrap();
    assert_eq!(rows.scope::<1>(&key(7, 0)).count().unwrap(), 0);

    let (_tx, selection) = db.tx().delete_scope(&scope, 5).unwrap();
    assert_eq!(
        selection,
        Selection {
            queued: 0,
            done: true
        }
    );
}

/// The break this catches: a key-only count decodes or sweeps an expired row,
/// causing an expired-but-unswept source key to disappear from the result.
#[test]
fn bulk_count_includes_an_expired_unswept_ttl_key() {
    let (_dir, db) = open_db();
    let rows = db
        .open_typed_tree::<TtlRow, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    rows.insert(&key(7, 1), TtlRow { expires_at_ms: 1 })
        .unwrap();

    assert_eq!(rows.scope::<1>(&key(7, 0)).count().unwrap(), 1);
    rows.sweep_expired(2, 1).unwrap();
    assert_eq!(rows.scope::<1>(&key(7, 0)).count().unwrap(), 0);
}

/// The break this catches: selection uses the caller's limit without reserving
/// room for already queued source operations, or writes while building a Tx.
#[test]
fn bulk_tx_scope_stops_at_remaining_source_op_capacity() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(
        Db::builder(dir.path())
            .tx_config(TxConfig {
                max_ops: 2,
                ..TxConfig::default()
            })
            .open()
            .unwrap(),
    );
    let rows = db
        .open_typed_tree::<Row, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    seed_typed(&rows);

    let scope = rows.tx_scope::<1>(&key(7, 0));
    let tx = db.tx().delete(&rows, &key(9, 0)).unwrap();
    let (tx, selection) = tx.delete_scope(&scope, 5).unwrap();
    assert_eq!(
        selection,
        Selection {
            queued: 1,
            done: false
        }
    );
    assert_eq!(rows.scope::<1>(&key(7, 0)).count().unwrap(), 4);

    tx.commit().unwrap();
    assert_eq!(rows.scope::<1>(&key(7, 0)).count().unwrap(), 3);
}

/// The break this catches: `delete_scope` accepts a source-only page whose
/// index and counter rows overflow final admission, appends a log record before
/// rejecting it, or cannot be safely rebuilt at a smaller page size.
#[test]
fn bulk_tx_derived_admission_failure_writes_nothing_then_halved_retry_commits() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(
        Db::builder(dir.path())
            .tx_config(TxConfig {
                // Five source deletes fit builder admission. Their five index
                // deletes plus one collapsed counter write do not fit final
                // derived-row admission; two source deletes do (2 + 2 + 1).
                max_ops: 5,
                ..TxConfig::default()
            })
            .open()
            .unwrap(),
    );
    let (source, index, counts) = open_derived_source(&db);
    seed_derived(&source, 6);
    let scope = source.tx_scope::<1>(&key(7, 0));

    let (too_large, selected) = db.tx().delete_scope(&scope, 5).unwrap();
    assert_eq!(
        selected,
        Selection {
            queued: 5,
            done: false
        }
    );
    assert!(matches!(
        too_large.durable().commit(),
        Err(DbError::Client(
            "transaction exceeds max_ops (including derived secondary-index, counter and queue rows)"
        ))
    ));
    assert_eq!(source.scope::<1>(&key(7, 0)).count().unwrap(), 6);
    assert_eq!(index.scope::<1>(&key(7, 0)).count().unwrap(), 6);
    assert_eq!(
        counts.get(&[7; 8]).as_deref(),
        Some(&DerivedCount { rows: 6 })
    );

    let (retry, selected) = db.tx().delete_scope(&scope, 2).unwrap();
    assert_eq!(
        selected,
        Selection {
            queued: 2,
            done: false
        }
    );
    retry.durable().commit().unwrap();
    assert_eq!(source.scope::<1>(&key(7, 0)).count().unwrap(), 4);
    assert_eq!(index.scope::<1>(&key(7, 0)).count().unwrap(), 4);
    assert_eq!(
        counts.get(&[7; 8]).as_deref(),
        Some(&DerivedCount { rows: 4 })
    );
}

fn run_crash_child(dir: &std::path::Path) -> std::process::ExitStatus {
    Command::new(std::env::current_exe().unwrap())
        .arg("--exact")
        .arg("bulk_tx_crash_child_entry")
        .arg("--nocapture")
        .env("ARMDB_BULK_TX_CRASH_DIR", dir)
        .env("ARMDB_BULK_TX_CRASH_CHILD", "1")
        .status()
        .unwrap()
}

#[test]
fn bulk_tx_crash_child_entry() {
    let Ok(dir) = std::env::var("ARMDB_BULK_TX_CRASH_DIR") else {
        return;
    };
    if std::env::var("ARMDB_BULK_TX_CRASH_CHILD").is_err() {
        return;
    }

    let db = Arc::new(Db::open(&dir).unwrap());
    let rows = db
        .open_typed_tree::<Row, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    let scope = rows.tx_scope::<1>(&key(7, 0)).asc();
    let (first_chunk, selected) = db.tx().delete_scope(&scope, 2).unwrap();
    assert_eq!(selected.queued, 2);

    // Under `just faults`, exit immediately after the first chunk's durable log
    // append. The parent must reopen and replay it; a default-feature focused
    // run still exercises the bounded child entry via the hard exit below.
    #[cfg(feature = "test-faults")]
    failpoints::STAGE.store(
        failpoints::AFTER_LOG_APPEND,
        std::sync::atomic::Ordering::SeqCst,
    );
    first_chunk.durable().commit().unwrap();

    // The second chunk is deliberately never built or committed. `exit` skips
    // destructors, so the parent distinguishes the committed log record from a
    // graceful `Db` shutdown before reopening the directory.
    std::process::exit(9);
}

/// The break this catches: a committed bulk chunk is not replayed after a hard
/// process exit, or a later chunk becomes durable without its own commit.
#[test]
fn bulk_tx_crash_between_committed_chunks_replays_the_first_chunk_only() {
    let dir = tempfile::tempdir().unwrap();
    {
        let db = Arc::new(Db::open(dir.path()).unwrap());
        let rows = db
            .open_typed_tree::<Row, RapiraCodec, _>(Config::test(), NoHook, &[])
            .unwrap();
        seed_typed(&rows);
    }

    assert_eq!(run_crash_child(dir.path()).code(), Some(9));

    let db = Arc::new(Db::open(dir.path()).unwrap());
    let rows = db
        .open_typed_tree::<Row, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    assert_eq!(rows.scope::<1>(&key(7, 0)).count().unwrap(), 2);
    assert!(rows.get(&key(7, 0)).is_none());
    assert!(rows.get(&key(7, 1)).is_none());
    assert_eq!(rows.get(&key(7, 2)).as_deref(), Some(&Row(2)));
    assert_eq!(rows.get(&key(7, 3)).as_deref(), Some(&Row(3)));
}

/// The break this catches: any managed family queues eagerly, loses its durable
/// commit path, or exposes `tx_scope` without actually deleting the selected
/// rows after commit.
#[test]
fn bulk_tx_scope_commits_for_each_managed_family() {
    let (_dir, db) = open_db();
    let mut iterable_config = Config::test();
    iterable_config.iterable = true;
    let typed_map = db
        .open_typed_map::<MapRow, RapiraCodec, _>(iterable_config.clone(), NoHook, &[])
        .unwrap();
    let var_tree = db
        .open_var_typed_tree::<VarTreeRow, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    let var_map = db
        .open_var_typed_map::<VarMapRow, RapiraCodec, _>(iterable_config.clone(), NoHook, &[])
        .unwrap();
    let zero_tree = db
        .open_zero_tree::<ZeroTreeRow, { std::mem::size_of::<ZeroTreeRow>() }, _>(
            Config::test(),
            NoHook,
            &[],
        )
        .unwrap();
    let zero_map = db
        .open_zero_map::<ZeroMapRow, { std::mem::size_of::<ZeroMapRow>() }, _>(
            iterable_config,
            NoHook,
            &[],
        )
        .unwrap();
    let fixed_zero = db
        .open_zero_tree_fixed::<FixedZeroRow, { std::mem::size_of::<FixedZeroRow>() }, _>(
            FixedConfig::test(),
            NoHook,
            &[],
        )
        .unwrap();

    macro_rules! assert_tx_scope_commit {
        ($seed:expr, $tx_scope:expr, $scope:expr) => {{
            for row in 0..4 {
                $seed(row);
            }
            let tx_scope = $tx_scope;
            let (tx, selection) = db.tx().delete_scope(&tx_scope, 2).unwrap();
            assert_eq!(
                selection,
                Selection {
                    queued: 2,
                    done: false
                }
            );
            assert_eq!($scope.count().unwrap(), 4);
            tx.durable().commit().unwrap();
            assert_eq!($scope.count().unwrap(), 2);
        }};
    }

    assert_tx_scope_commit!(
        |row| typed_map.insert(&key(7, row), MapRow(row)).unwrap(),
        typed_map.tx_scope::<1>(&key(7, 0)).unwrap(),
        typed_map.scope::<1>(&key(7, 0)).unwrap()
    );
    assert_tx_scope_commit!(
        |row| var_tree.insert(&key(7, row), &VarTreeRow(row)).unwrap(),
        var_tree.tx_scope::<1>(&key(7, 0)),
        var_tree.scope::<1>(&key(7, 0))
    );
    assert_tx_scope_commit!(
        |row| var_map.insert(&key(7, row), &VarMapRow(row)).unwrap(),
        var_map.tx_scope::<1>(&key(7, 0)).unwrap(),
        var_map.scope::<1>(&key(7, 0)).unwrap()
    );
    assert_tx_scope_commit!(
        |row| zero_tree.insert(&key(7, row), &ZeroTreeRow(row)).unwrap(),
        zero_tree.tx_scope::<1>(&key(7, 0)),
        zero_tree.scope::<1>(&key(7, 0))
    );
    assert_tx_scope_commit!(
        |row| zero_map.insert(&key(7, row), &ZeroMapRow(row)).unwrap(),
        zero_map.tx_scope::<1>(&key(7, 0)).unwrap(),
        zero_map.scope::<1>(&key(7, 0)).unwrap()
    );
    assert_tx_scope_commit!(
        |row| fixed_zero.insert(&key(7, row), &FixedZeroRow(row)).unwrap(),
        fixed_zero.tx_scope::<1>(&key(7, 0)),
        fixed_zero.scope::<1>(&key(7, 0))
    );
}