fix(zset): a rollback that took a rebuild's chunks interrupts the job for good (improvise-bw4p.5)

A monotone count of rollbacks that undid a step, recorded by the job before its first chunk.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VdfEGXA7aJ6ftMsCc9PJze

crates/zset-ffi/tests/rollback_abi.rs

@@ -0,0 +1,169 @@
+//! Transactions through the ABI: a rebuild across a rollback, callbacks, and poison.
+//! From the review of #355. Evidence: `context/test-evidence/zset-ffi.md`.
+
+use std::ffi::c_void;
+use std::ptr::null_mut;
+use zset_ffi::engine::Code;
+use zset_ffi::*;
+
+const OK: i32 = 0;
+
+fn zset(entries: &[(&[u8], i64)]) -> *mut ZsZSet {
+    let mut z = null_mut();
+    assert_eq!(unsafe { zs_zset_new(&mut z) }, OK);
+    for (k, w) in entries {
+        assert_eq!(unsafe { zs_zset_insert(z, k.as_ptr(), k.len(), *w) }, OK);
+    }
+    z
+}
+
+fn weight(z: *const ZsZSet, k: &[u8]) -> i64 {
+    let mut w = 0;
+    assert_eq!(
+        unsafe { zs_zset_weight(z, k.as_ptr(), k.len(), &mut w) },
+        OK
+    );
+    w
+}
+
+fn rig() -> (*mut ZsCircuit, u32, u32) {
+    let mut c = null_mut();
+    assert_eq!(unsafe { zs_circuit_new(&mut c) }, OK);
+    let mut i = 0;
+    assert_eq!(unsafe { zs_circuit_input(c, &mut i) }, OK);
+    let mut int = 0;
+    assert_eq!(unsafe { zs_circuit_integrate(c, i, &mut int) }, OK);
+    let mut s = 0;
+    assert_eq!(unsafe { zs_circuit_output(c, int, &mut s) }, OK);
+    (c, i, s)
+}
+
+extern "C" fn out_of_time(_: *mut c_void) -> i32 {
+    1
+}
+
+fn take(c: *mut ZsCircuit, s: u32) -> *mut ZsZSet {
+    let mut z = null_mut();
+    assert_eq!(unsafe { zs_circuit_take(c, s, &mut z) }, OK);
+    z
+}
+
+#[test]
+fn f1_rebuild_begun_before_a_txn_continued_inside_and_rolled_back() {
+    let (c, i, s) = rig();
+    let snap = zset(&[(b"a", 1), (b"b", 1), (b"c", 1), (b"d", 1)]);
+    let mut job = null_mut();
+    assert_eq!(unsafe { zs_rebuild_new(c, i, snap, 1, &mut job) }, OK);
+    let mut done = 0;
+    assert_eq!(
+        unsafe { zs_rebuild_step(job, c, Some(out_of_time), null_mut(), &mut done) },
+        OK
+    );
+    assert_eq!(unsafe { zs_circuit_begin(c) }, OK);
+    assert_eq!(
+        unsafe { zs_rebuild_step(job, c, Some(out_of_time), null_mut(), &mut done) },
+        OK
+    );
+    assert_eq!(unsafe { zs_circuit_rollback(c) }, OK);
+    let rc = unsafe { zs_rebuild_step(job, c, None, null_mut(), &mut done) };
+    let z = take(c, s);
+    let got = (
+        rc,
+        done,
+        weight(z, b"a"),
+        weight(z, b"b"),
+        weight(z, b"c"),
+        weight(z, b"d"),
+    );
+    unsafe { zs_zset_free(z) };
+    unsafe { zs_rebuild_free(job) };
+    unsafe { zs_zset_free(snap) };
+    unsafe { zs_circuit_free(c) };
+    assert_eq!(
+        got.0,
+        Code::Interrupted as i32,
+        "[f1] header: a rebuild stepped in a rolled-back txn reports ZS_ERR_INTERRUPTED; got (rc, done, a, b, c, d) = {got:?}"
+    );
+}
+
+extern "C" fn rolls_back(ctx: *mut c_void, _: *const u8, _: usize) -> i32 {
+    let (c, seen) = unsafe { &mut *(ctx as *mut (*mut ZsCircuit, i32)) };
+    *seen = unsafe { zs_circuit_rollback(*c) };
+    1
+}
+
+#[test]
+fn f2_a_rollback_from_inside_a_callback_is_refused_and_the_txn_stays_open() {
+    let (c, i, s) = rig();
+    let mut state: (*mut ZsCircuit, i32) = (c, 0);
+    let mut f = 0;
+    let ctx = (&mut state as *mut (*mut ZsCircuit, i32)).cast();
+    assert_eq!(
+        unsafe { zs_circuit_filter(c, i, Some(rolls_back), ctx, &mut f) },
+        OK
+    );
+    assert_eq!(unsafe { zs_circuit_begin(c) }, OK);
+    let z = zset(&[(b"k", 1)]);
+    assert_eq!(unsafe { zs_circuit_push(c, i, z) }, OK);
+    assert_eq!(unsafe { zs_circuit_step(c) }, OK);
+    assert_eq!(
+        state.1,
+        Code::Poisoned as i32,
+        "[f2] rollback inside a callback"
+    );
+    assert_eq!(
+        unsafe { zs_circuit_commit(c) },
+        OK,
+        "[f2] txn still open, commit works"
+    );
+    let out = take(c, s);
+    assert_eq!(weight(out, b"k"), 1);
+    unsafe { zs_zset_free(out) };
+    unsafe { zs_zset_free(z) };
+    unsafe { zs_circuit_free(c) };
+}
+
+#[test]
+fn f3_begin_on_a_circuit_poisoned_outside_a_txn_is_poisoned() {
+    let (c, i, _s) = rig();
+    let z = zset(&[(b"k", i64::MAX)]);
+    assert_eq!(unsafe { zs_circuit_push(c, i, z) }, OK);
+    assert_eq!(unsafe { zs_circuit_step(c) }, OK);
+    assert_eq!(unsafe { zs_circuit_push(c, i, z) }, OK);
+    assert_eq!(unsafe { zs_circuit_step(c) }, Code::Overflow as i32);
+    assert_eq!(unsafe { zs_circuit_begin(c) }, Code::Poisoned as i32);
+    assert_eq!(unsafe { zs_circuit_rollback(c) }, Code::Transaction as i32);
+    assert_eq!(unsafe { zs_circuit_commit(c) }, Code::Transaction as i32);
+    unsafe { zs_zset_free(z) };
+    unsafe { zs_circuit_free(c) };
+}
+
+#[test]
+fn f4_a_failed_rebuild_step_in_a_txn_then_rollback_then_resume() {
+    let (c, i, s) = rig();
+    let big = zset(&[(b"a", i64::MAX)]);
+    assert_eq!(unsafe { zs_circuit_push(c, i, big) }, OK);
+    assert_eq!(unsafe { zs_circuit_step(c) }, OK);
+    let o = take(c, s);
+    unsafe { zs_zset_free(o) };
+    let snap = zset(&[(b"a", 1), (b"b", 1), (b"c", 1)]);
+    let mut job = null_mut();
+    assert_eq!(unsafe { zs_rebuild_new(c, i, snap, 1, &mut job) }, OK);
+    assert_eq!(unsafe { zs_circuit_begin(c) }, OK);
+    let mut done = 0;
+    assert_eq!(
+        unsafe { zs_rebuild_step(job, c, None, null_mut(), &mut done) },
+        Code::Overflow as i32
+    );
+    assert_eq!(unsafe { zs_circuit_rollback(c) }, OK);
+    let rc = unsafe { zs_rebuild_step(job, c, None, null_mut(), &mut done) };
+    unsafe { zs_rebuild_free(job) };
+    unsafe { zs_zset_free(snap) };
+    unsafe { zs_zset_free(big) };
+    unsafe { zs_circuit_free(c) };
+    assert_eq!(
+        rc,
+        Code::Interrupted as i32,
+        "[f4] header: a rebuild stepped in a rolled-back txn reports ZS_ERR_INTERRUPTED; got rc={rc} done={done}"
+    );
+}

