smarm/src/raw_mutex.rs

//! A minimal futex-based mutex that cannot poison.
//!
//! `std::sync::Mutex` poisons on unwind, turning one panic into a cascade of
//! `lock().unwrap()` panics in every later user. The runtime's internal
//! critical sections must never unwind anyway (the stop sentinel is gated
//! behind `PREEMPTION_ENABLED`, and the guard below disables preemption), so
//! poisoning buys nothing and costs a failure mode. This mutex has no poison
//! state by construction.
//!
//! Two further properties the runtime wants:
//!
//! - **The guard enters `NoPreempt`.** A timeslice switch while holding an OS
//!   mutex would suspend the actor with the lock held, stalling every other
//!   OS thread that touches it until the actor is resumed. Disabling
//!   preemption for the (short) critical section keeps lock hold times
//!   bounded. It also closes the unwind hole structurally: with
//!   `PREEMPTION_ENABLED` false, `maybe_preempt` neither yields nor raises
//!   the stop sentinel, so no allocation inside the critical section can
//!   unwind it.
//! - **No std machinery.** One `AtomicU32` and two futex syscalls; friendlier
//!   to an eventual embedded port than `std::sync::Mutex` (swap the futex for
//!   a spin or WFE backend).
//!
//! Algorithm: the classic three-state futex mutex (Drepper, "Futexes Are
//! Tricky", mutex3). 0 = unlocked, 1 = locked, 2 = locked with (possible)
//! waiters. Uncontended lock/unlock is one CAS / one swap, no syscall.
//!
//! Lock-order position — two classes (see [`LockClass`]):
//!
//! - **Leaf**: slot cold locks, the free list, the stack pool, the name
//!   registry. Mutual leaves — never hold two at once.
//! - **Channel**: a channel's internal lock. May be acquired *under* a Leaf
//!   (finalize clones the supervisor/trap senders, and `monitor()` clones the
//!   Down sender, all under a cold lock — structural, the sender lives in the
//!   slot), but nothing may be acquired under a Channel lock: channel
//!   critical sections call only the lock-free unpark protocol.
//!
//! So the total order is Leaf → Channel, one of each at most. Holding either
//! while pushing to the run queue is permitted (unpark from inside a cold or
//! channel section); the reverse — taking any `RawMutex` from inside a
//! run-queue op — cannot arise (queue ops call nothing).

use std::cell::UnsafeCell;
use std::ops::{Deref, DerefMut};
use std::sync::atomic::{AtomicU32, Ordering};

const UNLOCKED: u32 = 0;
const LOCKED: u32 = 1;
const CONTENDED: u32 = 2;

/// How many `pause` spins to burn before falling back to the futex. Critical
/// sections under this lock are tens of nanoseconds (push to a Vec, clone a
/// sender), so a short spin almost always avoids the syscall.
const SPIN_LIMIT: u32 = 64;

/// Which rung of the two-rung lock order a `RawMutex` occupies. Debug builds
/// enforce the order mechanically (see the module docs); release builds carry
/// no state and no checks.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum LockClass {
    /// Runtime cold data: slot cold locks, free list, stack pool, registry.
    /// Mutual leaves among themselves; a Channel lock may be taken under one.
    Leaf,
    /// A channel's internal lock. One at a time, nothing acquired under it;
    /// may itself be acquired under a Leaf.
    Channel,
}

// The ordering rules, mechanically enforced (debug builds): a deadlock from a
// violated order is a hang waiting for the right interleaving, so fail at the
// acquisition that violates it, not in the eventual hang.
#[cfg(debug_assertions)]
thread_local! {
    static LEAVES_HELD: std::cell::Cell<u32> = const { std::cell::Cell::new(0) };
    static CHANNELS_HELD: std::cell::Cell<u32> = const { std::cell::Cell::new(0) };
}

