cpu_local/

preempt.rs

use core::{
    marker::PhantomData,
    mem::ManuallyDrop,
    ptr::NonNull,
    sync::atomic::{AtomicU32, Ordering},
};

use crate::{CpuLocalError, CpuPin};

const PREEMPT_NO_PENDING: u32 = 1 << 31;
const PREEMPT_DEPTH_MASK: u32 = !PREEMPT_NO_PENDING;

/// Snapshot of one architecture-selected preemption state.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PreemptionSnapshot {
    depth: u32,
    pending: bool,
}

impl PreemptionSnapshot {
    pub(crate) const fn from_raw(state: u32) -> Self {
        Self {
            depth: state & PREEMPT_DEPTH_MASK,
            pending: state & PREEMPT_NO_PENDING == 0,
        }
    }

    /// Returns the number of active preemption exclusions.
    pub const fn depth(self) -> u32 {
        self.depth
    }

    /// Returns whether work is pending at the next preemptible boundary.
    pub const fn is_pending(self) -> bool {
        self.pending
    }
}

/// Linear proof of one entered preemption exclusion.
#[must_use = "every entered preemption token must be finished exactly once"]
pub struct PreemptionToken {
    #[cfg(any(test, not(target_arch = "x86_64"), feature = "host-test"))]
    owner: NonNull<PreemptionState>,
    _not_send_or_sync: PhantomData<*mut ()>,
}

/// Linear proof that the final preemption depth is reserved for a safe point.
#[must_use = "pending preemption must be released by the external safe-point owner"]
pub struct PendingPreemption {
    owner: NonNull<PreemptionState>,
    _not_send_or_sync: PhantomData<*mut ()>,
}

/// Result of finishing one preemption exclusion.
#[must_use]
pub enum PreemptionExit {
    /// A nested depth was consumed; preemption remains excluded.
    Nested,
    /// The final depth was consumed and execution is preemptible.
    Enabled,
    /// The final depth remains reserved until an external safe point claims it.
    Pending(PendingPreemption),
}

/// Architecture-neutral preemption word.
///
/// The high bit uses inverted pending polarity, allowing a newly initialized
/// context to start enabled without a runtime initialization write. The low
/// bits contain the exclusion depth. All updates are local to the selected CPU
/// or pinned context; external safe-point serialization provides ordering for
/// protected state, so this word does not publish cross-CPU data.
#[repr(transparent)]
pub(crate) struct PreemptionState(AtomicU32);

impl PreemptionState {
    pub(crate) const fn new() -> Self {
        Self(AtomicU32::new(PREEMPT_NO_PENDING))
    }

    pub(crate) const fn bootstrap_disabled() -> Self {
        Self(AtomicU32::new(PREEMPT_NO_PENDING | 1))
    }

    pub(crate) fn snapshot(&self) -> PreemptionSnapshot {
        PreemptionSnapshot::from_raw(self.0.load(Ordering::Relaxed))
    }

    fn set_pending(&self) {
        self.0.fetch_and(PREEMPT_DEPTH_MASK, Ordering::Relaxed);
    }

    fn clear_pending(&self) {
        self.0.fetch_or(PREEMPT_NO_PENDING, Ordering::Relaxed);
    }

    #[cfg(all(test, target_arch = "x86_64", not(feature = "host-test")))]
    pub(crate) fn as_mut_ptr(&self) -> *mut u32 {
        self.0.as_ptr()
    }

