armdb 0.11.2

//! `Tx::insert` — insert-if-absent inside a crash-atomic transaction (#139).
//!
//! Spec: `docs/superpowers/specs/26-08-10-tx-insert.md`.
#![cfg(all(
    feature = "armour",
    feature = "typed-tree",
    feature = "rapira-codec",
    feature = "var-collections"
))]

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

#[cfg(feature = "encryption")]
use std::path::Path;

#[cfg(feature = "encryption")]
use armdb::armour::Collection;
use armdb::armour::Db;
use armdb::{
    Codec, CollectionMeta, Config, DbError, NoHook, RapiraCodec, TtlHook, TxConfig, TypedMap,
    VarTypedTree,
};
use armour_core::GetType;
use rapira::Rapira;

type Id = [u8; 8];

fn id(n: u64) -> Id {
    n.to_be_bytes()
}

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

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

type Claims = Arc<TypedMap<Id, Claim, RapiraCodec, TtlHook<Claim, NoHook>>>;
type Audits = Arc<TypedMap<Id, Audit, RapiraCodec, TtlHook<Audit, NoHook>>>;

fn open_pair(db: &Arc<Db>) -> (Claims, Audits) {
    let claims = db
        .open_typed_map::<Claim, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    let audits = db
        .open_typed_map::<Audit, RapiraCodec, _>(Config::test(), NoHook, &[])
        .unwrap();
    (claims, audits)
}

/// The happy path: a free key is claimed, and a second collection in the same
/// transaction lands with it.
#[test]
fn insert_writes_an_absent_key() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let (claims, audits) = open_pair(&db);

    db.tx()
        .insert(&claims, &id(1), &Claim { owner: 7 })
        .unwrap()
        .insert(&audits, &id(1), &Audit { at: 99 })
        .unwrap()
        .commit()
        .unwrap();

    assert_eq!(claims.get(&id(1)).map(|c| c.owner), Some(7));
    assert_eq!(audits.get(&id(1)).map(|a| a.at), Some(99));
}

/// The defining failure: an occupied key aborts the whole transaction **before
/// the log point**. Nothing is written, no shard is fenced, and — unlike a
/// post-log failure — the error is raw and the log holds no pending record.
#[test]
fn insert_on_an_occupied_key_fails_before_the_log_point() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let (claims, audits) = open_pair(&db);

    claims.insert(&id(1), Claim { owner: 1 }).unwrap();

    let err = db
        .tx()
        .insert(&claims, &id(1), &Claim { owner: 2 })
        .unwrap()
        .insert(&audits, &id(1), &Audit { at: 5 })
        .unwrap()
        .commit()
        .unwrap_err();

    assert!(matches!(err, DbError::KeyExists), "{err}");
    assert!(
        !matches!(err, DbError::TxFailedAfterLog { .. }),
        "a pre-log failure must not be wrapped: {err}"
    );
    assert!(
        db.pending_transactions().is_empty(),
        "nothing reached the log"
    );
    assert_eq!(
        claims.get(&id(1)).map(|c| c.owner),
        Some(1),
        "the occupied row keeps its original value"
    );
    assert!(
        audits.get(&id(1)).is_none(),
        "the other collection in the failed transaction must not be written"
    );

    // The shard is not fenced: an ordinary write still goes through.
    claims.put(&id(1), Claim { owner: 3 }).unwrap();
    assert_eq!(claims.get(&id(1)).map(|c| c.owner), Some(3));
}

/// `insert` sees the transaction's own pending writes before storage, so a
/// `put` earlier in the same transaction already occupies the key.
#[test]
fn insert_sees_an_earlier_put_in_the_same_transaction() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let (claims, _audits) = open_pair(&db);

    let err = db
        .tx()
        .put(&claims, &id(1), &Claim { owner: 1 })
        .unwrap()
        .insert(&claims, &id(1), &Claim { owner: 2 })
        .unwrap()
        .commit()
        .unwrap_err();

    assert!(matches!(err, DbError::KeyExists), "{err}");
    assert!(
        claims.get(&id(1)).is_none(),
        "the whole transaction aborted"
    );
}

