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.

src/lib.rs

@@ -42,3 +42,25 @@ pub use scheduler::{
     block_on_io, run, self_pid, sleep, spawn, spawn_under, yield_now, JoinError, JoinHandle,
 };
 pub use supervisor::Signal;
+
+// ---------------------------------------------------------------------------
+// check!() — explicit preemption point for tight no-alloc loops.
+// ---------------------------------------------------------------------------
+
+/// Voluntarily check whether this actor's timeslice has expired, yielding
+/// if so. Drop this into hot compute loops that don't allocate (heap or
+/// large stack frames) — without it, such loops monopolise the scheduler
+/// until they return.
+///
+/// Decrements the same per-actor event counter as the heap allocator's
+/// preemption hook, so the check rate is identical regardless of whether
+/// the actor is alloc-heavy, check-heavy, or mixed.
+///
+/// No-op outside an actor (the runtime's `PREEMPTION_ENABLED` flag is
+/// false there).
+#[macro_export]
+macro_rules! check {
+    () => {
+        $crate::preempt::maybe_preempt()
+    };
+}

src/preempt.rs

@@ -6,10 +6,16 @@
 //! `switch_to_scheduler` to yield. Resetting the counter to `ALLOC_INTERVAL`
 //! amortises the RDTSC across many cheap events.
 //!
-//! Events today are heap allocations (via `PreemptingAllocator`). v0.2 will
+//! Two event sources today:
-//! add stack-frame entries as a second event source — frames are stack
+//!   - `PreemptingAllocator` — heap allocations.
-//! allocations, the counter naming still fits — sharing this same counter
+//!   - `smarm::check!()` — explicit preemption point for tight no-alloc
-//! so both routes behave consistently.
+//!     loops, since stable Rust gives us no transparent way to preempt
+//!     such loops (`__rust_probestack` is emitted inline by LLVM and not
+//!     called at runtime).
+//!
+//! Both sources share `ALLOC_COUNT`, so the timeslice check fires at the
+//! same rate regardless of whether the actor is alloc-heavy, check-heavy,
+//! or mixed.
 //!
 //! All state is thread-local. The scheduler enables preemption on resume
 //! and disables it on the return path, so the scheduler can never preempt
@@ -80,9 +86,17 @@ unsafe impl GlobalAlloc for PreemptingAllocator {
 }
 
 /// Shared preemption check. Called by every preemption event source — the
-/// heap allocator today, the stack-frame entry hook in v0.2. Decrements
+/// heap allocator today, `smarm::check!()` for tight no-alloc loops.
-/// `ALLOC_COUNT`; every `ALLOC_INTERVAL` calls reads the timeslice clock
+/// Decrements `ALLOC_COUNT`; every `ALLOC_INTERVAL` calls reads the
-/// and yields if expired.
+/// timeslice clock and yields if expired.
+///
+/// **Invariant**: must not be called inside a "prep-to-park" region —
+/// e.g. between registering as a channel's parked receiver and calling
+/// `park_current()`. A preemption-driven yield in that window would
+/// reach the scheduler with state=Runnable, the unparker would no-op,
+/// the actor would then park, and the wakeup would be lost. Library
+/// code that touches the parking primitives must keep its prep-to-park
+/// regions allocation-free and check!()-free.
 #[inline(always)]
 pub fn maybe_preempt() {
     ALLOC_COUNT.with(|c| {

tests/preempt.rs

@@ -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);
+}