armdb 0.11.2

use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::process::Command;
use std::time::Duration;

use armdb::armour::Db;
use armdb::{DbError, LeaseToken, Queue, QueueConfig, queue};
use armour_core::GetType;
use proptest::prelude::*;

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

queue!(
    ModelJobs,
    payload = ModelJob,
    name = "queue_state_model_jobs"
);

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum State {
    Ready,
    Leased(u32),
    Dead,
    Acked,
}

#[derive(Clone, Debug)]
enum Operation {
    Push(u16),
    Lease,
    Ack { stale: bool, pick: u8 },
    Nack { stale: bool, pick: u8 },
    CleanReopen,
}

fn operations() -> impl Strategy<Value = Vec<Operation>> {
    prop::collection::vec(
        prop_oneof![
            35 => any::<u16>().prop_map(Operation::Push),
            25 => Just(Operation::Lease),
            15 => (any::<bool>(), any::<u8>())
                .prop_map(|(stale, pick)| Operation::Ack { stale, pick }),
            15 => (any::<bool>(), any::<u8>())
                .prop_map(|(stale, pick)| Operation::Nack { stale, pick }),
            10 => Just(Operation::CleanReopen),
        ],
        1..80,
    )
}

fn crash_operations() -> impl Strategy<Value = Vec<Operation>> {
    prop::collection::vec(0u8..3, 1..16).prop_map(|outcomes| {
        outcomes
            .into_iter()
            .enumerate()
            .flat_map(|(id, outcome)| {
                let mut operations = vec![Operation::Push(id as u16), Operation::Lease];
                match outcome {
                    0 => operations.push(Operation::Ack {
                        stale: false,
                        pick: 0,
                    }),
                    1 => operations.extend([
                        Operation::Nack {
                            stale: false,
                            pick: 0,
                        },
                        Operation::Lease,
                        Operation::Ack {
                            stale: false,
                            pick: 0,
                        },
                    ]),
                    _ => operations.extend([
                        Operation::Nack {
                            stale: false,
                            pick: 0,
                        },
                        Operation::Lease,
                        Operation::Nack {
                            stale: false,
                            pick: 0,
                        },
                    ]),
                }
                // Each outcome leaves a stale token and no active lease. This
                // is the process boundary: tokens never cross it, but stale
                // token fencing is exercised before and after the crash.
                operations.push(Operation::Ack {
                    stale: true,
                    pick: 0,
                });
                operations
            })
            .collect()
    })
}

struct QueueSystem {
    dir: tempfile::TempDir,
    db: Option<Db>,
    queue: Option<std::sync::Arc<Queue<ModelJobs>>>,
}

impl QueueSystem {
    fn new(config: QueueConfig) -> Self {
        let dir = tempfile::tempdir().expect("create database directory");
        let db = Db::open_test(dir.path()).expect("open database");
        let queue = db.open_queue::<ModelJobs>(config).expect("open queue");
        Self {
            dir,
            db: Some(db),
            queue: Some(queue),
        }
    }

    fn queue(&self) -> &std::sync::Arc<Queue<ModelJobs>> {
        self.queue.as_ref().expect("queue is open")
    }

    /// Model a graceful close and reopen within this process.
    ///
    /// This is intentionally not a crash operation: dropping `Db` can run its
    /// normal shutdown path. The hard-exit trace below owns crash recovery.
    fn clean_reopen(&mut self) {
        drop(self.queue.take());
        drop(self.db.take());
        let db = Db::open_test(self.dir.path()).expect("reopen database");
        let queue = db
            .open_queue::<ModelJobs>(model_config())
            .expect("reopen queue");
        self.db = Some(db);
        self.queue = Some(queue);
    }
}

fn model_config() -> QueueConfig {
    QueueConfig::balanced()
        .max_attempts(2)
        .backoff_secs(&[0])
        .lease_timeout(Duration::from_secs(60))
        .build()
        .expect("valid deterministic queue config")
}

fn pick_token(tokens: &BTreeMap<u64, LeaseToken>, pick: u8) -> Option<(u64, LeaseToken)> {
    tokens
        .iter()
        .nth(usize::from(pick) % tokens.len().max(1))
        .map(|(&id, &token)| (id, token))
}