/// The other direction: a delete earlier in the same transaction frees the key,
/// so delete-then-insert is legal in one commit.
#[test]
fn insert_after_a_delete_in_the_same_transaction_succeeds() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let (claims, _audits) = open_pair(&db);

    claims.insert(&id(1), Claim { owner: 1 }).unwrap();

    db.tx()
        .delete(&claims, &id(1))
        .unwrap()
        .insert(&claims, &id(1), &Claim { owner: 2 })
        .unwrap()
        .commit()
        .unwrap();

    assert_eq!(claims.get(&id(1)).map(|c| c.owner), Some(2));
}

/// Two inserts on one key inside a transaction: the first one's pending write
/// occupies the key for the second.
#[test]
fn a_second_insert_on_the_same_key_fails() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let (claims, _audits) = open_pair(&db);

    let err = db
        .tx()
        .insert(&claims, &id(1), &Claim { owner: 1 })
        .unwrap()
        .insert(&claims, &id(1), &Claim { owner: 2 })
        .unwrap()
        .commit()
        .unwrap_err();

    assert!(matches!(err, DbError::KeyExists), "{err}");
    assert!(
        claims.get(&id(1)).is_none(),
        "the whole transaction aborted"
    );
}

/// A `put` after an `insert` on the same key collapses onto the same plan slot
/// and wins — `insert` claims the key, `put` decides its final value.
#[test]
fn a_put_after_an_insert_on_the_same_key_wins() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let (claims, _audits) = open_pair(&db);

    db.tx()
        .insert(&claims, &id(1), &Claim { owner: 1 })
        .unwrap()
        .put(&claims, &id(1), &Claim { owner: 2 })
        .unwrap()
        .commit()
        .unwrap();

    assert_eq!(claims.get(&id(1)).map(|c| c.owner), Some(2));
}

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

type Docs = Arc<VarTypedTree<Id, Doc, RapiraCodec, TtlHook<Doc, NoHook>>>;

fn open_docs(db: &Arc<Db>) -> Docs {
    db.open_var_typed_tree::<Doc, RapiraCodec, NoHook>(Config::test(), NoHook, &[])
        .unwrap()
}

static FAIL_EXISTING_ENCODE: AtomicBool = AtomicBool::new(false);

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

#[derive(Clone, Default)]
struct ToggleEncodeCodec;

impl Codec<EncodeRow> for ToggleEncodeCodec {
    fn encode_to(&self, value: &EncodeRow, out: &mut Vec<u8>) -> armdb::DbResult<()> {
        if FAIL_EXISTING_ENCODE.load(Ordering::SeqCst) && value.value == 7 {
            return Err(DbError::CorruptedEntry { offset: 0 });
        }
        out.clear();
        out.extend_from_slice(&value.value.to_le_bytes());
        Ok(())
    }

    fn decode_from(&self, bytes: &[u8]) -> armdb::DbResult<EncodeRow> {
        let bytes: [u8; 8] = bytes
            .try_into()
            .map_err(|_| DbError::CorruptedEntry { offset: 0 })?;
        Ok(EncodeRow {
            value: u64::from_le_bytes(bytes),
        })
    }

    fn size(&self, _value: &EncodeRow) -> usize {
        8
    }
}

type EncodeRows = Arc<TypedMap<Id, EncodeRow, ToggleEncodeCodec, TtlHook<EncodeRow, NoHook>>>;

fn open_encode_rows(db: &Arc<Db>) -> EncodeRows {
    db.open_typed_map::<EncodeRow, ToggleEncodeCodec, _>(Config::test(), NoHook, &[])
        .unwrap()
}

/// An occupied typed row must reject by presence even if the codec can no
/// longer re-encode the resident value for the transaction's raw read.
#[test]
fn insert_rejects_an_occupied_typed_row_without_reencoding_it() {
    FAIL_EXISTING_ENCODE.store(false, Ordering::SeqCst);
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let rows = open_encode_rows(&db);
    rows.put(&id(1), EncodeRow { value: 7 }).unwrap();
    FAIL_EXISTING_ENCODE.store(true, Ordering::SeqCst);

    let err = db
        .tx()
        .insert(&rows, &id(1), &EncodeRow { value: 9 })
        .unwrap()
        .commit()
        .unwrap_err();

    FAIL_EXISTING_ENCODE.store(false, Ordering::SeqCst);
    assert!(matches!(err, DbError::KeyExists), "{err}");
    assert_eq!(rows.get(&id(1)).map(|row| row.value), Some(7));
    assert!(db.pending_transactions().is_empty());
}

