kill_recovery.rs raw

//! Proves the queue's crash-tolerance contract against a real SIGKILL.
//!
//! `queue_survives_sigkill` re-invokes this test binary as a child process
//! running the `kill_worker` workload, kills it without warning at a random
//! point, then reopens the queue and checks the delivery guarantees:
//!
//! - every value the worker reported as durably enqueued (a `SYNC` line)
//!   and never acknowledged is still present (at-least-once delivery)
//! - no value whose acknowledgement was covered by a sync reappears
//! - the surviving values are intact and still in FIFO order

use std::{
    collections::HashSet,
    env,
    io::{Read as _, Write as _},
    process::{Child, Command, Stdio},
    time::Duration,
};

use persistent_queue::PersistentQueue;

const ENV_QUEUE: &str = "PERSISTENT_QUEUE_KILL_TEST";

/// Kills the child if the parent test panics before reaching its own kill.
struct ChildGuard(Child);

impl Drop for ChildGuard {
    fn drop(&mut self) {
        let _ = self.0.kill();
        let _ = self.0.wait();
    }
}

/// Not a real test: this is the child workload. It only does anything when
/// the parent test sets `ENV_QUEUE`; under a normal `cargo test` run it
/// returns immediately. As a child it churns the queue forever — putting,
/// leasing, releasing, acking, and syncing — while reporting durability
/// milestones on stdout, until the parent kills it.
#[test]
fn kill_worker() {
    let Some(queue_path) = env::var_os(ENV_QUEUE) else {
        return;
    };

    let queue: PersistentQueue<u64> = PersistentQueue::open(queue_path).unwrap();
    let stdout = std::io::stdout();
    let mut out = stdout.lock();

    let mut state = u64::from(std::process::id()) | 1;
    let mut random = move || {
        state = state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        state >> 33
    };

    let mut next_value = 0u64;

    loop {
        for _ in 0..random() % 8 {
            queue.put(&next_value).unwrap();
            next_value += 1;
        }

        for _ in 0..random() % 8 {
            let Some(lease) = queue.get().unwrap() else {
                break;
            };

            if random() % 4 == 0 {
                lease.release();
            } else {
                let value = lease.ack();
                // Flushed before any sync can make the ack durable, so the
                // parent always hears about acks that might have persisted
                writeln!(out, "ACK {value}").unwrap();
                out.flush().unwrap();
            }
        }

        if random() % 4 == 0 {
            queue.sync().unwrap();
            writeln!(out, "SYNC {next_value}").unwrap();
            out.flush().unwrap();
        }
    }
}

#[test]
fn queue_survives_sigkill() {
    for round in 0..5u64 {
        let dir = tempfile::tempdir().unwrap();
        let queue_path = dir.path().join("queue");

        let mut child = ChildGuard(
            Command::new(env::current_exe().unwrap())
                .args(["kill_worker", "--exact", "--nocapture"])
                .env(ENV_QUEUE, &queue_path)
                .stdout(Stdio::piped())
                .spawn()
                .unwrap(),
        );

        std::thread::sleep(Duration::from_millis(80 + round * 130));

        child.0.kill().unwrap();
        child.0.wait().unwrap();

        let mut output = String::new();
        child
            .0
            .stdout
            .take()
            .unwrap()
            .read_to_string(&mut output)
            .unwrap();

        // The kill may have torn the final line mid-write
        if !output.ends_with('\n') {
            output.truncate(output.rfind('\n').map_or(0, |pos| pos + 1));
        }

        // Reconstruct what the worker knew to be durable. ACK lines are
        // durable only once a later SYNC line covers them; SYNC reports
        // how many puts the sync made durable.
        let mut synced_puts = 0u64;
        let mut all_acks = HashSet::new();
        let mut synced_acks = HashSet::new();
        let mut unsynced_acks = Vec::new();

        for line in output.lines() {
            if let Some(value) = line.strip_prefix("ACK ") {
                let value: u64 = value.parse().unwrap();
                unsynced_acks.push(value);
                all_acks.insert(value);
            } else if let Some(count) = line.strip_prefix("SYNC ") {
                synced_puts = count.parse().unwrap();
                synced_acks.extend(unsynced_acks.drain(..));
            }
            // Anything else is libtest harness noise
        }

        let queue: PersistentQueue<u64> = PersistentQueue::open(&queue_path).unwrap();

        let mut remaining = Vec::new();
        while let Some(lease) = queue.get().unwrap() {
            remaining.push(lease.ack());
        }

        assert!(
            remaining.windows(2).all(|pair| pair[0] < pair[1]),
            "round {round}: recovered values are not in FIFO order: {remaining:?}"
        );

        for value in &remaining {
            assert!(
                !synced_acks.contains(value),
                "round {round}: value {value} was durably acknowledged but reappeared"
            );
        }

        let present: HashSet<u64> = remaining.iter().copied().collect();

        for value in 0..synced_puts {
            assert!(
                all_acks.contains(&value) || present.contains(&value),
                "round {round}: value {value} was durably enqueued, never \
                 acknowledged, but is missing after recovery"
            );
        }

        queue.sync().unwrap();
        assert!(queue.is_empty());
    }
}