fn assert_counts(
    model: &BTreeMap<u64, State>,
    seen_ids: &BTreeSet<u64>,
    queue: &Queue<ModelJobs>,
    op: usize,
) {
    let live = model
        .values()
        .filter(|state| matches!(state, State::Ready | State::Leased(_)))
        .count() as u64;
    let dead = model
        .values()
        .filter(|state| matches!(state, State::Dead))
        .count() as u64;
    assert_eq!(
        queue.len(),
        live,
        "live count mismatch after operation {op}"
    );
    assert_eq!(
        queue.dead_len(),
        dead,
        "dead count mismatch after operation {op}"
    );
    // Conservation: every id the queue ever handed out is in exactly one state,
    // and the live/dead halves are the ones the engine reports above. Without
    // this, `queue.len()` and `queue.dead_len()` agreeing with the model says
    // nothing about a job that silently left both.
    let acked = model
        .values()
        .filter(|state| matches!(state, State::Acked))
        .count() as u64;
    assert_eq!(
        live + dead + acked,
        seen_ids.len() as u64,
        "a job left the model without being acked, dead-lettered or counted after operation {op}"
    );
    // `seen_ids` collects ids from `push` *and* from every `lease` return, so
    // the two sides of this comparison come from different sources: the ids the
    // test handed in, and the ids the engine handed back.
    assert_eq!(
        model.keys().copied().collect::<BTreeSet<_>>(),
        seen_ids.clone(),
        "the engine returned a queue id the model never pushed after operation {op}"
    );
}

fn run_model(operations: Vec<Operation>) {
    // Two delivery attempts allow both a retry with a new epoch and one final
    // nack. A zero retry delay keeps this state machine independent of clocks.
    let mut system = QueueSystem::new(model_config());
    let mut model = BTreeMap::<u64, State>::new();
    let mut seen_ids = BTreeSet::<u64>::new();
    let mut attempts = HashMap::<u64, u16>::new();
    let mut active = BTreeMap::<u64, LeaseToken>::new();
    let mut stale = Vec::<LeaseToken>::new();
    let mut final_nacks = HashMap::<u64, u8>::new();

    for (op_index, operation) in operations.into_iter().enumerate() {
        match operation {
            Operation::Push(value) => {
                let id = system
                    .queue()
                    .push(ModelJob(u64::from(value)), None)
                    .expect("push model job");
                assert_eq!(
                    model.insert(id, State::Ready),
                    None,
                    "id reuse at operation {op_index}"
                );
                assert!(seen_ids.insert(id), "id {id} was observed twice on push");
                attempts.insert(id, 0);
            }
            Operation::Lease => {
                let expected_count = model
                    .iter()
                    .filter(|(_, state)| matches!(state, State::Ready))
                    .count()
                    .min(2);
                let leased = system
                    .queue()
                    .lease(2, Some(Duration::from_secs(60)))
                    .expect("lease model jobs");
                let returned: Vec<_> = leased.iter().map(|entry| entry.id()).collect();
                assert_eq!(
                    returned.len(),
                    expected_count,
                    "lease skipped a due ready row or returned more than its limit at operation {op_index}"
                );
                assert!(
                    returned
                        .iter()
                        .all(|id| matches!(model.get(id), Some(State::Ready))),
                    "lease returned an id that was not ready at operation {op_index}"
                );
                for entry in leased {
                    let id = entry.id();
                    let token = entry.token();
                    seen_ids.insert(id);
                    assert!(
                        matches!(model.get(&id), Some(State::Ready)),
                        "lease returned id {id} that was not ready at operation {op_index}"
                    );
                    assert!(
                        active.insert(id, token).is_none(),
                        "lease duplicated active token for id {id} at operation {op_index}"
                    );
                    model.insert(id, State::Leased(token.epoch()));
                }
            }
            Operation::Ack {
                stale: use_stale,
                pick,
            } if use_stale && !stale.is_empty() => {
                let token = stale[usize::from(pick) % stale.len()];
                assert!(
                    matches!(system.queue().ack(token), Err(DbError::LeaseExpired)),
                    "stale ack succeeded at operation {op_index}"
                );
            }
            Operation::Ack { stale: _, pick } => {
                if let Some((id, token)) = pick_token(&active, pick) {
                    system.queue().ack(token).expect("ack active lease");
                    active.remove(&id);
                    stale.push(token);
                    assert_eq!(
                        model.insert(id, State::Acked),
                        Some(State::Leased(token.epoch()))
                    );
                }
            }
            Operation::Nack { stale: true, pick } if !stale.is_empty() => {
                let token = stale[usize::from(pick) % stale.len()];
                assert!(
                    matches!(system.queue().nack(token), Err(DbError::LeaseExpired)),
                    "stale nack succeeded at operation {op_index}"
                );
            }
            Operation::Nack { stale: _, pick } => {
                if let Some((id, token)) = pick_token(&active, pick) {
                    system.queue().nack(token).expect("nack active lease");
                    active.remove(&id);
                    stale.push(token);
                    let next_attempt = attempts.get_mut(&id).expect("attempt state for leased id");
                    *next_attempt += 1;
                    if *next_attempt == 2 {
                        let final_count = final_nacks.entry(id).or_default();
                        *final_count += 1;
                        assert_eq!(*final_count, 1, "id {id} received more than one final nack");
                        assert_eq!(
                            model.insert(id, State::Dead),
                            Some(State::Leased(token.epoch()))
                        );
                    } else {
                        assert_eq!(
                            model.insert(id, State::Ready),
                            Some(State::Leased(token.epoch()))
                        );
                    }
                }
            }
            Operation::CleanReopen => {
                system.clean_reopen();
                // Lease tokens are process-local by design: a reopen mints a new
                // queue instance, so every outstanding token is fenced and now
                // behaves exactly like a stale one. The rows they named stay
                // live and leased on disk — unleasable until their deadline,
                // which `model_config` puts well past this test's runtime — so
                // their model state is deliberately left alone.
                stale.extend(active.values().copied());
                active.clear();
            }
        }

        assert_counts(&model, &seen_ids, system.queue(), op_index);
        for (&id, &token) in &active {
            assert_eq!(
                model.get(&id),
                Some(&State::Leased(token.epoch())),
                "active lease token diverged from model at operation {op_index}"
            );
        }
    }

    assert!(
        final_nacks.values().all(|&count| count == 1),
        "every dead-letter transition must come from exactly one final nack"
    );
}

