smarm/src/slot_state.rs

//! The per-slot scheduling state machine, as a standalone unit.
//!
//! 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
//!   `loom_tests` module; built with `RUSTFLAGS="--cfg loom"`), and
//! - 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)`).

use crate::sync_shim::{AtomicU64, Ordering};

pub(crate) const ST_VACANT: u64 = 0;
pub(crate) const ST_QUEUED: u64 = 1;
pub(crate) const ST_RUNNING: u64 = 2;
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, 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
}

/// What an unpark amounted to. The caller owns the side effects (enqueue,
/// trace events) — this module is pure state.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub(crate) enum Unpark {
    /// Parked → Queued: the caller must enqueue the pid.
    Enqueue,
    /// Running → RunningNotified: the scheduler's park-return will re-queue.
    Notified,
    /// Stale generation, stale epoch, already queued/notified, done, or
    /// vacant.
    Noop,
}

/// A pid's-eye view of the slot, for cold paths that hold the slot lock.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub(crate) enum Status {
    /// The generation no longer matches: the slot was reclaimed (and possibly
    /// reused) — the pid is stale.
    Stale,
    /// The actor terminated; its outcome is (or was) in the slot.
    Done,
    /// Alive in some scheduling state (Queued / Running / Notified / Parked).
    Live,
}

pub(crate) struct StateWord(AtomicU64);

impl StateWord {
    pub(crate) fn new() -> Self {
        Self(AtomicU64::new(pack(0, 0, ST_VACANT)))
    }

    #[inline]
    pub(crate) fn load(&self) -> u64 {
        self.0.load(Ordering::Acquire)
    }

    #[inline]
    pub(crate) fn generation(&self) -> u32 {
        word_gen(self.load())
    }

    #[inline]
    pub(crate) fn status_for(&self, gen: u32) -> Status {
        let w = self.load();
        if word_gen(w) != gen {
            return Status::Stale;
        }
        match word_state(w) {
            ST_DONE => Status::Done,
            // A matching generation on a Vacant slot is unreachable for any
            // ISSUED pid — reclaim bumps the generation in the very store
            // that vacates, and install publishes Queued before the pid
            // escapes. But `Pid::new` is public, so a forged / never-issued
            // pid (e.g. `Pid::new(5, 0)` against a fresh slab) can land
            // here; for those, "no such actor" is the correct total answer.
            ST_VACANT => Status::Stale,
            _ => Status::Live,
        }
    }

    /// 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. 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, 0, ST_VACANT),
            "publish over a non-vacant slot"
        );
        self.0.store(pack(gen, 0, ST_QUEUED), Ordering::Release);
    }

    /// 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 {
        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, 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) {
        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. Epoch preserved on both paths (the notify already consumed
    /// it).
    #[must_use]
    pub(crate) fn park_return(&self, gen: u32) -> bool {
        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, 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, bumped, ST_QUEUED),
                            Ordering::AcqRel,
                            Ordering::Acquire,
                        )
                        .is_ok()
                    {
                        return Unpark::Enqueue;
                    }
                }
                ST_RUNNING => {
                    if self
                        .0
                        .compare_exchange(
                            w,
                            pack(gen, bumped, ST_RUNNING_NOTIFIED),
                            Ordering::AcqRel,
                            Ordering::Acquire,
                        )
                        .is_ok()
                    {
                        return Unpark::Notified;
                    }
                }
                _ => return Unpark::Noop, // Queued | Notified | Done | Vacant
            }
        }
    }

    /// Eat a pending notification: RunningNotified → Running, epoch
    /// preserved; no-op on Running. Called by the RUNNING actor itself, on
    /// the no-park exit of a wait it registered for but never parked on
    /// (`select` returning a ready arm at registration time), AFTER bumping
    /// the epoch and BEFORE re-checking its stop flag:
    ///
    /// - post-bump, the only wakers that can have set RunningNotified are
    ///   ones stamped with the just-retired epoch (a select arm) or a
    ///   terminal wildcard (`request_stop`);
    /// - the caller's stop-flag check AFTER the clear catches the terminal
    ///   case (the flag is set before the wake fires), so eating its
    ///   notification loses nothing — and a stop arriving later re-notifies
    ///   a Running word as usual;
    /// - what remains eaten is exactly the stale arm wake that would
    ///   otherwise fault the actor's next one-shot park.
    ///
    /// Returns whether a notification was eaten.
    pub(crate) fn clear_notify(&self, gen: u32) -> bool {
        loop {
            let w = self.load();
            debug_assert!(
                matches!(word_state(w), ST_RUNNING | ST_RUNNING_NOTIFIED)
                    && word_gen(w) == gen,
                "clear_notify from invalid word {w:#x}"
            );
            if word_state(w) != ST_RUNNING_NOTIFIED {
                return false;
            }
            if self
                .0
                .compare_exchange(
                    w,
                    pack(gen, word_epoch(w), ST_RUNNING),
                    Ordering::AcqRel,
                    Ordering::Acquire,
                )
                .is_ok()
            {
                return true;
            }
        }
    }

    /// 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, 0, ST_DONE), Ordering::AcqRel);
        debug_assert!(
            matches!(word_state(prev), ST_RUNNING | ST_RUNNING_NOTIFIED)
                && word_gen(prev) == gen,
            "finalize from invalid word {prev:#x}"
        );
    }

    /// Reclaim: Done → Vacant(gen + 1). The generation bump IS the reclaim:
    /// every stale pid is dead from this store onwards. Caller holds the cold
    /// lock and has verified eligibility (asserted).
    pub(crate) fn reclaim(&self, gen: u32) {
        debug_assert_eq!(
            self.load(),
            pack(gen, 0, ST_DONE),
            "reclaim of a non-Done slot"
        );
        self.0
            .store(pack(gen.wrapping_add(1), 0, ST_VACANT), Ordering::Release);
    }
}