crates/zset/tests/rollback_rebuild.rs

@@ -0,0 +1,144 @@
+//! 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.clone(), 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(i.current());
+    }
+    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.clone(), 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 {:?}",
+        i.current()
+    );
+}
+
+#[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.clone(), 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(i.current());
+    }
+    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);
+    input.push(z(&[(1, i64::MAX)]));
+    c.step();
+    let snap = z(&[(1, 1), (2, 1), (3, 1)]);
+    let mut job = Rebuild::new(input.clone(), 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 {
+        input.push(z(&[(1, -5)])); // make room so the retry does not overflow
+        while job.step(&mut c, || false) == Progress::Pending {}
+        got = Some(i.current());
+    }
+    assert!(
+        interrupted,
+        "[r4] after a panicked first chunk + rollback, the job is not interrupted; finishing gives {:?}",
+        got
+    );
+}
+
+#[test]
+fn b1_a_closure_that_reads_a_stream_during_eval_does_not_panic() {
+    let mut c = Circuit::new();
+    let (input, s) = c.input::<u8>();
+    let s2 = s.clone();
+    let i = c.integrate(&s);
+    let i2 = i.clone();
+    // map reads its own source and another node's output during eval
+    let m = c.map(&s, move |k| {
+        k.wrapping_add(s2.current().len() as u8 + i2.current().len() as u8)
+    });
+    let m2 = m.clone();
+    // a later node reads an earlier node's output inside its own eval
+    let _f = c.filter(&i, move |_| m2.current().len() < 1000);
+    input.push(z(&[(1, 1), (2, 1)]));
+    c.step();
+    c.begin();
+    input.push(z(&[(3, 1)]));
+    c.step_continue();
+    c.rollback();
+    c.step();
+}