    #[cfg(any(test, not(target_arch = "x86_64"), feature = "host-test"))]
    #[inline(always)]
    fn compare_exchange_local(&self, current: u32, next: u32) -> bool {
        #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
        {
            // SAFETY: a live token retains positive depth on this exact
            // CPU-owned word. No remote CPU may update it; a local interrupt
            // can update it only before or after the single instruction.
            unsafe { crate::register::compare_exchange_x86_preemption_state(self, current, next) }
        }

        #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
        {
            self.0
                .compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed)
                .is_ok()
        }
    }

    #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
    fn enter(&self) {
        let previous = self.0.fetch_add(1, Ordering::Relaxed);
        assert_ne!(
            previous & PREEMPT_DEPTH_MASK,
            PREEMPT_DEPTH_MASK,
            "preemption nesting overflow"
        );
    }

    #[cfg(any(test, not(target_arch = "x86_64"), feature = "host-test"))]
    fn finish(&self) -> PreemptionExit {
        loop {
            let state = self.0.load(Ordering::Relaxed);
            let depth = state & PREEMPT_DEPTH_MASK;
            assert!(depth > 0, "unbalanced preemption exit");

            if depth == 1 && state & PREEMPT_NO_PENDING == 0 {
                return PreemptionExit::Pending(PendingPreemption::new(self));
            }

            // A local interrupt may update the sticky pending bit, and weak
            // LL/SC compare-exchange may also fail spuriously. Retry at a
            // fixed stack depth, matching Linux's bounded IRQ-return loop.
            if self.compare_exchange_local(state, state - 1) {
                return if depth == 1 {
                    PreemptionExit::Enabled
                } else {
                    PreemptionExit::Nested
                };
            }
        }
    }

    fn release_pending(&self) {
        assert_eq!(
            self.0
                .compare_exchange(1, 0, Ordering::Relaxed, Ordering::Relaxed),
            Ok(1),
            "pending preemption no longer owns the final depth"
        );
    }

    fn release_bootstrap(&self) {
        loop {
            let state = self.0.load(Ordering::Relaxed);
            assert!(
                state == (PREEMPT_NO_PENDING | 1) || state == 1,
                "bootstrap preemption depth must be released exactly once"
            );
            if self
                .0
                .compare_exchange_weak(state, state - 1, Ordering::Relaxed, Ordering::Relaxed)
                .is_ok()
            {
                return;
            }
        }
    }

    #[cfg(any(all(target_arch = "x86_64", not(feature = "host-test")), test))]
    fn release_initial_switch(&self) -> bool {
        loop {
            let state = self.0.load(Ordering::Relaxed);
            if state == PREEMPT_NO_PENDING {
                return false;
            }
            if state != (PREEMPT_NO_PENDING | 1) && state != 1 {
                panic!("initial context switch found invalid preemption state {state:#x}");
            }
            // A nested outgoing guard may mirror its owner's pending bit after
            // the switch guard was entered. The incoming context has its own
            // upper-layer publication, so consume the inherited depth and
            // reset that stale mirror as one transition.
            if self
                .0
                .compare_exchange_weak(
                    state,
                    PREEMPT_NO_PENDING,
                    Ordering::Relaxed,
                    Ordering::Relaxed,
                )
                .is_ok()
            {
                return true;
            }
        }
    }
}

impl PreemptionToken {
    #[cfg(any(test, not(target_arch = "x86_64"), feature = "host-test"))]
    fn new(owner: &PreemptionState) -> Self {
        Self {
            owner: NonNull::from(owner),
            _not_send_or_sync: PhantomData,
        }
    }

    #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
    fn new_current() -> Self {
        Self {
            _not_send_or_sync: PhantomData,
        }
    }

    /// Converts this token into an opaque value for an ABI that transports
    /// guard state as one machine word.
    #[doc(hidden)]
    pub fn into_raw(self) -> usize {
        #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
        {
            let _token = ManuallyDrop::new(self);
            1
        }
        #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
        {
            let token = ManuallyDrop::new(self);
            token.owner.as_ptr() as usize
        }
    }

    /// Reconstructs a token produced by [`Self::into_raw`].
    ///
    /// # Safety
    ///
    /// `raw` must come from one still-live token, refer to a live CPU-local
    /// preemption owner, and be reconstructed and finished exactly once.
    #[doc(hidden)]
    pub unsafe fn from_raw(raw: usize) -> Option<Self> {
        #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
        {
            (raw == 1).then(Self::new_current)
        }
        #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
        {
            if !raw.is_multiple_of(core::mem::align_of::<PreemptionState>()) {
                return None;
            }
            NonNull::new(raw as *mut PreemptionState).map(|owner| Self {
                owner,
                _not_send_or_sync: PhantomData,
            })
        }
    }

