From d3ab81b833aed20c2ed294833685dc76ee537b40 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 22 May 2026 05:37:04 +0000 Subject: [PATCH] preempt: explicit check!() macro for no-alloc loops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 22 ++++++++++++++++ src/preempt.rs | 28 +++++++++++++++----- tests/preempt.rs | 66 ++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 7 deletions(-) create mode 100644 tests/preempt.rs diff --git a/src/lib.rs b/src/lib.rs index 33896c0..d7a93ba 100644 --- a/src/lib.rs +++ b/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() + }; +} diff --git a/src/preempt.rs b/src/preempt.rs index a99cc70..4ef52b6 100644 --- a/src/preempt.rs +++ b/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 -//! add stack-frame entries as a second event source — frames are stack -//! allocations, the counter naming still fits — sharing this same counter -//! so both routes behave consistently. +//! Two event sources today: +//! - `PreemptingAllocator` — heap allocations. +//! - `smarm::check!()` — explicit preemption point for tight no-alloc +//! 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 -/// `ALLOC_COUNT`; every `ALLOC_INTERVAL` calls reads the timeslice clock -/// and yields if expired. +/// heap allocator today, `smarm::check!()` for tight no-alloc loops. +/// Decrements `ALLOC_COUNT`; every `ALLOC_INTERVAL` calls reads the +/// 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| { diff --git a/tests/preempt.rs b/tests/preempt.rs new file mode 100644 index 0000000..3a04288 --- /dev/null +++ b/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>> = 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); +}