armdb 0.11.1

sharded bitcask key-value storage optimized for NVMe
Documentation
//! Replay of pending transaction records into a collection that is being
//! opened, before its migrations run.

use crate::armour::Db;
use crate::error::{DbError, DbResult};
use crate::txlog::record::{OpKind, TxOp};
use crate::txn::participant::TxParticipant;

pub(crate) fn replay_pending<S: TxParticipant + ?Sized>(
    db: &Db,
    coll: &S,
    stored_fingerprint: Option<u64>,
) -> DbResult<usize> {
    let pending = db.txlog().pending();
    if pending.is_empty() {
        return Ok(0);
    }
    let name = coll.name();
    let mut applied_ops = 0usize;

    for record in pending {
        let mine: Vec<&TxOp> = record.ops.iter().filter(|o| o.name == name).collect();
        if mine.is_empty() {
            continue;
        }

        let mut shards: Vec<usize> = Vec::new();
        for op in &mine {
            let sid = coll.shard_of(&op.key)?;
            if !shards.contains(&sid) {
                shards.push(sid);
            }
        }
        shards.sort_unstable();
        let mut locked = coll.lock_shards(&shards)?;
        let mut discard = false;
        for op in &mine {
            if !locked.needs_replay(&op.key, op.kind, op.version)? {
                continue;
            }
            // The record holds bytes encoded by the schema live at commit time.
            // Replay runs before migrations, so the collection is still in its
            // stored shape — the two must agree.
            //
            // Asked per operation, and only once the operation is known to need
            // replay. A record can outlive the migration that replayed it: a
            // participant that never opens keeps it pending, the collection
            // migrates, and every later open would then be refused over bytes
            // nobody is going to decode. The gate guards *decoding*, so it
            // belongs where the decoding would happen.
            if let Some(stored) = stored_fingerprint
                && op.schema_fingerprint != stored
            {
                if !db.tx_config().discard_pending_on_schema_mismatch {
                    return Err(DbError::TxSchemaMismatch {
                        collection: name,
                        tx_id: record.tx_id,
                        stored,
                        logged: op.schema_fingerprint,
                    });
                }
                tracing::error!(
                    collection = name,
                    tx_id = record.tx_id,
                    "dropping a pending transaction: schema fingerprint changed"
                );
                // Stop this collection's siblings from replaying the record too.
                //
                // Only forward, though: collections open one at a time, so a
                // sibling opened *earlier* has already replayed its share, and
                // `mark_applied` then fixes a partially applied transaction in
                // place. That is the bargain this flag is — it trades atomicity
                // for the ability to start at all, which is why it is off by
                // default and why the fail-closed path above is the norm.
                discard = true;
                break;
            }
            // The event count is not used here: replay fires every buffered
            // event in order after the loop, it does not interleave them.
            let _events = match op.kind {
                OpKind::Put => locked.apply_put(&op.key, &op.value, op.version)?,
                OpKind::Delete => locked.apply_delete(&op.key, op.version)?,
            };
            applied_ops += 1;
        }
        // Replay never marks itself logged, so it never fences; the sink stays
        // empty and is discarded.
        // Replay has no live subscribers to wake: it runs at open, before the
        // collection is published.
        let mut deferred = locked.finish(
            &mut Vec::new(),
            &mut crate::watch::WatchWakeBatch::default(),
        )?;
        for i in 0..deferred.event_count() {
            deferred.replay_one(i);
        }
        if discard {
            db.txlog().mark_applied(record.tx_id);
        }
    }
    Ok(applied_ops)
}