fn crash_logical_id(next: &mut u64) -> u64 {
    let id = *next;
    *next += 1;
    id
}

fn encode_operations(operations: &[Operation]) -> String {
    operations
        .iter()
        .map(|operation| match operation {
            Operation::Push(value) => format!("p:{value}"),
            Operation::Lease => "l".into(),
            Operation::Ack { stale, pick } => format!("a:{}:{pick}", u8::from(*stale)),
            Operation::Nack { stale, pick } => format!("n:{}:{pick}", u8::from(*stale)),
            Operation::CleanReopen => "r".into(),
        })
        .collect::<Vec<_>>()
        .join(",")
}

fn decode_operations(trace: &str) -> Vec<Operation> {
    trace
        .split(',')
        .map(|part| {
            let mut fields = part.split(':');
            match fields.next().expect("operation tag") {
                "p" => Operation::Push(fields.next().expect("push value").parse().expect("u16")),
                "l" => Operation::Lease,
                "r" => Operation::CleanReopen,
                "a" => Operation::Ack {
                    stale: fields.next().expect("stale flag") == "1",
                    pick: fields.next().expect("pick").parse().expect("u8"),
                },
                "n" => Operation::Nack {
                    stale: fields.next().expect("stale flag") == "1",
                    pick: fields.next().expect("pick").parse().expect("u8"),
                },
                tag => panic!("unknown crash trace operation {tag}"),
            }
        })
        .collect()
}

fn run_crash_child(operations: Vec<Operation>) -> ! {
    let dir = std::env::var("ARMDB_QUEUE_MODEL_CRASH_DIR").expect("crash directory");
    let db = Db::open_test(&dir).expect("open crash model database");
    let queue = db
        .open_queue::<ModelJobs>(model_config())
        .expect("open crash model queue");
    let mut active = BTreeMap::<u64, LeaseToken>::new();
    let mut stale = Vec::<LeaseToken>::new();
    let mut next_id = 0;
    for operation in operations {
        match operation {
            Operation::Push(_) => {
                let id = crash_logical_id(&mut next_id);
                queue.push(ModelJob(id), None).expect("child push");
            }
            Operation::Lease => {
                for entry in queue
                    .lease(2, Some(Duration::from_secs(60)))
                    .expect("child lease")
                {
                    assert!(active.insert(entry.payload().0, entry.token()).is_none());
                }
            }
            Operation::Ack { stale: true, pick } if !stale.is_empty() => {
                assert!(matches!(
                    queue.ack(stale[usize::from(pick) % stale.len()]),
                    Err(DbError::LeaseExpired)
                ));
            }
            Operation::Ack { pick, .. } => {
                if let Some((id, token)) = pick_token(&active, pick) {
                    queue.ack(token).expect("child ack");
                    active.remove(&id);
                    stale.push(token);
                }
            }
            Operation::Nack { stale: true, pick } if !stale.is_empty() => {
                assert!(matches!(
                    queue.nack(stale[usize::from(pick) % stale.len()]),
                    Err(DbError::LeaseExpired)
                ));
            }
            Operation::Nack { pick, .. } => {
                if let Some((id, token)) = pick_token(&active, pick) {
                    queue.nack(token).expect("child nack");
                    active.remove(&id);
                    stale.push(token);
                }
            }
            Operation::CleanReopen => unreachable!("crash trace has no clean reopen"),
        }
    }
    unsafe { libc::_exit(9) }
}

