armdb 0.12.0

sharded bitcask key-value storage optimized for NVMe
Documentation
//! Transaction-log checkpoint: durable-flush participants, then retire applied
//! records at or below the cut.

use crate::armour::Db;
use crate::error::DbResult;
use crate::sync::lock;

/// Retire everything that is applied, at or below the cut, and durable.
///
/// The cut is one below the lowest transaction id currently in flight, so a
/// commit that started before the checkpoint cannot have its record retired
/// while its apply is still running.
pub(crate) fn run(db: &Db) -> DbResult<usize> {
    let cut = db.txlog().in_flight_floor();
    if cut == 0 {
        return Ok(0);
    }
    durable_flush_all(&db.collections)?;
    db.txlog().checkpoint(cut)
}

/// Durable-flush every published collection, holding the registry lock only
/// long enough to copy the handles out of it.
///
/// **The snapshot is the whole point.** A checkpoint takes the registry lock
/// and then, through `durable_flush`, a shard lock; `Tx::commit` takes them in
/// the opposite order — it holds the shard locks it locked for the commit and
/// then resolves collection ids under them, in `projection_commit::admit` ->
/// `TxContext::collection_id`. Flushing while still holding the registry closes
/// that cycle and deadlocks both threads (#149). Every other registry holder in
/// `Db` already snapshots for the same reason.
///
/// A collection opened after the snapshot cannot owe this cut a flush: `cut` is
/// read before the snapshot is taken, a record at or below it was logged before
/// that read, and a participant is published before a transaction naming it can
/// commit — so the snapshot already contains every participant of every record
/// this checkpoint may retire. The registry is append-only, so nothing can
/// leave it either.
pub(crate) fn durable_flush_all(
    collections: &crate::sync::Mutex<Vec<std::sync::Arc<dyn crate::armour::Collection>>>,
) -> DbResult<()> {
    let snapshot = lock(collections).clone();
    for c in snapshot.iter() {
        c.durable_flush()?;
    }
    Ok(())
}