#[cfg(feature = "encryption")]
fn corrupt_variable_tags(collection_path: &Path) {
    let mut found = false;
    for shard in std::fs::read_dir(collection_path).unwrap() {
        let shard = shard.unwrap().path();
        if !shard.is_dir() {
            continue;
        }
        for file in std::fs::read_dir(shard).unwrap() {
            let path = file.unwrap().path();
            if path.extension().and_then(|ext| ext.to_str()) != Some("tags") {
                continue;
            }
            let mut bytes = std::fs::read(&path).unwrap();
            if bytes.is_empty() {
                continue;
            }
            for byte in &mut bytes {
                *byte ^= 0xFF;
            }
            std::fs::write(path, bytes).unwrap();
            found = true;
        }
    }
    assert!(found, "no non-empty encrypted tag file found");
}

#[cfg(feature = "encryption")]
#[test]
fn insert_rejects_an_occupied_variable_row_when_payload_read_fails() {
    const DATA_KEY: [u8; 32] = [0x71; 32];
    const LOG_KEY: [u8; 32] = [0x72; 32];

    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(
        Db::open_with_tx_config(
            dir.path(),
            TxConfig {
                log_encryption_key: Some(LOG_KEY),
                ..TxConfig::default()
            },
        )
        .unwrap(),
    );
    let mut config = Config::test();
    config.encryption_key = Some(DATA_KEY);
    let docs = db
        .open_var_typed_tree::<Doc, RapiraCodec, _>(config, NoHook, &[])
        .unwrap();
    docs.put(&id(1), &Doc { size: 42 }).unwrap();
    docs.flush().unwrap();
    assert!(docs.try_get(&id(1)).is_ok());
    corrupt_variable_tags(&db.tree_path(Doc::NAME));
    assert!(matches!(
        docs.try_get(&id(1)),
        Err(DbError::EncryptionError(_))
    ));

    let err = db
        .tx()
        .insert(&docs, &id(1), &Doc { size: 99 })
        .unwrap()
        .commit()
        .unwrap_err();
    assert!(matches!(err, DbError::KeyExists), "{err}");
    assert!(db.pending_transactions().is_empty());
}

/// `insert` decides on raw bytes and never decodes the stored value, so an
/// occupied key whose value no longer decodes is still `KeyExists` — the same
/// answer a collection-level `insert` gives, since it rejects on the index
/// before any codec runs.
///
/// The second half is the control: the `guard` + `put` form people reach for
/// instead **does** decode, and fails with the codec's error. The two forms
/// diverge here on purpose; both halves are pinned so neither can drift.
#[test]
fn insert_rejects_an_undecodable_occupied_row_without_decoding_it() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let docs = open_docs(&db);

    // One byte where the codec expects a `u64` — present, but undecodable.
    // NoHook means the write-path adapter does not decode, so the raw bytes
    // land on disk as-is.
    docs.as_inner().put(&id(1), &[0x01]).unwrap();
    assert!(
        matches!(docs.try_get(&id(1)), Err(DbError::CorruptedEntry { .. })),
        "the fixture must actually be undecodable"
    );

    let err = db
        .tx()
        .insert(&docs, &id(1), &Doc { size: 42 })
        .unwrap()
        .commit()
        .unwrap_err();
    assert!(
        matches!(err, DbError::KeyExists),
        "insert checks existence, not decodability: {err}"
    );

    // Control: the guard form decodes first and therefore reports corruption.
    let guard_err = db
        .tx()
        .guard(&docs, &id(1), |cur: Option<&Doc>| {
            if cur.is_none() {
                Ok(())
            } else {
                Err(DbError::KeyExists)
            }
        })
        .unwrap()
        .put(&docs, &id(1), &Doc { size: 42 })
        .unwrap()
        .commit()
        .unwrap_err();
    assert!(
        matches!(guard_err, DbError::CorruptedEntry { .. }),
        "guard decodes before its closure, so it reports the codec error: {guard_err}"
    );

    // Neither attempt wrote anything.
    assert!(matches!(
        docs.try_get(&id(1)),
        Err(DbError::CorruptedEntry { .. })
    ));
    assert!(db.pending_transactions().is_empty());
}