    #[cfg(any(test, not(target_arch = "x86_64"), feature = "host-test"))]
    fn state(&self) -> &PreemptionState {
        // SAFETY: construction and raw reconstruction require the selected
        // owner to remain live through this token's single finish operation.
        unsafe { self.owner.as_ref() }
    }

    #[cfg(any(all(target_arch = "x86_64", not(feature = "host-test")), test))]
    fn handoff_after_context_switch(self, resumed_owner: &PreemptionState) -> Self {
        #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
        {
            let _ = resumed_owner;
            self
        }
        #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
        {
            if self.owner == NonNull::from(resumed_owner) {
                self
            } else {
                // The old CPU's incoming context consumed the depth represented by
                // this proof. Consume the proof itself and adopt the equivalent
                // depth left by the outgoing context on the resumed CPU.
                Self::new(resumed_owner)
            }
        }
    }
}

impl PendingPreemption {
    fn new(owner: &PreemptionState) -> Self {
        Self {
            owner: NonNull::from(owner),
            _not_send_or_sync: PhantomData,
        }
    }

    /// Atomically consumes the final reserved preemption depth.
    pub fn release(self) {
        // SAFETY: the linear token retains its exact owner until this consume.
        unsafe { self.owner.as_ref() }.release_pending();
    }
}

/// Enters preemption exclusion on the owner selected by this architecture.
#[inline(always)]
pub fn enter_preemption() -> PreemptionToken {
    #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
    {
        // Increment through GS before resolving the token owner. Once the
        // increment is visible this execution cannot migrate away from it.
        unsafe { crate::register::enter_x86_preemption() };
        PreemptionToken::new_current()
    }

    #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
    {
        let current = unsafe { crate::current_context_unpinned() }
            .unwrap_or_else(|_| crate::register::fatal_register_invariant());
        // SAFETY: the architecture current register identifies the executing
        // context itself. An interrupt may suspend and later migrate that
        // context between this read and the increment, but it cannot resume
        // this instruction stream as a different context; the pinned header
        // stays live and therefore remains the exact token owner.
        let owner = unsafe { current.as_ref() }.preemption_state();
        owner.enter();
        PreemptionToken::new(owner)
    }
}

/// Observes the current architecture-selected preemption state.
#[inline(always)]
pub fn preemption_snapshot(pin: &CpuPin<'_>) -> Result<PreemptionSnapshot, CpuLocalError> {
    Ok(selected_state(pin)?.snapshot())
}

/// Tests the current execution context's advisory preemption-pending state.
///
/// This is an architecture-local safe-point query. It neither publishes nor
/// consumes external work. A `false` result is only a snapshot; a later
/// interrupt may publish work immediately after the observation.
#[inline(always)]
pub fn current_preemption_pending() -> Result<bool, CpuLocalError> {
    Ok(crate::register::current_preemption_snapshot()?.is_pending())
}

/// Marks work pending at the current preemptible boundary.
#[inline(always)]
pub fn set_preemption_pending(pin: &CpuPin<'_>) -> Result<(), CpuLocalError> {
    selected_state(pin)?.set_pending();
    Ok(())
}

/// Clears the pending mark after the external owner has drained its work.
#[inline(always)]
pub fn clear_preemption_pending(pin: &CpuPin<'_>) -> Result<(), CpuLocalError> {
    selected_state(pin)?.clear_pending();
    Ok(())
}

/// Finishes the exact owner captured by [`enter_preemption`].
#[inline(always)]
pub fn finish_preemption(token: PreemptionToken) -> PreemptionExit {
    #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
    {
        let _token = ManuallyDrop::new(token);
        finish_current_x86_preemption()
    }
    #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
    {
        token.state().finish()
    }
}

