armdb 0.12.0

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::sync::mpsc;
use std::time::Duration;

use armdb::armour::Db;
use armdb::{QueueConfig, queue};
use armour_core::GetType;
use rapira::Rapira;

#[derive(
    Clone,
    Copy,
    GetType,
    zerocopy::FromBytes,
    zerocopy::IntoBytes,
    zerocopy::Immutable,
    zerocopy::KnownLayout,
)]
#[repr(transparent)]
pub struct Job(u64);

queue!(
    RetentionJobs,
    payload = Job,
    name = "queue_retention_fixed_jobs"
);

#[derive(Clone, Debug, PartialEq, Eq, GetType, Rapira)]
pub struct VarJob(String);

queue!(
    RetentionVarJobs,
    payload = VarJob,
    name = "queue_retention_var_jobs",
    var
);

fn open_retention_queue(
    dead_retention: Option<Duration>,
) -> (
    tempfile::TempDir,
    Db,
    std::sync::Arc<armdb::Queue<RetentionJobs>>,
) {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let q = db
        .open_queue::<RetentionJobs>(
            QueueConfig::balanced()
                .max_attempts(1)
                .dead_retention(dead_retention)
                .build()
                .unwrap(),
        )
        .unwrap();
    (dir, db, q)
}

fn open_var_retention_queue() -> (
    tempfile::TempDir,
    Db,
    std::sync::Arc<armdb::Queue<RetentionVarJobs>>,
) {
    let dir = tempfile::tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let q = db
        .open_queue::<RetentionVarJobs>(
            QueueConfig::balanced()
                .max_attempts(1)
                .dead_retention(Some(Duration::from_millis(1)))
                .build()
                .unwrap(),
        )
        .unwrap();
    (dir, db, q)
}

fn push_then_lease(q: &armdb::Queue<RetentionVarJobs>) -> armdb::LeaseToken {
    q.push(VarJob("payload".into()), None).unwrap();
    q.lease(1, None).unwrap().pop().unwrap().token()
}

#[test]
fn final_nack_sets_an_absolute_deadline_and_ttl_removes_the_dead_row() {
    let (_dir, db, q) = open_retention_queue(Some(Duration::from_millis(1)));
    q.push(Job(1), None).unwrap();
    let leased = q.lease(1, None).unwrap().pop().unwrap();
    q.nack(leased.token()).unwrap();
    assert_eq!(q.len(), 0);
    assert_eq!(q.dead_len(), 1);
    std::thread::sleep(Duration::from_millis(3));
    db.sweep_expired_now().unwrap();
    assert_eq!(q.dead_len(), 0);
}

#[test]
fn failed_var_payload_delete_is_retried_without_restart() {
    let (_dir, db, q) = open_var_retention_queue();
    let token = push_then_lease(&q);
    q.debug_fail_next_payload_delete();
    q.nack(token).unwrap();
    assert_eq!(q.debug_payload_len(), 1);
    std::thread::sleep(Duration::from_millis(3));
    db.sweep_expired_now().unwrap();
    assert_eq!(q.debug_payload_len(), 0);
}

#[test]
fn retry_failure_preserves_all_pending_payload_cleanup_ids() {
    let (_dir, db, q) = open_var_retention_queue();
    for _ in 0..3 {
        q.push(VarJob("payload".into()), None).unwrap();
    }
    let tokens: Vec<_> = q
        .lease(3, None)
        .unwrap()
        .into_iter()
        .map(|leased| leased.token())
        .collect();
    for token in tokens {
        q.nack(token).unwrap();
    }
    q.debug_fail_payload_deletes(4);
    std::thread::sleep(Duration::from_millis(3));
    assert!(db.sweep_expired_now().is_err());
    assert_eq!(q.debug_payload_len(), 3);

    q.debug_fail_payload_deletes(1);
    assert!(db.sweep_expired_now().is_err());
    db.sweep_expired_now().unwrap();
    assert_eq!(q.debug_payload_len(), 0);
}

#[test]
fn transactional_var_push_drains_pending_payload_cleanup() {
    let (_dir, db, q) = open_var_retention_queue();
    let token = push_then_lease(&q);
    q.nack(token).unwrap();
    q.debug_fail_payload_deletes(2);
    std::thread::sleep(Duration::from_millis(3));
    assert!(db.sweep_expired_now().is_err());
    assert_eq!(q.debug_payload_len(), 1);
    q.debug_fail_payload_deletes(0);

    db.tx()
        .push(&q, VarJob("next".into()), None)
        .unwrap()
        .commit()
        .unwrap();

    assert_eq!(q.debug_payload_len(), 1, "the old dead payload was retried");
    assert_eq!(q.len(), 1);
}

#[test]
fn a_persistently_failing_payload_delete_never_rejects_queue_work() {
    let (_dir, db, q) = open_var_retention_queue();
    let token = push_then_lease(&q);
    q.nack(token).unwrap();
    q.debug_fail_payload_deletes(usize::MAX);
    std::thread::sleep(Duration::from_millis(3));
    assert!(
        db.sweep_expired_now().is_err(),
        "the sweeper still surfaces the failure"
    );
    let stuck = q.debug_payload_len();
    assert_eq!(stuck, 1, "the dead payload is retained for a later retry");

    // Retention frees bytes belonging to rows that are already gone. Its
    // failure must never reach a user operation on the queue.
    q.push(VarJob("after".into()), None).unwrap();
    db.tx()
        .push(&q, VarJob("in a transaction".into()), None)
        .unwrap()
        .commit()
        .unwrap();
    let leased = q.lease(2, None).unwrap();
    assert_eq!(leased.len(), 2);
    for entry in leased {
        q.nack(entry.token()).unwrap();
    }

    q.debug_fail_payload_deletes(0);
    std::thread::sleep(Duration::from_millis(3));
    db.sweep_expired_now().unwrap();
    assert_eq!(q.debug_payload_len(), 0);
    assert_eq!(q.dead_len(), 0);
}

#[test]
fn no_retention_keeps_dead_rows_across_ttl_sweeps() {
    let (_dir, db, q) = open_retention_queue(None);
    q.push(Job(1), None).unwrap();
    let leased = q.lease(1, None).unwrap().pop().unwrap();
    q.nack(leased.token()).unwrap();
    std::thread::sleep(Duration::from_millis(3));
    db.sweep_expired_now().unwrap();
    assert_eq!(q.dead_len(), 1);
}

#[test]
fn full_orphan_pass_defers_while_payload_first_push_holds_a_producer_permit() {
    let (_dir, _db, q) = open_var_retention_queue();
    let (written, release) = q.debug_pause_after_payload();
    let pushing = std::sync::Arc::clone(&q);
    let push = std::thread::spawn(move || pushing.push(VarJob("kept".into()), None));
    written.wait();

    let retaining = std::sync::Arc::clone(&q);
    let (done_tx, done_rx) = mpsc::channel();
    let retention = std::thread::spawn(move || {
        let result = retaining.debug_force_full_orphan_pass();
        done_tx.send(result).unwrap();
    });
    let result = done_rx.recv_timeout(Duration::from_secs(1));
    assert_eq!(
        q.debug_payload_len(),
        1,
        "deferred maintenance must not delete an in-flight payload"
    );
    release.wait();
    push.join().unwrap().unwrap();
    retention.join().unwrap();
    result
        .expect("full orphan maintenance must defer instead of waiting for a producer")
        .unwrap();
    assert_eq!(q.lease(1, None).unwrap().pop().unwrap().payload().0, "kept");
}