/// Two transactions racing for the same free key: exactly one commits, the
/// other gets `KeyExists`, and the stored value is the winner's. This is what
/// "the check runs under the shard lock" means operationally — a builder-time
/// check would let both through.
#[test]
fn two_transactions_racing_for_one_key_produce_exactly_one_winner() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(Db::open_test(dir.path()).unwrap());
    let (claims, _audits) = open_pair(&db);

    let barrier = std::sync::Barrier::new(2);
    let outcomes: Vec<(u32, Result<(), DbError>)> = std::thread::scope(|s| {
        let handles: Vec<_> = [10u32, 20u32]
            .into_iter()
            .map(|owner| {
                let db = Arc::clone(&db);
                let claims = Arc::clone(&claims);
                let barrier = &barrier;
                s.spawn(move || {
                    let tx = db
                        .tx()
                        .insert(&claims, &id(1), &Claim { owner })
                        .expect("builder accepts the op");
                    barrier.wait();
                    (owner, tx.commit())
                })
            })
            .collect();
        handles.into_iter().map(|h| h.join().unwrap()).collect()
    });

    let winners: Vec<u32> = outcomes
        .iter()
        .filter(|(_, r)| r.is_ok())
        .map(|(o, _)| *o)
        .collect();
    let losers: Vec<&DbError> = outcomes
        .iter()
        .filter_map(|(_, r)| r.as_ref().err())
        .collect();

    assert_eq!(
        winners.len(),
        1,
        "exactly one commit may succeed: {outcomes:?}"
    );
    assert_eq!(losers.len(), 1, "the other must fail: {outcomes:?}");
    assert!(matches!(losers[0], DbError::KeyExists), "{:?}", losers[0]);
    assert_eq!(
        claims.get(&id(1)).map(|c| c.owner),
        Some(winners[0]),
        "the stored value belongs to the transaction that won"
    );
}

/// `insert` is one operation against `max_ops`, exactly like `put` — the
/// property that a `guard` + `put` desugaring would not have.
#[test]
fn insert_costs_one_op_against_max_ops_just_like_put() {
    let dir = tempfile::tempdir().unwrap();
    let db = Arc::new(
        Db::builder(dir.path())
            .for_tests()
            .tx_config(TxConfig {
                max_ops: 2,
                ..TxConfig::default()
            })
            .open()
            .unwrap(),
    );
    let (claims, _audits) = open_pair(&db);

    // Two inserts fit; the third is refused by the builder, before any lock.
    let tx = db
        .tx()
        .insert(&claims, &id(1), &Claim { owner: 1 })
        .unwrap()
        .insert(&claims, &id(2), &Claim { owner: 2 })
        .unwrap();
    let err = tx
        .insert(&claims, &id(3), &Claim { owner: 3 })
        .err()
        .expect("third insert must exceed max_ops");
    assert!(
        matches!(err, DbError::Client("transaction exceeds max_ops")),
        "{err}"
    );

    // `put` refuses at the very same call index — that is the parity claim.
    let tx = db
        .tx()
        .put(&claims, &id(1), &Claim { owner: 1 })
        .unwrap()
        .put(&claims, &id(2), &Claim { owner: 2 })
        .unwrap();
    let err = tx
        .put(&claims, &id(3), &Claim { owner: 3 })
        .err()
        .expect("third put must exceed max_ops");
    assert!(
        matches!(err, DbError::Client("transaction exceeds max_ops")),
        "{err}"
    );

    // And two inserts really do commit under that limit.
    db.tx()
        .insert(&claims, &id(1), &Claim { owner: 1 })
        .unwrap()
        .insert(&claims, &id(2), &Claim { owner: 2 })
        .unwrap()
        .commit()
        .unwrap();
    assert_eq!(claims.get(&id(2)).map(|c| c.owner), Some(2));
}