feat(slot_state): park-epoch in the slot word; consuming, epoch-matched unparks

Word repacks to (gen << 32) | (epoch << 8) | state. begin_wait opens a wait
identity; every successful wake consumes it, so at most one wake lands per
wait by construction. unpark grows an Option<epoch> match; RuntimeInner gains
unpark_at. Claim/yield/park transitions become epoch-preserving CAS loops.

Loom: all four prior theorems re-proved on the new word, plus
consumed_epoch_unpark_never_lands — a late waker stamped with a consumed
epoch can neither enqueue nor notify, in every interleaving (the property
select's loser arms and every one-shot park site lean on).

No callers stamped yet; behaviour identical. Two transient dead_code
warnings (begin_wait, unpark_at) are consumed by the next commit.
This commit is contained in:
smarm
2026-06-10 07:07:44 +00:00
parent aa295582c4
commit 4913835c02
2 changed files with 267 additions and 79 deletions
+238 -70
View File
@@ -1,9 +1,9 @@
//! The per-slot scheduling state machine, as a standalone unit.
//!
//! One atomic word packs `(generation << 32) | state`; every transition is a
//! CAS on the packed word, so the generation check is atomic with the
//! transition — no ABA, no acting on a recycled slot. The diagram and the
//! full protocol rationale live in `runtime.rs`; this module is the
//! One atomic word packs `(generation << 32) | (epoch << 8) | state`; every
//! transition is a CAS on the packed word, so the generation check is atomic
//! with the transition — no ABA, no acting on a recycled slot. The diagram
//! and the full protocol rationale live in `runtime.rs`; this module is the
//! mechanism, factored out so that:
//!
//! - loom can model-check the production transitions directly (see the
@@ -11,6 +11,34 @@
//! - every method asserts the precondition it relies on (`debug_assert!` —
//! these are hot paths), per the assert-the-invariants house rule.
//!
//! ## The park-epoch (wait identity)
//!
//! The middle 24 bits carry the slot's *park-epoch*: the identity of the
//! actor's current (or most recent) wait. The rules:
//!
//! - [`begin_wait`](StateWord::begin_wait) bumps the epoch and returns it;
//! the actor calls it once per wait, *before* registering itself with any
//! waker. Registrations carry `(pid, epoch)`.
//! - A wake may be **epoch-matched** (`unpark(gen, Some(epoch))`): it lands
//! only if the word still carries that epoch. Wakers whose registration
//! handle can outlive the wait it was created for (channel senders, mutex
//! grants, wait-timers) MUST use this form.
//! - Every successful wake **consumes** the epoch — `Parked(e) → Queued(e+1)`,
//! `Running(e) → RunningNotified(e+1)` — so at most one wake can ever land
//! per wait, by construction. A loser in a multi-waker race (e.g. the
//! non-winning arms of a `select`) fails the epoch check and no-ops; it can
//! neither steal a future wait's wake nor leave a pending notification that
//! would fault a later one-shot park (`Mutex::lock_timeout`, `sleep`,
//! `block_on_io`, `wait_fd` all rely on wakes being *meaningful*).
//! - The wildcard form (`unpark(gen, None)`) also consumes, and is reserved
//! for terminal wakes — `request_stop` — which never return control to the
//! code that parked.
//!
//! Epoch wrap (24 bits = 16.7M waits) is harmless: a collision would require
//! a taken registration to stay in flight across a full wrap of the *same
//! actor's* waits, and registrations are consumed at take-time under their
//! primitive's lock — the exposure is the taker's instruction window.
//!
//! Atomics come from `sync_shim` (std normally, `loom::sync` under
//! `cfg(loom)`).
@@ -23,15 +51,23 @@ pub(crate) const ST_RUNNING_NOTIFIED: u64 = 3;
pub(crate) const ST_PARKED: u64 = 4;
pub(crate) const ST_DONE: u64 = 5;
/// Park-epoch width: 24 bits, packed at word bits 8..32.
pub(crate) const EPOCH_MASK: u32 = 0x00FF_FFFF;
#[inline]
pub(crate) const fn pack(gen: u32, st: u64) -> u64 {
((gen as u64) << 32) | st
pub(crate) const fn pack(gen: u32, epoch: u32, st: u64) -> u64 {
debug_assert!(epoch & !EPOCH_MASK == 0);
((gen as u64) << 32) | ((epoch as u64) << 8) | st
}
#[inline]
pub(crate) const fn word_gen(w: u64) -> u32 {
(w >> 32) as u32
}
#[inline]
pub(crate) const fn word_epoch(w: u64) -> u32 {
((w >> 8) as u32) & EPOCH_MASK
}
#[inline]
pub(crate) const fn word_state(w: u64) -> u64 {
w & 0xFF
}
@@ -44,7 +80,8 @@ pub(crate) enum Unpark {
Enqueue,
/// Running → RunningNotified: the scheduler's park-return will re-queue.
Notified,
/// Stale generation, already queued/notified, done, or vacant.
/// Stale generation, stale epoch, already queued/notified, done, or
/// vacant.
Noop,
}
@@ -64,7 +101,7 @@ pub(crate) struct StateWord(AtomicU64);
impl StateWord {
pub(crate) fn new() -> Self {
Self(AtomicU64::new(pack(0, ST_VACANT)))
Self(AtomicU64::new(pack(0, 0, ST_VACANT)))
}
#[inline]
@@ -99,91 +136,164 @@ impl StateWord {
/// Spawn-side publish: Vacant → Queued. The caller owns the vacant slot
/// exclusively (it popped the index from the free list), so this is a
/// plain Release store; it is the moment the actor becomes visible to
/// pops, unparks, and stops.
/// pops, unparks, and stops. The epoch starts at 0 for each occupancy
/// (`set_done` zeroes it; wait identity never crosses a lifetime).
pub(crate) fn publish_queued(&self, gen: u32) {
debug_assert_eq!(
self.load(),
pack(gen, ST_VACANT),
pack(gen, 0, ST_VACANT),
"publish over a non-vacant slot"
);
self.0.store(pack(gen, ST_QUEUED), Ordering::Release);
self.0.store(pack(gen, 0, ST_QUEUED), Ordering::Release);
}
/// Scheduler pop-side claim: Queued → Running. `false` means the popped
/// pid is stale — by the at-most-once-enqueued invariant, a generation
/// mismatch is the only possible failure (asserted).
/// Scheduler pop-side claim: Queued → Running, epoch preserved. `false`
/// means the popped pid is stale — by the at-most-once-enqueued
/// invariant, a generation mismatch is the only possible failure
/// (asserted). Nothing can move a matching-gen word off Queued (wakes
/// no-op on Queued), so the CAS loop is single-shot in practice.
#[must_use]
pub(crate) fn try_claim(&self, gen: u32) -> bool {
match self.0.compare_exchange(
pack(gen, ST_QUEUED),
pack(gen, ST_RUNNING),
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => true,
Err(actual) => {
debug_assert_ne!(
word_gen(actual),
gen,
"queued pid found in unexpected state {} — double enqueue?",
word_state(actual)
);
false
loop {
let w = self.load();
if word_gen(w) != gen {
return false;
}
debug_assert_eq!(
word_state(w),
ST_QUEUED,
"queued pid found in unexpected state {} — double enqueue?",
word_state(w)
);
if self
.0
.compare_exchange(
w,
pack(gen, word_epoch(w), ST_RUNNING),
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok()
{
return true;
}
}
}
/// Yield return path: Running | RunningNotified → Queued. A notification
/// that arrived mid-run coalesces into the re-queue. Caller must enqueue.
/// Yield return path: Running | RunningNotified → Queued, epoch
/// preserved. A notification that arrived mid-run coalesces into the
/// re-queue. Caller must enqueue. CAS loop because a notify can bump the
/// epoch between the read and the exchange.
pub(crate) fn yield_return(&self, gen: u32) {
let prev = self.0.swap(pack(gen, ST_QUEUED), Ordering::AcqRel);
debug_assert!(
matches!(word_state(prev), ST_RUNNING | ST_RUNNING_NOTIFIED)
&& word_gen(prev) == gen,
"yield return from invalid word {prev:#x}"
);
loop {
let w = self.load();
debug_assert!(
matches!(word_state(w), ST_RUNNING | ST_RUNNING_NOTIFIED)
&& word_gen(w) == gen,
"yield return from invalid word {w:#x}"
);
if self
.0
.compare_exchange(
w,
pack(gen, word_epoch(w), ST_QUEUED),
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok()
{
return;
}
}
}
/// Park return path. `true` = actually parked. `false` = an unpark landed
/// in the prep-to-park window (RunningNotified); the word is already back
/// to Queued and the caller must enqueue — the lost-wakeup window, closed.
/// to Queued and the caller must enqueue — the lost-wakeup window,
/// closed. Epoch preserved on both paths (the notify already consumed
/// it).
#[must_use]
pub(crate) fn park_return(&self, gen: u32) -> bool {
match self.0.compare_exchange(
pack(gen, ST_RUNNING),
pack(gen, ST_PARKED),
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => true,
Err(actual) => {
debug_assert_eq!(
actual,
pack(gen, ST_RUNNING_NOTIFIED),
"park return from invalid word {actual:#x}"
);
self.0.store(pack(gen, ST_QUEUED), Ordering::Release);
false
loop {
let w = self.load();
debug_assert_eq!(word_gen(w), gen, "park return with stale gen");
let target = match word_state(w) {
ST_RUNNING => ST_PARKED,
ST_RUNNING_NOTIFIED => ST_QUEUED,
st => unreachable!("park return from invalid state {st}"),
};
if self
.0
.compare_exchange(
w,
pack(gen, word_epoch(w), target),
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok()
{
return target == ST_PARKED;
}
}
}
/// Open a new wait: bump the park-epoch and return it. Called by the
/// waiting actor itself (so the state is Running, or RunningNotified if
/// a terminal wake is already pending — the bump preserves the pending
/// notification), once per wait, BEFORE registering `(pid, epoch)` with
/// any waker.
#[must_use]
pub(crate) fn begin_wait(&self, gen: u32) -> u32 {
loop {
let w = self.load();
debug_assert!(
matches!(word_state(w), ST_RUNNING | ST_RUNNING_NOTIFIED)
&& word_gen(w) == gen,
"begin_wait from invalid word {w:#x}"
);
let next = word_epoch(w).wrapping_add(1) & EPOCH_MASK;
if self
.0
.compare_exchange(
w,
pack(gen, next, word_state(w)),
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok()
{
return next;
}
}
}
/// The unpark protocol — the one way anything outside the scheduler makes
/// an actor runnable. See [`Unpark`] for the caller's obligations.
///
/// `want = Some(epoch)` is the epoch-matched form: lands only if the word
/// still carries that epoch (i.e. the wait it was registered for is still
/// the current, un-woken wait). `want = None` is the wildcard, reserved
/// for terminal wakes. Both forms CONSUME the epoch on success.
#[must_use]
pub(crate) fn unpark(&self, gen: u32) -> Unpark {
pub(crate) fn unpark(&self, gen: u32, want: Option<u32>) -> Unpark {
loop {
let w = self.load();
if word_gen(w) != gen {
return Unpark::Noop;
}
if let Some(e) = want {
if word_epoch(w) != e {
return Unpark::Noop;
}
}
let bumped = word_epoch(w).wrapping_add(1) & EPOCH_MASK;
match word_state(w) {
ST_PARKED => {
if self
.0
.compare_exchange(
w,
pack(gen, ST_QUEUED),
pack(gen, bumped, ST_QUEUED),
Ordering::AcqRel,
Ordering::Acquire,
)
@@ -197,7 +307,7 @@ impl StateWord {
.0
.compare_exchange(
w,
pack(gen, ST_RUNNING_NOTIFIED),
pack(gen, bumped, ST_RUNNING_NOTIFIED),
Ordering::AcqRel,
Ordering::Acquire,
)
@@ -211,12 +321,13 @@ impl StateWord {
}
}
/// Finalize: Running | RunningNotified → Done. Called by the scheduler
/// that just ran the actor to completion (so those are the only legal
/// prior states), under the slot's cold lock so join's check-or-register
/// is linearized against it.
/// Finalize: Running | RunningNotified → Done, epoch zeroed (wait
/// identity never crosses an occupancy). Called by the scheduler that
/// just ran the actor to completion (so those are the only legal prior
/// states), under the slot's cold lock so join's check-or-register is
/// linearized against it.
pub(crate) fn set_done(&self, gen: u32) {
let prev = self.0.swap(pack(gen, ST_DONE), Ordering::AcqRel);
let prev = self.0.swap(pack(gen, 0, ST_DONE), Ordering::AcqRel);
debug_assert!(
matches!(word_state(prev), ST_RUNNING | ST_RUNNING_NOTIFIED)
&& word_gen(prev) == gen,
@@ -230,11 +341,11 @@ impl StateWord {
pub(crate) fn reclaim(&self, gen: u32) {
debug_assert_eq!(
self.load(),
pack(gen, ST_DONE),
pack(gen, 0, ST_DONE),
"reclaim of a non-Done slot"
);
self.0
.store(pack(gen.wrapping_add(1), ST_VACANT), Ordering::Release);
.store(pack(gen.wrapping_add(1), 0, ST_VACANT), Ordering::Release);
}
}
@@ -260,17 +371,18 @@ mod loom_tests {
let word = Arc::new(StateWord::new());
word.publish_queued(0);
assert!(word.try_claim(0)); // scheduler claimed: actor Running
let epoch = word.begin_wait(0); // actor opens the wait
let ready = Arc::new(AtomicBool::new(false));
let enqueues = Arc::new(AtomicUsize::new(0));
// Waker: make the condition true, then unpark.
// Waker: make the condition true, then wake the registered wait.
let w = word.clone();
let r = ready.clone();
let e = enqueues.clone();
let waker = thread::spawn(move || {
r.store(true, O::SeqCst);
if w.unpark(0) == Unpark::Enqueue {
if w.unpark(0, Some(epoch)) == Unpark::Enqueue {
e.fetch_add(1, O::SeqCst);
}
});
@@ -303,13 +415,15 @@ mod loom_tests {
}
/// Two concurrent unparkers, one parked actor: exactly one wins the
/// enqueue (at-most-once), regardless of interleaving.
/// enqueue (at-most-once), regardless of interleaving. Both stamped with
/// the live epoch — the consuming bump is what serializes them.
#[test]
fn two_unparkers_one_enqueue() {
loom::model(|| {
let word = Arc::new(StateWord::new());
word.publish_queued(0);
assert!(word.try_claim(0));
let epoch = word.begin_wait(0);
assert!(word.park_return(0)); // actor parked
let enqueues = Arc::new(AtomicUsize::new(0));
@@ -318,7 +432,7 @@ mod loom_tests {
let w = word.clone();
let e = enqueues.clone();
hs.push(thread::spawn(move || {
if w.unpark(0) == Unpark::Enqueue {
if w.unpark(0, Some(epoch)) == Unpark::Enqueue {
e.fetch_add(1, O::SeqCst);
}
}));
@@ -331,6 +445,60 @@ mod loom_tests {
});
}
/// The stale-epoch theorem — what `select`'s loser arms lean on. An
/// actor opens a wait, two registered wakers race it (against the park
/// itself, covering the prep-to-park window); afterwards the actor is
/// runnable exactly once, and a LATE waker still stamped with the
/// consumed epoch can neither enqueue nor notify — in every
/// interleaving. (Under wildcard semantics the late waker would corrupt
/// the actor's NEXT one-shot park; this is the theorem that buys
/// `Mutex::lock_timeout`/`sleep`/`block_on_io` their unchanged code.)
#[test]
fn consumed_epoch_unpark_never_lands() {
loom::model(|| {
let word = Arc::new(StateWord::new());
word.publish_queued(0);
assert!(word.try_claim(0));
let epoch = word.begin_wait(0);
// Two arms race the wake, concurrent with the park itself.
let enqueues = Arc::new(AtomicUsize::new(0));
let mut hs = Vec::new();
for _ in 0..2 {
let w = word.clone();
let e = enqueues.clone();
hs.push(thread::spawn(move || {
if w.unpark(0, Some(epoch)) == Unpark::Enqueue {
e.fetch_add(1, O::SeqCst);
}
}));
}
let mut runnable_via_notify = false;
if !word.park_return(0) {
runnable_via_notify = true; // notified in prep-to-park
}
for h in hs {
h.join().unwrap();
}
// Exactly one path made the actor runnable.
let direct = enqueues.load(O::SeqCst);
if runnable_via_notify {
assert_eq!(direct, 0, "woken twice: notify AND enqueue");
} else {
assert_eq!(direct, 1, "parked forever, or woken twice");
}
assert_eq!(word_state(word.load()), ST_QUEUED);
// The actor runs again. A waker still holding the OLD epoch —
// a select loser arm firing later — must be a strict no-op,
// not a pending notification.
assert!(word.try_claim(0));
assert_eq!(word.unpark(0, Some(epoch)), Unpark::Noop);
assert_eq!(word_state(word.load()), ST_RUNNING, "stale epoch notified a live run");
});
}
/// The ABA theorem: a stale-generation unpark racing reclaim + reuse can
/// never touch the slot's new occupant.
#[test]
@@ -342,7 +510,7 @@ mod loom_tests {
assert!(word.try_claim(0));
let w = word.clone();
let stale = thread::spawn(move || w.unpark(0));
let stale = thread::spawn(move || w.unpark(0, None));
// Scheduler: finalize, reclaim, and a new spawn reuses the slot.
word.set_done(0);
@@ -355,7 +523,7 @@ mod loom_tests {
// Either way it must never claim an enqueue.
assert_ne!(stale.join().unwrap(), Unpark::Enqueue);
// And the new occupant is exactly where its spawn put it.
assert_eq!(word.load(), pack(1, ST_QUEUED));
assert_eq!(word.load(), pack(1, 0, ST_QUEUED));
});
}
@@ -369,7 +537,7 @@ mod loom_tests {
word.publish_queued(0);
let w = word.clone();
let unparker = thread::spawn(move || w.unpark(0));
let unparker = thread::spawn(move || w.unpark(0, None));
assert!(word.try_claim(0)); // the entry is ours; claim must win
let r = unparker.join().unwrap();