#[inline]
fn order_check_acquire(class: LockClass) {
    #[cfg(debug_assertions)]
    match class {
        LockClass::Leaf => LEAVES_HELD.with(|l| {
            debug_assert_eq!(
                l.get(),
                0,
                "lock order violated: acquiring a Leaf RawMutex while already \
                 holding one (cold locks / free list / stack pool / registry \
                 are mutual leaves)"
            );
            CHANNELS_HELD.with(|c| {
                debug_assert_eq!(
                    c.get(),
                    0,
                    "lock order violated: acquiring a Leaf RawMutex under a \
                     channel lock (order is Leaf -> Channel, never the reverse)"
                );
            });
            l.set(l.get() + 1);
        }),
        LockClass::Channel => CHANNELS_HELD.with(|c| {
            debug_assert_eq!(
                c.get(),
                0,
                "lock order violated: acquiring a channel lock while already \
                 holding one (channel locks are mutual leaves)"
            );
            c.set(c.get() + 1);
        }),
    }
    #[cfg(not(debug_assertions))]
    let _ = class;
}

#[inline]
fn order_check_release(class: LockClass) {
    #[cfg(debug_assertions)]
    match class {
        LockClass::Leaf => LEAVES_HELD.with(|c| c.set(c.get() - 1)),
        LockClass::Channel => CHANNELS_HELD.with(|c| c.set(c.get() - 1)),
    }
    #[cfg(not(debug_assertions))]
    let _ = class;
}

pub(crate) struct RawMutex<T> {
    state: AtomicU32,
    class: LockClass,
    data: UnsafeCell<T>,
}

// SAFETY: standard mutex argument — exclusive access to `data` is mediated by
// `state`; `T: Send` suffices for both because `&RawMutex` only ever hands out
// access to one thread at a time.
unsafe impl<T: Send> Send for RawMutex<T> {}
unsafe impl<T: Send> Sync for RawMutex<T> {}

impl<T> RawMutex<T> {
    /// A Leaf-class mutex — the default for runtime cold data.
    pub(crate) const fn new(data: T) -> Self {
        Self::with_class(data, LockClass::Leaf)
    }

    /// A Channel-class mutex — for channel internals only.
    pub(crate) const fn new_channel(data: T) -> Self {
        Self::with_class(data, LockClass::Channel)
    }

    pub(crate) const fn with_class(data: T, class: LockClass) -> Self {
        Self {
            state: AtomicU32::new(UNLOCKED),
            class,
            data: UnsafeCell::new(data),
        }
    }

    #[inline]
    pub(crate) fn lock(&self) -> RawMutexGuard<'_, T> {
        // Enter NoPreempt *before* acquiring, so a preemption can't fire
        // between acquisition and guard construction.
        let prev_preempt = crate::preempt::PREEMPTION_ENABLED.with(|c| c.replace(false));
        order_check_acquire(self.class);
        if self
            .state
            .compare_exchange(UNLOCKED, LOCKED, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            self.lock_slow();
        }
        RawMutexGuard { m: self, prev_preempt }
    }

    #[cold]
    fn lock_slow(&self) {
        // Bounded spin first: the expected hold time is far below the cost of
        // a futex round trip.
        let mut spins = 0;
        loop {
            let s = self.state.load(Ordering::Relaxed);
            if s == UNLOCKED
                && self
                    .state
                    .compare_exchange_weak(UNLOCKED, LOCKED, Ordering::Acquire, Ordering::Relaxed)
                    .is_ok()
            {
                return;
            }
            spins += 1;
            if spins >= SPIN_LIMIT {
                break;
            }
            std::hint::spin_loop();
        }
        // Futex path. Mark contended and sleep until woken; on wake, retake
        // by swapping to CONTENDED (we cannot know whether other waiters
        // remain, so we must conservatively keep the contended marker).
        while self.state.swap(CONTENDED, Ordering::Acquire) != UNLOCKED {
            futex_wait(&self.state, CONTENDED);
        }
    }

    #[inline]
    fn unlock(&self) {
        if self.state.swap(UNLOCKED, Ordering::Release) == CONTENDED {
            futex_wake(&self.state, 1);
        }
    }
}

pub(crate) struct RawMutexGuard<'a, T> {
    m: &'a RawMutex<T>,
    prev_preempt: bool,
}

impl<T> Deref for RawMutexGuard<'_, T> {
    type Target = T;
    #[inline]
    fn deref(&self) -> &T {
        // SAFETY: guard existence implies exclusive ownership of the lock.
        unsafe { &*self.m.data.get() }
    }
}

