smarm/tests/preempt.rs

//! Tests for explicit preemption via `smarm::check!()`.

use smarm::{run, spawn};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

#[test]
fn check_yields_when_timeslice_expired() {
    // A single actor that drives the timeslice clock to zero manually,
    // then calls check!() and expects to yield. The scheduler has nothing
    // else to run, so it just re-queues us. To prove we actually yielded,
    // observe the run counter on the slot... we don't have one. So
    // instead: spawn a second actor that increments a counter and joins
    // it; verify both actors made progress in interleaved order under
    // forced timeslice expiry.
    let order: Arc<std::sync::Mutex<Vec<u8>>> = Arc::new(std::sync::Mutex::new(Vec::new()));
    let o1 = order.clone();
    let o2 = order.clone();

    run(move || {
        let a = spawn(move || {
            o1.lock().unwrap().push(b'A');
            // Force the timeslice to be considered expired.
            smarm::preempt::expire_timeslice_for_test();
            smarm::check!();
            o1.lock().unwrap().push(b'a');
        });
        let b = spawn(move || {
            o2.lock().unwrap().push(b'B');
            smarm::preempt::expire_timeslice_for_test();
            smarm::check!();
            o2.lock().unwrap().push(b'b');
        });
        a.join().unwrap();
        b.join().unwrap();
    });

    // FIFO scheduling + forced preemption: A starts, expires, yields to B;
    // B starts, expires, yields to A; A finishes, B finishes.
    // Required: both uppercase letters appear before either lowercase.
    let v = order.lock().unwrap();
    let pos_big_a = v.iter().position(|&c| c == b'A').unwrap();
    let pos_big_b = v.iter().position(|&c| c == b'B').unwrap();
    let pos_lit_a = v.iter().position(|&c| c == b'a').unwrap();
    let pos_lit_b = v.iter().position(|&c| c == b'b').unwrap();
    assert!(
        pos_big_a < pos_lit_a,
        "A's tail ran before B's head: {:?}",
        *v
    );
    assert!(
        pos_big_b < pos_lit_b,
        "B's tail ran before A's head: {:?}",
        *v
    );
    assert!(
        pos_big_a.max(pos_big_b) < pos_lit_a.min(pos_lit_b),
        "preemption didn't interleave: {:?}",
        *v
    );
}

#[test]
fn check_is_a_noop_when_timeslice_not_expired() {
    // After a fresh resume, check!() should be cheap and not yield. Run
    // a single actor that calls check!() many times; it should complete
    // promptly.
    let count = Arc::new(AtomicU64::new(0));
    let c = count.clone();
    run(move || {
        for _ in 0..1_000 {
            smarm::check!();
            c.fetch_add(1, Ordering::Relaxed);
        }
    });
    assert_eq!(count.load(Ordering::Relaxed), 1_000);
}

// ---------------------------------------------------------------------------
// XMM-not-saved assumption (context.rs): the context switch saves no SSE
// state. The justification is that every yield crosses a Rust `call` boundary
// (SysV AMD64: XMM0-15 caller-saved), so the compiler spills live XMM before
// the switch and reloads after. The non-obvious path is PREEMPTION: `check!()`
// inlines `maybe_preempt`, so the yield can fire mid-floating-point-loop, not
// at a syntactic call site. This is adversarial because the yield is buried
// inside the actor's own hot loop — but `switch_to_scheduler` is still an
// `extern "C"` call, so the spill still happens. Verified by disasm: the FP
// accumulators are spilled to (%rsp) immediately before the preempt path and
// reloaded after. This test guards against a future change to the yield path
// that bypasses that call boundary (which WOULD require saving XMM).
#[test]
fn live_xmm_survives_preemptive_switch() {
    // A floating-point loop with several live f64 accumulators the compiler
    // wants resident in XMM across iterations. We force a timeslice expiry on
    // a chosen iteration so a preemptive switch_to_scheduler fires while XMM
    // is live, then compare against the same workload run with no preemption.
    #[inline(never)]
    fn fp_workload(preempt_at: u64) -> f64 {
        let mut a = 1.0000001_f64;
        let mut b = 0.9999999_f64;
        let mut acc = 0.0_f64;
        for i in 0..50_000u64 {
            a = a * 1.0000003 + 0.0000001;
            b = b * 0.9999997 + 0.0000002;
            acc += a - b + (i as f64) * 1e-9;
            if i == preempt_at {
                smarm::preempt::expire_timeslice_for_test();
            }
            smarm::check!();
        }
        acc
    }

    // Reference: a preempt_at past the loop end => check!() never yields.
    let reference = fp_workload(u64::MAX);

    let got = Arc::new(AtomicU64::new(0));
    let g = got.clone();
    run(move || {
        // A second actor so the scheduler has somewhere to switch on yield.
        let _other = spawn(|| {
            for _ in 0..8 {
                smarm::preempt::expire_timeslice_for_test();
                smarm::check!();
            }
        });
        let h = spawn(move || {
            // Yield from several distinct points so the switch lands on
            // different live-XMM states.
            let mut total = 0.0_f64;
            for k in 0..8 {
                total += fp_workload(k * 6_000 + 100);
            }
            g.store(total.to_bits(), Ordering::SeqCst);
        });
        h.join().unwrap();
    });

    let reference_total: f64 = (0..8).map(|_| reference).sum();
    let got = f64::from_bits(got.load(Ordering::SeqCst));
    assert_eq!(
        got.to_bits(),
        reference_total.to_bits(),
        "XMM state diverged across preemptive switch: got {got:.12}, want {reference_total:.12}"
    );
}