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: slot cold locks, the free list, and the stack pool
//! are all `RawMutex`es and all *leaves among themselves* — never hold two at
//! once. Holding a `RawMutex` while taking the run-queue mutex (`with_shared`)
//! is permitted (e.g. unpark from inside a cold section); the reverse —
//! taking any `RawMutex` from inside `with_shared` — is forbidden.

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;

pub(crate) struct RawMutex<T> {
    state: AtomicU32,
    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> {
    pub(crate) const fn new(data: T) -> Self {
        Self {
            state: AtomicU32::new(UNLOCKED),
            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));
        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();
        // 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);
    }
}