#[cfg(all(test, feature = "rapira-codec"))]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
    use std::sync::mpsc;
    use std::thread::JoinHandle;
    use std::time::Duration;

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

    use crate::armour::{Collection, Db};
    use crate::sync::{Mutex, lock};
    use crate::{CollectionMeta, Config, DbResult, NoHook, RapiraCodec, TxConfig};

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

    /// Reports, from a second thread, whether the collection registry is still
    /// locked while this collection's `durable_flush` runs.
    ///
    /// The peer thread stands in for a committing transaction: `Tx::commit`
    /// holds shard locks and then takes the registry lock in
    /// `projection_commit::admit`. A checkpoint that flushes while holding that
    /// lock closes the cycle (#149).
    struct RegistryLockProbe {
        collections: Arc<Mutex<Vec<Arc<dyn Collection>>>>,
        flush_calls: Arc<AtomicUsize>,
        peer_acquired: Arc<AtomicBool>,
        peer: Mutex<Option<JoinHandle<()>>>,
    }

    impl Collection for RegistryLockProbe {
        fn name(&self) -> &str {
            "registry_lock_probe"
        }
        fn len(&self) -> usize {
            0
        }
        fn compact(&self) -> DbResult<usize> {
            Ok(0)
        }
        fn flush(&self) -> DbResult<()> {
            Ok(())
        }
        fn periodic_flush(&self) -> DbResult<()> {
            Ok(())
        }
        fn clean_shutdown(&self) -> DbResult<()> {
            Ok(())
        }

        fn durable_flush(&self) -> DbResult<()> {
            self.flush_calls.fetch_add(1, Ordering::SeqCst);
            let collections = Arc::clone(&self.collections);
            let (tx, rx) = mpsc::channel();
            // Joined by the caller *after* the checkpoint returns. Joining here
            // would hang on the very deadlock this probe exists to detect.
            let peer = std::thread::spawn(move || {
                let _guard = lock(&collections);
                let _ = tx.send(());
            });
            let acquired = rx.recv_timeout(Duration::from_secs(2)).is_ok();
            self.peer_acquired.store(acquired, Ordering::SeqCst);
            *lock(&self.peer) = Some(peer);
            Ok(())
        }
    }

    /// Runs `checkpoint` against a database holding one retirable record and a
    /// probe collection. Returns whether the registry was free while the probe
    /// was flushed; panics if the checkpoint never reached its flush loop, so a
    /// checkpoint that returns early cannot pass for a checkpoint that behaves.
    fn registry_free_during_flush(checkpoint: impl FnOnce(&Arc<Db>)) -> bool {
        let dir = tempfile::tempdir().unwrap();
        // The default threshold keeps `commit`'s own `maybe_checkpoint` inert,
        // so the record below is still retirable when the test checkpoints.
        let db = Arc::new(Db::open_with_tx_config(dir.path(), TxConfig::default()).unwrap());
        let rows = db
            .open_typed_map::<Row, RapiraCodec, _>(Config::test(), NoHook, &[])
            .unwrap();
        // A checkpoint with an empty log returns before the flush loop.
        db.tx()
            .put(&rows, &[0u8; 8], &Row { n: 1 })
            .unwrap()
            .commit()
            .unwrap();

        let flush_calls = Arc::new(AtomicUsize::new(0));
        let peer_acquired = Arc::new(AtomicBool::new(false));
        let probe = Arc::new(RegistryLockProbe {
            collections: Arc::clone(&db.collections),
            flush_calls: Arc::clone(&flush_calls),
            peer_acquired: Arc::clone(&peer_acquired),
            peer: Mutex::new(None),
        });
        lock(&db.collections).push(Arc::clone(&probe) as Arc<dyn Collection>);

        checkpoint(&db);

        if let Some(peer) = lock(&probe.peer).take() {
            peer.join().unwrap();
        }
        assert!(
            flush_calls.load(Ordering::SeqCst) > 0,
            "the checkpoint returned before flushing — this test proves nothing \
             about lock order until it reaches the flush loop"
        );
        peer_acquired.load(Ordering::SeqCst)
    }

    #[test]
    fn checkpoint_run_releases_the_registry_lock_before_flushing() {
        assert!(
            registry_free_during_flush(|db| {
                super::run(db).unwrap();
            }),
            "checkpoint::run held db.collections across durable_flush: a \
             concurrent Tx::commit holding shard locks deadlocks against it (#149)"
        );
    }

    /// End-to-end shape from #149: several threads committing while the
    /// checkpoints their own commits trigger run concurrently. Pre-fix this
    /// parks every thread in `futex_wait` within a few hundred commits.
    #[test]
    fn concurrent_commits_do_not_deadlock_against_checkpoints() {
        const THREADS: u64 = 8;
        const PER_THREAD: u64 = 200;

        let dir = tempfile::tempdir().unwrap();
        let db = Arc::new(
            Db::open_with_tx_config(
                dir.path(),
                TxConfig {
                    // Small enough that most commits reach the checkpoint.
                    checkpoint_bytes: 4096,
                    segment_bytes: 4096,
                    ..TxConfig::default()
                },
            )
            .unwrap(),
        );
        let rows = db
            .open_typed_map::<Row, RapiraCodec, _>(Config::test(), NoHook, &[])
            .unwrap();

        let done = Arc::new(AtomicUsize::new(0));
        for t in 0..THREADS {
            let db = Arc::clone(&db);
            let rows = Arc::clone(&rows);
            let done = Arc::clone(&done);
            std::thread::spawn(move || {
                for i in 0..PER_THREAD {
                    let key = (t * PER_THREAD + i).to_be_bytes();
                    db.tx()
                        .put(&rows, &key, &Row { n: i })
                        .unwrap()
                        .commit()
                        .unwrap();
                    done.fetch_add(1, Ordering::SeqCst);
                }
            });
        }

        let deadline = std::time::Instant::now() + Duration::from_secs(30);
        let total = (THREADS * PER_THREAD) as usize;
        while done.load(Ordering::SeqCst) < total {
            assert!(
                std::time::Instant::now() < deadline,
                "commits stalled at {}/{total} — checkpoint/commit deadlock (#149)",
                done.load(Ordering::SeqCst)
            );
            std::thread::sleep(Duration::from_millis(20));
        }
    }

    #[test]
    fn queue_owned_checkpoint_releases_the_registry_lock_before_flushing() {
        assert!(
            registry_free_during_flush(|db| {
                // A queue-owned context over this `Db`'s log and registry, with
                // a threshold low enough that `maybe_checkpoint` reaches its
                // flush loop instead of returning on the byte check.
                let base = &db.tx_runtime;
                let runtime = Arc::new(crate::txn::TxRuntime {
                    config: TxConfig {
                        checkpoint_bytes: 0,
                        ..TxConfig::default()
                    },
                    log: Arc::clone(&base.log),
                    collections: Arc::clone(&base.collections),
                    ranks: Arc::clone(&base.ranks),
                    post_log_queue_producer_permits: Mutex::new(Vec::new()),
                    _root_lease: Arc::clone(&base._root_lease),
                });
                crate::txn::TxContext::Runtime(runtime)
                    .maybe_checkpoint()
                    .unwrap();
            }),
            "TxContext::Runtime::maybe_checkpoint held the registry across \
             durable_flush — the same cycle as #149, on the queue-owned path"
        );
    }
}