#[cfg(test)]
#[cfg(all(feature = "armour", feature = "rapira-codec"))]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::Ordering;

    use armour_core::GetType;
    use rapira::Rapira;

    use crate::armour::{Db, DbInfo};
    use crate::txn::commit_tests::{Counter, User, UserId, open_pair, user_id};
    use crate::{CollectionMeta, Config, DbResult, NoHook, RapiraCodec, TtlHook, TypedMap};

    fn open_users(
        db: &Arc<Db>,
    ) -> DbResult<Arc<TypedMap<UserId, User, RapiraCodec, TtlHook<User, NoHook>>>> {
        db.open_typed_map::<User, RapiraCodec, _>(Config::test(), NoHook, &[])
    }

    fn corrupt_stored_typ_hash(dir: &std::path::Path, name: &str) {
        let path = dir.join("db.info");
        let raw = std::fs::read_to_string(&path).unwrap();
        let mut info: DbInfo = serde_json::from_str(&raw).unwrap();
        if let Some(ci) = info.collections.get_mut(name) {
            ci.typ_hash = 0xdeadbeef;
        }
        std::fs::write(&path, serde_json::to_string(&info).unwrap()).unwrap();
    }

    #[test]
    fn tx_replay_finishes_a_transaction_whose_apply_never_ran() {
        let dir = tempfile::tempdir().unwrap();
        {
            let db = Arc::new(Db::open_test(dir.path()).unwrap());
            let (users, _counters) = open_pair(&db);
            crate::txn::SKIP_APPLY.store(true, Ordering::SeqCst);
            db.tx()
                .put(&users, &user_id(9), &User { active: true })
                .unwrap()
                .durable()
                .commit()
                .unwrap();
            crate::txn::SKIP_APPLY.store(false, Ordering::SeqCst);
            assert!(users.get(&user_id(9)).is_none(), "apply was suppressed");
        }

        let db = Arc::new(Db::open_test(dir.path()).unwrap());
        let (users, _counters) = open_pair(&db);
        assert_eq!(
            users.get(&user_id(9)).map(|u| u.active),
            Some(true),
            "opening the collection replays the committed record"
        );
        // Strategy A: once every participant named in the record has completed
        // open-time replay, the record is marked applied and can be checkpointed.
        db.checkpoint_now().unwrap();
        assert!(
            db.pending_transactions().is_empty(),
            "settled replayed records must be eligible for retirement"
        );
    }

    #[test]
    fn tx_replay_skips_an_operation_that_already_landed() {
        let dir = tempfile::tempdir().unwrap();
        {
            let db = Arc::new(Db::open_test(dir.path()).unwrap());
            let (users, _counters) = open_pair(&db);
            db.tx()
                .put(&users, &user_id(10), &User { active: true })
                .unwrap()
                .durable()
                .commit()
                .unwrap();
            users.put(&user_id(10), User { active: false }).unwrap();
        }

        let db = Arc::new(Db::open_test(dir.path()).unwrap());
        let (users, _counters) = open_pair(&db);
        assert_eq!(
            users.get(&user_id(10)).map(|u| u.active),
            Some(false),
            "replay must not resurrect the transaction's older value"
        );
    }

    /// Replay writes an entry carrying a GSN reserved *before* the crash, while
    /// open only seeds the collection counter to `recovered_max + 1`. If the
    /// counter is not pulled up past a replayed entry, the next transaction
    /// reserves a GSN that is already on disk, and the replay rule
    /// (`gsn > durable_recovered_gsn`) reads that transaction as "already
    /// applied" — dropping a committed transaction on the floor.
    #[test]
    fn tx_replay_keeps_the_gsn_counter_ahead_of_replayed_entries() {
        use crate::txn::participant::TxParticipant;

        let dir = tempfile::tempdir().unwrap();

        // Both keys must land on the same shard: the watermark is per-shard, so
        // a GSN collision is only observable within one.
        let probe = tempfile::tempdir().unwrap();
        let (k1, k2) = {
            let db = Arc::new(Db::open_test(probe.path()).unwrap());
            let users = open_users(&db).unwrap();
            let first = user_id(20);
            let sid = TxParticipant::shard_of(users.as_ref(), &first).unwrap();
            let second = (21u64..256)
                .map(user_id)
                .find(|k| TxParticipant::shard_of(users.as_ref(), k).unwrap() == sid)
                .expect("a second key on the same shard");
            (first, second)
        };

        // Session 1: the record is logged and fsynced, the apply never runs.
        {
            let db = Arc::new(Db::open_test(dir.path()).unwrap());
            let users = open_users(&db).unwrap();
            crate::txn::SKIP_APPLY.store(true, Ordering::SeqCst);
            db.tx()
                .put(&users, &k1, &User { active: true })
                .unwrap()
                .durable()
                .commit()
                .unwrap();
            crate::txn::SKIP_APPLY.store(false, Ordering::SeqCst);
            db.close().unwrap();
        }

        // Session 2: open replays k1 (writing the reserved GSN), then a fresh
        // transaction is logged — again without applying. Its reserved GSN must
        // be above the one replay just wrote.
        {
            let db = Arc::new(Db::open_test(dir.path()).unwrap());
            let users = open_users(&db).unwrap();
            assert_eq!(
                users.get(&k1).map(|u| u.active),
                Some(true),
                "open must replay the first transaction"
            );
            crate::txn::SKIP_APPLY.store(true, Ordering::SeqCst);
            db.tx()
                .put(&users, &k2, &User { active: true })
                .unwrap()
                .durable()
                .commit()
                .unwrap();
            crate::txn::SKIP_APPLY.store(false, Ordering::SeqCst);
            // Make the replayed entry durable, so session 3 recovers a watermark
            // that actually covers it.
            db.close().unwrap();
        }

        // Session 3: the second transaction must still be replayable.
        let db = Arc::new(Db::open_test(dir.path()).unwrap());
        let users = open_users(&db).unwrap();
        assert_eq!(users.get(&k1).map(|u| u.active), Some(true));
        assert_eq!(
            users.get(&k2).map(|u| u.active),
            Some(true),
            "a committed transaction was dropped: replay left the GSN counter \
             behind, so this transaction reused a GSN at or below the watermark"
        );
    }

    /// A collection whose schema advanced past a still-pending record.
    #[derive(Clone, Copy, Debug, PartialEq, Eq, GetType, Rapira)]
    struct Doc {
        n: u32,
    }
    impl CollectionMeta for Doc {
        type SelfId = [u8; 8];
        const NAME: &'static str = "docs";
        const VERSION: u16 = 1;
    }

    fn stored_schema_of(dir: &std::path::Path, name: &str) -> (u16, u64) {
        let raw = std::fs::read_to_string(dir.join("db.info")).unwrap();
        let info: DbInfo = serde_json::from_str(&raw).unwrap();
        let ci = info.collections.get(name).unwrap();
        (ci.version, ci.typ_hash)
    }

    /// Rewind `name` to a pre-migration schema, the state a session that
    /// migrated *after* replaying a record leaves behind for the next open.
    fn rewind_stored_schema(dir: &std::path::Path, name: &str) {
        let path = dir.join("db.info");
        let raw = std::fs::read_to_string(&path).unwrap();
        let mut info: DbInfo = serde_json::from_str(&raw).unwrap();
        if let Some(ci) = info.collections.get_mut(name) {
            ci.version = 0;
            ci.typ_hash = 0xdeadbeef;
        }
        std::fs::write(&path, serde_json::to_string(&info).unwrap()).unwrap();
    }

    /// The fingerprint gate exists to stop a record's **stale bytes** from being
    /// decoded by a migrated collection. An operation that no longer needs
    /// replay hands over no bytes at all, so a schema change must not turn it
    /// into a permanently unopenable database.
    ///
    /// The state is reachable without a crash: a record whose other participant
    /// was never opened stays pending across a session that replayed *and*
    /// migrated this collection. Every subsequent open then sees a stored
    /// fingerprint the record predates.
    #[test]
    fn tx_replay_ignores_a_schema_change_for_operations_that_no_longer_need_replay() {
        let dir = tempfile::tempdir().unwrap();
        let key = [3u8; 8];

        // Session 1: the record reaches the log, the apply never runs.
        {
            let db = Arc::new(Db::open_test(dir.path()).unwrap());
            let docs = db
                .open_typed_map::<Doc, RapiraCodec, _>(Config::test(), NoHook, &[])
                .unwrap();
            let (_users, counters) = open_pair(&db);
            crate::txn::SKIP_APPLY.store(true, Ordering::SeqCst);
            db.tx()
                .put(&docs, &key, &Doc { n: 7 })
                .unwrap()
                .put(&counters, &user_id(1), &Counter { n: 1 })
                .unwrap()
                .durable()
                .commit()
                .unwrap();
            crate::txn::SKIP_APPLY.store(false, Ordering::SeqCst);
            db.close().unwrap();
        }

        // Session 2: only `docs` is opened, so it replays its operation while
        // the record stays pending — `counters` never reports replay done.
        {
            let db = Arc::new(Db::open_test(dir.path()).unwrap());
            let docs = db
                .open_typed_map::<Doc, RapiraCodec, _>(Config::test(), NoHook, &[])
                .unwrap();
            assert_eq!(docs.get(&key).map(|d| d.n), Some(7), "replay applied it");
            assert_eq!(
                db.pending_transactions().len(),
                1,
                "the record stays pending until every participant has opened"
            );
            db.close().unwrap();
        }

        // The migration `docs` would run next stores a new schema; the record
        // predates it and is still in the log.
        rewind_stored_schema(dir.path(), Doc::NAME);

        // Session 3: `docs` migrates. Its pending operation is already on disk,
        // so replay decodes nothing — the schema change is irrelevant to it.
        let db = Arc::new(Db::open_test(dir.path()).unwrap());
        let docs = db
            .open_typed_map::<Doc, RapiraCodec, _>(
                Config::test(),
                NoHook,
                &[(
                    0u16,
                    (|_k: &[u8; 8], v: &Doc| crate::MigrateAction::Update(Doc { n: v.n })) as _,
                )],
            )
            .expect("an already-applied operation must not block the open");
        assert_eq!(docs.get(&key).map(|d| d.n), Some(7));
        let (version, _) = stored_schema_of(dir.path(), Doc::NAME);
        assert_eq!(version, 1, "the migration must have run");
    }

    #[test]
    fn tx_replay_fails_closed_on_a_schema_fingerprint_change() {
        let dir = tempfile::tempdir().unwrap();
        {
            let db = Arc::new(Db::open_test(dir.path()).unwrap());
            let (users, _counters) = open_pair(&db);
            crate::txn::SKIP_APPLY.store(true, Ordering::SeqCst);
            db.tx()
                .put(&users, &user_id(11), &User { active: true })
                .unwrap()
                .durable()
                .commit()
                .unwrap();
            crate::txn::SKIP_APPLY.store(false, Ordering::SeqCst);
        }
        corrupt_stored_typ_hash(dir.path(), User::NAME);

        let db = Arc::new(Db::open_test(dir.path()).unwrap());
        let err = match open_users(&db) {
            Err(e) => e,
            Ok(_) => panic!("open should fail on schema mismatch"),
        };
        assert!(
            matches!(err, crate::DbError::TxSchemaMismatch { .. }),
            "{err}"
        );
    }
}