impl<T> DerefMut for RawMutexGuard<'_, T> {
    #[inline]
    fn deref_mut(&mut self) -> &mut T {
        // SAFETY: as above, plus &mut self.
        unsafe { &mut *self.m.data.get() }
    }
}

impl<T> Drop for RawMutexGuard<'_, T> {
    #[inline]
    fn drop(&mut self) {
        self.m.unlock();
        order_check_release(self.m.class);
        // Restore preemption only after the lock is released.
        crate::preempt::PREEMPTION_ENABLED.with(|c| c.set(self.prev_preempt));
    }
}

// ---------------------------------------------------------------------------
// futex (x86-64 Linux; master is x86-only, see arm-port branch)
// ---------------------------------------------------------------------------

fn futex_wait(state: &AtomicU32, expected: u32) {
    // SAFETY: `state` is a valid, aligned u32 for the duration of the call.
    // Spurious wakeups and EAGAIN (value already changed) are both handled by
    // the caller's retry loop.
    unsafe {
        libc::syscall(
            libc::SYS_futex,
            state.as_ptr(),
            libc::FUTEX_WAIT | libc::FUTEX_PRIVATE_FLAG,
            expected,
            std::ptr::null::<libc::timespec>(),
        );
    }
}

fn futex_wake(state: &AtomicU32, n: i32) {
    // SAFETY: as above.
    unsafe {
        libc::syscall(
            libc::SYS_futex,
            state.as_ptr(),
            libc::FUTEX_WAKE | libc::FUTEX_PRIVATE_FLAG,
            n,
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;

    #[test]
    fn uncontended_lock_unlock() {
        let m = RawMutex::new(0u64);
        for _ in 0..1000 {
            *m.lock() += 1;
        }
        assert_eq!(*m.lock(), 1000);
    }

    #[test]
    fn contended_counter_is_exact() {
        const THREADS: usize = 8;
        const PER: u64 = 50_000;
        let m = Arc::new(RawMutex::new(0u64));
        let hs: Vec<_> = (0..THREADS)
            .map(|_| {
                let m = m.clone();
                std::thread::spawn(move || {
                    for _ in 0..PER {
                        *m.lock() += 1;
                    }
                })
            })
            .collect();
        for h in hs {
            h.join().unwrap();
        }
        assert_eq!(*m.lock(), THREADS as u64 * PER);
    }

    #[test]
    fn no_poison_on_unwind() {
        let m = Arc::new(RawMutex::new(0u64));
        let m2 = m.clone();
        let _ = std::thread::spawn(move || {
            let _g = m2.lock();
            panic!("unwind while holding");
        })
        .join();
        // A std Mutex would now be poisoned; this one just works.
        *m.lock() += 1;
        assert_eq!(*m.lock(), 1);
    }

    #[test]
    fn channel_lock_nests_under_leaf() {
        // The permitted ordering: Leaf -> Channel (finalize/monitor clone a
        // sender under a cold lock). Must not trip the order check.
        let leaf = RawMutex::new(0u64);
        let chan = RawMutex::new_channel(0u64);
        let _l = leaf.lock();
        let _c = chan.lock();
    }

    #[cfg(debug_assertions)]
    #[test]
    #[should_panic(expected = "lock order violated")]
    fn leaf_under_channel_is_rejected() {
        let leaf = RawMutex::new(0u64);
        let chan = RawMutex::new_channel(0u64);
        let _c = chan.lock();
        let _l = leaf.lock(); // Channel -> Leaf: forbidden
    }

    #[cfg(debug_assertions)]
    #[test]
    #[should_panic(expected = "lock order violated")]
    fn two_leaves_are_rejected() {
        let a = RawMutex::new(0u64);
        let b = RawMutex::new(0u64);
        let _ga = a.lock();
        let _gb = b.lock(); // leaves are mutual: forbidden
    }

    #[cfg(debug_assertions)]
    #[test]
    #[should_panic(expected = "lock order violated")]
    fn two_channel_locks_are_rejected() {
        let a = RawMutex::new_channel(0u64);
        let b = RawMutex::new_channel(0u64);
        let _ga = a.lock();
        let _gb = b.lock(); // channel locks are mutual leaves: forbidden
    }
}