#[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
#[inline(always)]
fn finish_current_x86_preemption() -> PreemptionExit {
    // SAFETY: the consumed token proves a positive depth on the GS-selected
    // word until this transition completes.
    let state = unsafe { crate::register::read_current_x86_preemption_state_raw() };
    let depth = state & PREEMPT_DEPTH_MASK;
    assert!(depth > 0, "unbalanced preemption exit");

    if depth == 1 {
        if state & PREEMPT_NO_PENDING == 0 {
            // SAFETY: the retained final depth still pins the current owner.
            let owner = unsafe { crate::register::current_x86_preemption_state() };
            return PreemptionExit::Pending(PendingPreemption::new(owner));
        }
        return if unsafe {
            crate::register::compare_exchange_current_x86_preemption_state(state, state - 1)
        } {
            PreemptionExit::Enabled
        } else {
            finish_current_x86_preemption()
        };
    }

    // SAFETY: a nested depth remains positive after the single local update.
    unsafe { crate::register::decrement_current_x86_preemption_state() };
    PreemptionExit::Nested
}

/// Transfers a CPU-owned switch exclusion to the CPU where its context resumed.
///
/// Context-owned tokens keep their original owner across migration. A CPU-owned
/// token is replaced only when a context switch resumes it on another CPU: the
/// old CPU's incoming context has already consumed the old switch depth, while
/// this CPU's outgoing context left the matching depth for the resumed guard.
///
/// # Errors
///
/// Returns an error when the pinned CPU has no valid selected preemption owner.
///
/// # Panics
///
/// Panics when a context-owned architecture observes a different owner after
/// the context switch. Such architectures migrate the owner with the context.
#[doc(hidden)]
pub fn handoff_preemption_after_context_switch(
    pin: &CpuPin<'_>,
    token: PreemptionToken,
) -> Result<PreemptionToken, CpuLocalError> {
    let owner = selected_state(pin)?;
    #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
    {
        Ok(token.handoff_after_context_switch(owner))
    }

    #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
    {
        assert_eq!(
            token.owner,
            NonNull::from(owner),
            "context-owned preemption token changed owner across a context switch"
        );
        Ok(token)
    }
}

/// Releases the single preemption depth inherited by a bootstrap context.
///
/// The caller uses this only after the context and every local safe-point
/// dependency have been published.
#[doc(hidden)]
pub fn release_bootstrap_preemption(pin: &CpuPin<'_>) -> Result<(), CpuLocalError> {
    selected_state(pin)?.release_bootstrap();
    Ok(())
}

/// Releases the preemption exclusion transferred to a context on its first
/// architecture switch.
///
/// CPU-owned preemption architectures carry the outgoing switch exclusion
/// across the raw transfer, but a new context has no suspended caller whose
/// guard can finish it. Context-owned architectures need no action because the
/// new header begins enabled.
#[doc(hidden)]
pub fn release_initial_context_preemption(pin: &CpuPin<'_>) -> Result<bool, CpuLocalError> {
    #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
    {
        Ok(selected_state(pin)?.release_initial_switch())
    }

    #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
    {
        let snapshot = selected_state(pin)?.snapshot();
        assert_eq!(
            snapshot.depth(),
            0,
            "new context-owned preemption state must start enabled"
        );
        Ok(false)
    }
}

