clawhdf5/crates/clawhdf5-agent/tests/crash_recovery.rs

//! Crash-recovery matrix for `HDF5Memory`.
//!
//! A process crash leaves whatever reached the OS on disk. These tests build
//! the on-disk images such a crash can leave behind — after every operation,
//! inside the checkpoint window (new `.h5` in place, WAL not yet truncated),
//! and with the WAL torn at every possible length — then reopen each image
//! and check the recovered store against a model of what was acknowledged.
//!
//! Invariants:
//! * never a duplicated or invented record;
//! * an image taken between operations recovers *exactly* the acknowledged
//!   state;
//! * a torn WAL recovers the last checkpoint plus a prefix of the operations
//!   logged since.

use std::path::{Path, PathBuf};

use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
use tempfile::TempDir;

struct Rng(u64);

impl Rng {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }
    fn below(&mut self, n: usize) -> usize {
        (self.next() % n.max(1) as u64) as usize
    }
}

fn entry(chunk: &str, tags: &str) -> MemoryEntry {
    MemoryEntry {
        chunk: chunk.to_string(),
        embedding: vec![1.0, 0.0, 0.0, 0.0],
        source_channel: "test".into(),
        timestamp: 1.0,
        session_id: "s".into(),
        tags: tags.to_string(),
    }
}

fn wal_path(h5: &Path) -> PathBuf {
    h5.with_extension("h5.wal")
}

/// Copy the store (`.h5` + WAL) into a fresh directory, as a crash image.
fn image(h5: &Path, into: &TempDir, name: &str) -> PathBuf {
    let dest = into.path().join(format!("{name}.h5"));
    std::fs::copy(h5, &dest).unwrap();
    if wal_path(h5).exists() {
        std::fs::copy(wal_path(h5), wal_path(&dest)).unwrap();
    }
    dest
}

fn recovered(h5: &Path) -> Vec<String> {
    // Read-only: the image must not be modified, and no lock is needed.
    HDF5Memory::open_read_only(h5).unwrap().cache.chunks.clone()
}

/// Apply one random operation to the store and to the model.
fn step(mem: &mut HDF5Memory, model: &mut Vec<String>, rng: &mut Rng, n: usize) {
    match rng.below(6) {
        0 => mem.flush_wal().unwrap(),
        1 if !model.is_empty() => {
            // Update an existing record in place, addressed by its tag.
            let idx = rng.below(model.len());
            let chunk = format!("u{n}");
            assert_eq!(
                mem.save_or_update(entry(&chunk, &format!("tag{idx}")))
                    .unwrap(),
                idx
            );
            model[idx] = chunk;
        }
        _ => {
            let chunk = format!("c{n}");
            mem.save(entry(&chunk, &format!("tag{}", model.len())))
                .unwrap();
            model.push(chunk);
        }
    }
}

#[test]
fn image_after_every_operation_recovers_the_acknowledged_state() {
    for seed in 0..40u64 {
        let mut rng = Rng(seed);
        let dir = TempDir::new().unwrap();
        let images = TempDir::new().unwrap();
        let mut config = MemoryConfig::new(dir.path().join("store.h5"), "agent", 4);
        config.wal_enabled = true;
        config.wal_max_entries = 1 + rng.below(6); // force frequent checkpoints
        let h5 = config.path.clone();
        let mut mem = HDF5Memory::create(config).unwrap();
        let mut model = Vec::new();

        for n in 0..30 {
            step(&mut mem, &mut model, &mut rng, n);
            let img = image(&h5, &images, &format!("s{seed}-{n}"));
            assert_eq!(recovered(&img), model, "seed {seed}, after op {n}");
        }
    }
}

#[test]
fn crash_inside_the_checkpoint_window_never_duplicates() {
    for seed in 0..40u64 {
        let mut rng = Rng(seed ^ 0xABCD);
        let dir = TempDir::new().unwrap();
        let images = TempDir::new().unwrap();
        let mut config = MemoryConfig::new(dir.path().join("store.h5"), "agent", 4);
        config.wal_enabled = true;
        config.wal_max_entries = 1000; // checkpoints only when we ask
        let h5 = config.path.clone();
        let mut mem = HDF5Memory::create(config).unwrap();
        let mut model = Vec::new();

        for round in 0..4 {
            for n in 0..(1 + rng.below(6)) {
                step(&mut mem, &mut model, &mut rng, round * 100 + n);
            }
            // The WAL as it is just before the checkpoint...
            let stale_wal = images.path().join(format!("stale-{seed}-{round}.wal"));
            if wal_path(&h5).exists() {
                std::fs::copy(wal_path(&h5), &stale_wal).unwrap();
            }
            mem.flush_wal().unwrap();
            // ...put back next to the NEW .h5: the crash-in-the-window image.
            let img = image(&h5, &images, &format!("w{seed}-{round}"));
            if stale_wal.exists() {
                std::fs::copy(&stale_wal, wal_path(&img)).unwrap();
            }
            assert_eq!(recovered(&img), model, "seed {seed}, round {round}");
        }
    }
}

#[test]
fn torn_wal_recovers_checkpoint_plus_a_prefix() {
    let dir = TempDir::new().unwrap();
    let images = TempDir::new().unwrap();
    let mut config = MemoryConfig::new(dir.path().join("store.h5"), "agent", 4);
    config.wal_enabled = true;
    config.wal_max_entries = 1000;
    let h5 = config.path.clone();
    let mut mem = HDF5Memory::create(config).unwrap();

    for name in ["a", "b"] {
        mem.save(entry(name, name)).unwrap();
    }
    mem.flush_wal().unwrap();
    let checkpointed = vec!["a".to_string(), "b".to_string()];

    // States the store passes through as each later op is logged.
    let mut states = vec![checkpointed.clone()];
    let mut model = checkpointed.clone();
    mem.save(entry("c", "c")).unwrap();
    model.push("c".into());
    states.push(model.clone());
    mem.save_or_update(entry("a2", "a")).unwrap();
    model[0] = "a2".into();
    states.push(model.clone());
    mem.save(entry("d", "d")).unwrap();
    model.push("d".into());
    states.push(model.clone());

    let full_wal = std::fs::read(wal_path(&h5)).unwrap();
    let mut seen = std::collections::BTreeSet::new();
    for len in 0..=full_wal.len() {
        let img = image(&h5, &images, &format!("t{len}"));
        std::fs::write(wal_path(&img), &full_wal[..len]).unwrap();
        let got = recovered(&img);
        let which = states
            .iter()
            .position(|s| *s == got)
            .unwrap_or_else(|| panic!("WAL torn at {len} bytes recovered {got:?}"));
        seen.insert(which);
    }
    // Every intermediate state is reachable, and the full WAL gives the last.
    assert_eq!(seen.into_iter().collect::<Vec<_>>(), [0, 1, 2, 3]);
}