#[test]
fn queue_model_crash_child_executes_randomized_prefix_then_hard_exits() {
    let Ok(trace_path) = std::env::var("ARMDB_QUEUE_MODEL_CRASH_TRACE") else {
        return;
    };
    let trace = std::fs::read_to_string(trace_path).expect("read child crash trace");
    run_crash_child(decode_operations(&trace));
}

fn simulate_crash_prefix(
    operations: &[Operation],
) -> (BTreeMap<u64, State>, HashMap<u64, u16>, u64) {
    let mut model = BTreeMap::new();
    let mut attempts = HashMap::new();
    let mut active = BTreeMap::<u64, u32>::new();
    let mut stale_exists = false;
    let mut next_id = 0;
    for operation in operations {
        match *operation {
            Operation::Push(_) => {
                let id = crash_logical_id(&mut next_id);
                assert_eq!(model.insert(id, State::Ready), None);
                attempts.insert(id, 0);
            }
            Operation::Lease => {
                let (&id, _) = model
                    .iter()
                    .find(|(_, state)| matches!(state, State::Ready))
                    .expect("one ready row per crash trace lease");
                active.insert(id, 1);
                model.insert(id, State::Leased(1));
            }
            Operation::Ack { stale: true, .. } => assert!(stale_exists),
            Operation::Ack { .. } => {
                let (&id, _) = active.iter().next().expect("active row for ack");
                active.remove(&id);
                stale_exists = true;
                model.insert(id, State::Acked);
            }
            Operation::Nack { stale: true, .. } => assert!(stale_exists),
            Operation::Nack { .. } => {
                let (&id, _) = active.iter().next().expect("active row for nack");
                active.remove(&id);
                stale_exists = true;
                let attempt = attempts.get_mut(&id).expect("attempt state");
                *attempt += 1;
                model.insert(
                    id,
                    if *attempt == 2 {
                        State::Dead
                    } else {
                        State::Ready
                    },
                );
            }
            Operation::CleanReopen => unreachable!("crash trace has no clean reopen"),
        }
    }
    assert!(
        active.is_empty(),
        "crash trace must end at token-free boundary"
    );
    (model, attempts, next_id)
}

