preempt: explicit check!() macro for no-alloc loops

Stable Rust emits stack probes inline (subq/movq/jne loop) rather than
calling __rust_probestack, so there's no transparent hook for stack-
frame preemption. Override of __rust_probestack links cleanly but never
runs. Falling back to an explicit check!() that users drop into hot
compute loops.

check!() decrements the same ALLOC_COUNT counter as the heap path, so
both event sources fire timeslice checks at the same rate. Documents
the prep-to-park invariant on maybe_preempt — library code that
registers a wakeup and then parks must keep that window alloc-free and
check-free, or a preemption-driven yield in the middle would lose the
wakeup.
This commit is contained in:
Claude
2026-05-22 05:37:04 +00:00
parent 51bfccc3c2
commit d3ab81b833
3 changed files with 109 additions and 7 deletions

66
tests/preempt.rs Normal file
View File

@@ -0,0 +1,66 @@
//! 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);
}