improvise/crates/zset/tests/rollback_rebuild.rs

//! A `Rebuild` stepped across a rollback, and closures that read streams during an eval.
//! Evidence: `context/test-evidence/zset.md`.

use std::panic::{AssertUnwindSafe, catch_unwind};
use zset::{Circuit, Progress, Rebuild, ZSet};

type Z<K> = ZSet<K>;

fn quiet() {
    static O: std::sync::Once = std::sync::Once::new();
    O.call_once(|| {
        let prev = std::panic::take_hook();
        std::panic::set_hook(Box::new(move |info| {
            let s = format!("{info}");
            if !(s.contains("tripped") || s.contains("overflow")) {
                prev(info);
            }
        }));
    });
}

fn z(keys: &[(u8, i64)]) -> Z<u8> {
    quiet();
    keys.iter().copied().collect()
}

#[test]
fn r1_rebuild_begun_before_txn_continued_inside_then_rolled_back_reports_interrupted() {
    let mut c = Circuit::new();
    let (input, s) = c.input::<u8>();
    let i = c.integrate(&s);
    let snap = z(&[(1, 1), (2, 1), (3, 1), (4, 1)]);
    let mut job = Rebuild::new(input, snap.clone(), 1);
    assert_eq!(job.step(&mut c, || true), Progress::Pending);
    c.begin();
    assert_eq!(job.step(&mut c, || true), Progress::Pending);
    c.rollback();
    let interrupted = job.interrupted(&c);
    // Docs (zset.md Transactions 6, header): the job reports itself dead after a rollback.
    let mut finished = None;
    if !interrupted {
        while job.step(&mut c, || false) == Progress::Pending {}
        finished = Some(c.current(&i));
    }
    assert!(
        interrupted,
        "[r1] job not interrupted after rollback; finishing it gives {:?}, want {:?}",
        finished, snap
    );
}

#[test]
fn r2_rebuild_finished_inside_a_rolled_back_txn_is_reported_done_with_nothing_loaded() {
    let mut c = Circuit::new();
    let (input, s) = c.input::<u8>();
    let i = c.integrate(&s);
    let snap = z(&[(1, 1), (2, 1), (3, 1)]);
    let mut job = Rebuild::new(input, snap.clone(), 1);
    c.begin();
    assert_eq!(job.step(&mut c, || false), Progress::Done);
    c.rollback();
    let interrupted = job.interrupted(&c);
    assert!(
        interrupted,
        "[r2] after rollback: not interrupted, integral {:?}",
        c.current(&i)
    );
}

#[test]
fn r3_an_interrupted_rebuild_resurrects_after_rollback_and_a_later_step() {
    let mut c = Circuit::new();
    let (input, s) = c.input::<u8>();
    let i = c.integrate(&s);
    c.step(); // ticks 1
    let snap = z(&[(1, 1), (2, 1), (3, 1), (4, 1)]);
    let mut job = Rebuild::new(input, snap.clone(), 1);
    c.begin();
    assert_eq!(job.step(&mut c, || true), Progress::Pending); // began_at = 2
    c.rollback(); // ticks back to 1
    assert!(job.interrupted(&c), "[r3] dead right after rollback");
    c.step(); // ticks 2 again
    let dead_after_step = job.interrupted(&c);
    let mut got = None;
    if !dead_after_step {
        while job.step(&mut c, || false) == Progress::Pending {}
        got = Some(c.current(&i));
    }
    assert!(
        dead_after_step,
        "[r3] interrupted job came back to life after a later step; finishing it gives {:?}, want {:?}",
        got, snap
    );
}

#[test]
fn r4_rebuild_whose_first_chunk_panics_in_a_txn_then_rolled_back() {
    let mut c = Circuit::new();
    let (input, s) = c.input::<u8>();
    let i = c.integrate(&s);
    c.push(&input, z(&[(1, i64::MAX)]));
    c.step();
    let snap = z(&[(1, 1), (2, 1), (3, 1)]);
    let mut job = Rebuild::new(input, snap.clone(), 1);
    c.begin();
    let r = catch_unwind(AssertUnwindSafe(|| job.step(&mut c, || false)));
    assert!(r.is_err());
    c.rollback();
    let interrupted = job.interrupted(&c);
    let mut got = None;
    if !interrupted {
        c.push(&input, z(&[(1, -5)])); // make room so the retry does not overflow
        while job.step(&mut c, || false) == Progress::Pending {}
        got = Some(c.current(&i));
    }
    assert!(
        interrupted,
        "[r4] after a panicked first chunk + rollback, the job is not interrupted; finishing gives {:?}",
        got
    );
}

#[test]
fn p1_a_chunk_whose_push_panics_in_a_txn_is_counted() {
    let mut c = Circuit::new();
    let (input, s) = c.input::<u8>();
    let i = c.integrate(&s);
    c.push(&input, z(&[(1, i64::MAX)]));
    let mut job = Rebuild::new(input, z(&[(1, 1), (2, 1)]), 1);
    c.begin();
    let r = catch_unwind(AssertUnwindSafe(|| job.step(&mut c, || false)));
    assert!(r.is_err(), "the push overflows");
    c.rollback();
    let interrupted = job.interrupted(&c);
    let mut got = None;
    if !interrupted {
        while job.step(&mut c, || false) == Progress::Pending {}
        got = Some(c.current(&i));
    }
    assert!(
        interrupted,
        "[p1] the chunk left the job unstepped; finishing it gives {got:?}"
    );
}