fn run_crash_suffix(
    operations: &[Operation],
    mut model: BTreeMap<u64, State>,
    mut attempts: HashMap<u64, u16>,
    mut next_id: u64,
    queue: &Queue<ModelJobs>,
) -> (BTreeMap<u64, State>, HashMap<u64, u16>, u64) {
    let mut seen_ids: BTreeSet<u64> = model.keys().copied().collect();
    let mut active = BTreeMap::<u64, LeaseToken>::new();
    let mut stale = Vec::<LeaseToken>::new();
    let mut final_nacks = HashMap::<u64, u8>::new();
    for (op_index, operation) in operations.iter().enumerate() {
        match *operation {
            Operation::Push(_) => {
                let id = crash_logical_id(&mut next_id);
                queue.push(ModelJob(id), None).expect("push after crash");
                assert_eq!(model.insert(id, State::Ready), None);
                seen_ids.insert(id);
                attempts.insert(id, 0);
            }
            Operation::Lease => {
                let expected = model
                    .values()
                    .filter(|state| matches!(state, State::Ready))
                    .count()
                    .min(2);
                let leased = queue
                    .lease(2, Some(Duration::from_secs(60)))
                    .expect("lease after crash");
                assert_eq!(
                    leased.len(),
                    expected,
                    "recovered lease count at suffix operation {op_index}"
                );
                for entry in leased {
                    let id = entry.payload().0;
                    seen_ids.insert(id);
                    assert_eq!(
                        model.get(&id),
                        Some(&State::Ready),
                        "unknown recovered logical id {id}"
                    );
                    let token = entry.token();
                    active.insert(id, token);
                    model.insert(id, State::Leased(token.epoch()));
                }
            }
            Operation::Ack { stale: true, pick } if !stale.is_empty() => {
                assert!(matches!(
                    queue.ack(stale[usize::from(pick) % stale.len()]),
                    Err(DbError::LeaseExpired)
                ));
            }
            Operation::Ack { pick, .. } => {
                if let Some((id, token)) = pick_token(&active, pick) {
                    queue.ack(token).expect("ack after crash");
                    active.remove(&id);
                    stale.push(token);
                    model.insert(id, State::Acked);
                }
            }
            Operation::Nack { stale: true, pick } if !stale.is_empty() => {
                assert!(matches!(
                    queue.nack(stale[usize::from(pick) % stale.len()]),
                    Err(DbError::LeaseExpired)
                ));
            }
            Operation::Nack { pick, .. } => {
                if let Some((id, token)) = pick_token(&active, pick) {
                    queue.nack(token).expect("nack after crash");
                    active.remove(&id);
                    stale.push(token);
                    let attempt = attempts.get_mut(&id).expect("attempt state");
                    *attempt += 1;
                    if *attempt == 2 {
                        *final_nacks.entry(id).or_default() += 1;
                        model.insert(id, State::Dead);
                    } else {
                        model.insert(id, State::Ready);
                    }
                }
            }
            Operation::CleanReopen => {}
        }
        assert_counts(&model, &seen_ids, queue, op_index);
    }
    assert!(final_nacks.values().all(|&count| count == 1));
    (model, attempts, next_id)
}

fn run_crash_model(prefix: Vec<Operation>, suffix: Vec<Operation>) {
    let dir = tempfile::tempdir().expect("create randomized crash model directory");
    let trace_path = dir.path().join("queue-model-crash-trace");
    std::fs::write(&trace_path, encode_operations(&prefix)).expect("write child crash trace");
    let (model, attempts, next_id) = simulate_crash_prefix(&prefix);
    let status = Command::new(std::env::current_exe().expect("current test executable"))
        .arg("--exact")
        .arg("queue_model_crash_child_executes_randomized_prefix_then_hard_exits")
        .arg("--nocapture")
        .env("ARMDB_QUEUE_MODEL_CRASH_DIR", dir.path())
        .env("ARMDB_QUEUE_MODEL_CRASH_TRACE", &trace_path)
        .status()
        .expect("run randomized crash model child");
    assert_eq!(status.code(), Some(9), "crash operation must hard-exit");

    let db = Db::open_test(dir.path()).expect("reopen randomized crash model database");
    let queue = db
        .open_queue::<ModelJobs>(model_config())
        .expect("reopen randomized crash model queue");
    assert_counts(
        &model,
        &model.keys().copied().collect(),
        &queue,
        prefix.len(),
    );
    let _ = run_crash_suffix(&suffix, model, attempts, next_id, &queue);
}

#[test]
fn retry_fences_the_old_token_and_one_final_nack_dead_letters_the_job() {
    run_model(vec![
        Operation::Push(1),
        Operation::Lease,
        Operation::Nack {
            stale: false,
            pick: 0,
        },
        Operation::Ack {
            stale: true,
            pick: 0,
        },
        Operation::Nack {
            stale: true,
            pick: 0,
        },
        Operation::Lease,
        Operation::Nack {
            stale: false,
            pick: 0,
        },
        Operation::Nack {
            stale: true,
            pick: 1,
        },
    ]);
}

proptest! {
    // No child processes here, so this one can explore far more than the crash
    // model below; it is still the long pole of the default suite, hence 48
    // rather than proptest's default 256.
    #![proptest_config(ProptestConfig::with_cases(48))]

    #[test]
    fn fixed_queue_state_model(ops in operations()) {
        run_model(ops);
    }
}

proptest! {
    // One hard-exiting child process per case.
    #![proptest_config(ProptestConfig::with_cases(16))]

    #[test]
    fn randomized_queue_state_model_crash_replay_interleaves_with_queue_operations(
        prefix in crash_operations(),
        suffix in crash_operations(),
    ) {
        run_crash_model(prefix, suffix);
    }
}