// ---------------------------------------------------------------------------
// loom model tests — RUSTFLAGS="--cfg loom" cargo test --lib --release
// ---------------------------------------------------------------------------

#[cfg(all(test, loom))]
mod loom_tests {
    use super::*;
    use loom::sync::atomic::{AtomicBool, AtomicUsize};
    use loom::sync::Arc;
    use loom::thread;
    use std::sync::atomic::Ordering as O;

    /// THE lost-wakeup theorem. A waiter registers a condition check then
    /// parks (as every parking site does); a waker sets the condition then
    /// unparks. In every interleaving the waiter must end up runnable —
    /// parked-forever-with-condition-set must be unreachable.
    #[test]
    fn no_lost_wakeup_park_vs_unpark() {
        loom::model(|| {
            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 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, Some(epoch)) == Unpark::Enqueue {
                    e.fetch_add(1, O::SeqCst);
                }
            });

            // Waiter (as the scheduler executes it): re-check the condition,
            // park only if still false; a Notified park-return re-queues.
            let parked = if ready.load(O::SeqCst) {
                false // condition already visible: doesn't park at all
            } else if word.park_return(0) {
                true
            } else {
                enqueues.fetch_add(1, O::SeqCst); // notified → re-queued
                false
            };

            waker.join().unwrap();

            let w = word.load();
            if parked {
                // Parked is only a FINAL state if the waker's unpark moved it
                // back to Queued (+ one enqueue). Parked-and-stays-parked
                // would be the lost wakeup.
                assert_eq!(word_state(w), ST_QUEUED, "lost wakeup: parked forever");
                assert_eq!(enqueues.load(O::SeqCst), 1);
            } else {
                // Never more than one enqueue (at-most-once-enqueued).
                assert!(enqueues.load(O::SeqCst) <= 1);
            }
        });
    }

    /// Two concurrent unparkers, one parked actor: exactly one wins the
    /// 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));
            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);
                    }
                }));
            }
            for h in hs {
                h.join().unwrap();
            }
            assert_eq!(enqueues.load(O::SeqCst), 1);
            assert_eq!(word_state(word.load()), ST_QUEUED);
        });
    }

    /// 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 retire theorem — `select`'s no-park exit. An actor opens a wait
    /// and registers, then finds an arm ready and returns WITHOUT parking;
    /// a loser arm's waker fires concurrently, stamped with the live epoch.
    /// The exit retires the wait (bump, then eat): in every interleaving
    /// the run ends on a clean Running word — no pending notification
    /// survives to fault the actor's next one-shot park — and the waker
    /// never enqueues.
    #[test]
    fn retire_eats_late_arm_notification() {
        loom::model(|| {
            let word = Arc::new(StateWord::new());
            word.publish_queued(0);
            assert!(word.try_claim(0));
            let epoch = word.begin_wait(0); // select opens + registers

            let w = word.clone();
            let waker = thread::spawn(move || w.unpark(0, Some(epoch)));

            // No-park exit: bump (invalidates in-flight wakes), then eat
            // (consumes one that already landed).
            let _ = word.begin_wait(0);
            word.clear_notify(0);

            assert_ne!(waker.join().unwrap(), Unpark::Enqueue);
            assert_eq!(
                word_state(word.load()),
                ST_RUNNING,
                "stale arm wake survived the retire"
            );
        });
    }

    /// The ABA theorem: a stale-generation unpark racing reclaim + reuse can
    /// never touch the slot's new occupant.
    #[test]
    fn stale_unpark_never_hits_reused_slot() {
        loom::model(|| {
            let word = Arc::new(StateWord::new());
            // Gen-0 actor runs to completion.
            word.publish_queued(0);
            assert!(word.try_claim(0));

            let w = word.clone();
            let stale = thread::spawn(move || w.unpark(0, None));

            // Scheduler: finalize, reclaim, and a new spawn reuses the slot.
            word.set_done(0);
            word.reclaim(0);
            word.publish_queued(1);

            // The stale unpark may have squeezed in only while gen 0 was
            // still Running (→ Notified) — in which case set_done's swap
            // absorbed it — or it observed Done/Vacant/gen-1 and no-op'd.
            // 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, 0, ST_QUEUED));
        });
    }

    /// Unpark racing the claim itself: whatever the interleaving, the actor
    /// is Running or RunningNotified afterwards and nobody enqueued (it was
    /// never parked).
    #[test]
    fn unpark_vs_claim_coalesces() {
        loom::model(|| {
            let word = Arc::new(StateWord::new());
            word.publish_queued(0);

            let w = word.clone();
            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();
            assert_ne!(r, Unpark::Enqueue);
            let st = word_state(word.load());
            assert!(matches!(st, ST_RUNNING | ST_RUNNING_NOTIFIED));
        });
    }
}