improvise/crates/zset/tests/clone_panic.rs

//! A panicking `Clone` anywhere in a transactional step is rolled back, and `MemTrace::apply`
//! survives one. Evidence: `context/test-evidence/zset.md`.

use std::cell::Cell;
use std::panic::{AssertUnwindSafe, catch_unwind};
use zset::{Circuit, Input, MemTrace, Output, Stream, Trace, 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);
            }
        }));
    });
}

thread_local! { static CLONE_TRIP: Cell<Option<u32>> = const { Cell::new(None) }; }

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct Boom(u8);

impl Clone for Boom {
    fn clone(&self) -> Self {
        CLONE_TRIP.with(|t| {
            if let Some(n) = t.get() {
                if n == 0 {
                    t.set(None);
                    panic!("clone tripped");
                }
                t.set(Some(n - 1));
            }
        });
        Boom(self.0)
    }
}

type B = (Boom, Boom);

struct Rig {
    c: Circuit,
    a: Input<B>,
    b: Input<B>,
    streams: Vec<Stream<Z<B>>>,
    agg: Stream<Z<(Boom, u8)>>,
    sinks: Vec<Output<B>>,
}

fn rig(which: u8) -> Rig {
    let mut c = Circuit::new();
    let (a, sa) = c.input::<B>();
    let (b, sb) = c.input::<B>();
    let mut streams = vec![sa.clone(), sb.clone()];
    if which & 1 != 0 {
        streams.push(c.integrate(&sa));
    }
    if which & 2 != 0 {
        streams.push(c.distinct(&sa));
    }
    if which & 4 != 0 {
        streams.push(c.join_indexed(&sa, &sb, |k, x, y| (Boom(k.0), Boom(x.0 ^ y.0))));
    }
    if which & 8 != 0 {
        streams.push(c.delay(&sa));
    }
    let agg = c.aggregate(&sa, |_, vals: &[(Boom, i64)]| Some(vals.len() as u8));
    let sinks = streams.iter().map(|s| c.output(s)).collect();
    Rig {
        c,
        a,
        b,
        streams,
        agg,
        sinks,
    }
}

impl Rig {
    fn view(&self) -> String {
        format!(
            "{:?} {:?}",
            self.streams.iter().map(|s| s.current()).collect::<Vec<_>>(),
            self.agg.current()
        )
    }
    fn drain(&self) -> String {
        format!(
            "{:?}",
            self.sinks.iter().map(|s| s.take()).collect::<Vec<_>>()
        )
    }
}

fn bz(v: &[(u8, u8, i64)]) -> Z<B> {
    v.iter()
        .map(|(k, x, w)| ((Boom(*k), Boom(*x)), *w))
        .collect()
}

fn clone_panic_sweep(which: u8) -> Vec<String> {
    quiet();
    let mut failures = vec![];
    for n in 0..400u32 {
        let (mut x, mut y) = (rig(which), rig(which));
        for r in [&x, &y] {
            r.a.push(bz(&[(1, 1, 1), (2, 2, 1)]));
            r.b.push(bz(&[(1, 3, 1), (3, 3, 1)]));
        }
        x.c.step();
        y.c.step();
        x.drain();
        let (v0, d0) = (x.view(), x.drain());
        y.drain();
        x.c.begin();
        x.a.push(bz(&[(1, 5, 1), (4, 4, 1), (2, 2, -1)]));
        x.b.push(bz(&[(4, 6, 1), (1, 3, -1)]));
        CLONE_TRIP.with(|t| t.set(Some(n)));
        let r = catch_unwind(AssertUnwindSafe(|| x.c.step()));
        let fired = CLONE_TRIP.with(|t| t.get()).is_none();
        CLONE_TRIP.with(|t| t.set(None));
        if !fired {
            assert!(r.is_ok());
            break;
        }
        x.c.rollback();
        let (v1, d1) = (x.view(), x.drain());
        if v1 != v0 || d1 != d0 {
            failures.push(format!(
                "n={n}: rollback state differs\n  view {v1}\n  want {v0}\n  drain {d1}"
            ));
            continue;
        }
        for r in [&x, &y] {
            // The keys the failed step touched come back, so state leaked by one shows.
            r.a.push(bz(&[
                (1, 7, 1),
                (4, 1, 1),
                (1, 5, 1),
                (4, 4, 1),
                (2, 2, -1),
            ]));
            r.b.push(bz(&[(4, 2, 1), (1, 9, 1)]));
        }
        x.c.step();
        y.c.step();
        if x.view() != y.view() || x.drain() != y.drain() {
            failures.push(format!("n={n}: post tick differs"));
        }
    }
    failures
}

#[test]
fn c1_a_clone_panic_anywhere_in_a_step_is_rolled_back() {
    let mut all = vec![];
    for which in [1u8, 2, 4, 8, 15] {
        for f in clone_panic_sweep(which) {
            all.push(format!("which={which} {f}"));
        }
    }
    assert!(
        all.is_empty(),
        "{} failures:\n{}",
        all.len(),
        all.join("\n")
    );
}

#[test]
fn c2_memtrace_value_clone_panic_leaves_an_empty_group() {
    quiet();
    let mut t: MemTrace<u8, Boom> = MemTrace::new();
    CLONE_TRIP.with(|c| c.set(Some(0)));
    let r = catch_unwind(AssertUnwindSafe(|| t.apply(&1, &Boom(1), 1)));
    CLONE_TRIP.with(|c| c.set(None));
    assert!(r.is_err());
    assert_eq!(
        t.keys(),
        0,
        "[c2] MemTrace::apply panicked but left a group"
    );
}