#[inline(always)]
fn selected_state(pin: &CpuPin<'_>) -> Result<&'static PreemptionState, CpuLocalError> {
    #[cfg(all(target_arch = "x86_64", not(feature = "host-test")))]
    {
        Ok(pin.area().runtime_anchor().preemption_state())
    }

    #[cfg(any(not(target_arch = "x86_64"), feature = "host-test"))]
    {
        let current = crate::current_context(pin)?;
        // SAFETY: current publication keeps the pinned context live, and the
        // caller's CpuPin prevents it from changing during this observation.
        Ok(unsafe { current.as_ref() }.preemption_state())
    }
}

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

    fn enter_on(state: &PreemptionState) -> PreemptionToken {
        state.enter();
        PreemptionToken::new(state)
    }

    #[test]
    fn nested_and_final_exits_are_linear() {
        let state = PreemptionState::new();
        let outer = enter_on(&state);
        let inner = enter_on(&state);

        assert!(matches!(finish_preemption(inner), PreemptionExit::Nested));
        assert_eq!(state.snapshot().depth(), 1);
        assert!(matches!(finish_preemption(outer), PreemptionExit::Enabled));
        assert_eq!(state.snapshot().depth(), 0);
    }

    #[test]
    fn pending_exit_reserves_depth_until_release() {
        let state = PreemptionState::new();
        let token = enter_on(&state);
        state.set_pending();

        let PreemptionExit::Pending(pending) = finish_preemption(token) else {
            panic!("final pending exit must retain its depth");
        };
        assert_eq!(state.snapshot().depth(), 1);
        pending.release();
        assert_eq!(state.snapshot().depth(), 0);
        assert!(state.snapshot().is_pending());
    }

    #[test]
    fn bootstrap_starts_disabled() {
        let state = PreemptionState::bootstrap_disabled();
        assert_eq!(state.snapshot().depth(), 1);
        assert!(!state.snapshot().is_pending());
    }

    #[test]
    fn bootstrap_release_preserves_pending_external_work() {
        let state = PreemptionState::bootstrap_disabled();
        state.set_pending();

        state.release_bootstrap();

        assert_eq!(state.snapshot().depth(), 0);
        assert!(state.snapshot().is_pending());
    }

    #[test]
    fn initial_switch_discards_outgoing_pending_mirror() {
        let state = PreemptionState::bootstrap_disabled();
        state.set_pending();

        assert!(state.release_initial_switch());
        assert_eq!(state.snapshot().depth(), 0);
        assert!(!state.snapshot().is_pending());
    }

    #[test]
    #[should_panic(expected = "pending preemption no longer owns the final depth")]
    fn pending_depth_cannot_be_consumed_twice() {
        let state = PreemptionState(AtomicU32::new(1));
        let first = PendingPreemption::new(&state);
        let duplicate = PendingPreemption::new(&state);

        first.release();
        duplicate.release();
    }

    #[test]
    #[should_panic(expected = "bootstrap preemption depth must be released exactly once")]
    fn bootstrap_depth_cannot_be_released_twice() {
        let state = PreemptionState::bootstrap_disabled();

        state.release_bootstrap();
        state.release_bootstrap();
    }

    #[test]
    fn token_stays_bound_to_the_entry_owner() {
        let original_context = PreemptionState::new();
        let migrated_context = PreemptionState::new();
        let original_cpu = PreemptionState::new();
        let migrated_cpu = PreemptionState::new();

        let context_token = enter_on(&original_context);
        let cpu_token = enter_on(&original_cpu);
        migrated_context.enter();
        migrated_cpu.enter();

        assert!(matches!(
            finish_preemption(context_token),
            PreemptionExit::Enabled
        ));
        assert!(matches!(
            finish_preemption(cpu_token),
            PreemptionExit::Enabled
        ));
        assert_eq!(original_context.snapshot().depth(), 0);
        assert_eq!(original_cpu.snapshot().depth(), 0);
        assert_eq!(migrated_context.snapshot().depth(), 1);
        assert_eq!(migrated_cpu.snapshot().depth(), 1);
    }

    #[test]
    fn cpu_owned_switch_token_handoff_follows_the_resumed_cpu() {
        let original_cpu = PreemptionState::new();
        let resumed_cpu = PreemptionState::new();

        let suspended_switch = enter_on(&original_cpu);
        // Another incoming context consumes the switch depth left on the old
        // CPU while this context is suspended.
        assert!(original_cpu.release_initial_switch());
        // The outgoing context on the destination CPU leaves the equivalent
        // switch depth for the context that is about to resume there.
        resumed_cpu.enter();

        let resumed_switch = suspended_switch.handoff_after_context_switch(&resumed_cpu);
        assert!(matches!(
            finish_preemption(resumed_switch),
            PreemptionExit::Enabled
        ));
        assert_eq!(original_cpu.snapshot().depth(), 0);
        assert_eq!(resumed_cpu.snapshot().depth(), 0);
    }

    #[test]
    fn malformed_raw_owner_is_rejected() {
        // SAFETY: no token is reconstructed because the value is misaligned.
        assert!(unsafe { PreemptionToken::from_raw(1) }.is_none());
    }
}