Skip to main content

ax_task/sync/spin/
raw.rs

1//! Raw spin locks that implement [`lock_api::RawMutex`].
2//!
3//! Unlike [`BaseSpinLock`](super::base::BaseSpinLock), these locks do not own the
4//! protected data; they only provide the lock/unlock primitive so they can be
5//! plugged into foreign generic code that is parameterised over
6//! `lock_api::RawMutex` (for example the `kprobe` crate's `ProbeManager`).
7//!
8//! The guard semantics still come from a [`GuardState`]: acquiring the lock
9//! runs `G::acquire()` (e.g. disabling preemption and local IRQs) *before*
10//! spinning, and releasing it restores that state. This matches the behaviour
11//! of [`SpinLock::lock_irqsave`](super::SpinLock::lock_irqsave) and is what makes the lock safe to take
12//! from contexts that may be re-entered by interrupts or trap handlers.
13
14#[cfg(feature = "smp")]
15use core::sync::atomic::AtomicBool;
16use core::{cell::UnsafeCell, marker::PhantomData};
17
18use crate::sync::context::{GuardState, PreemptIrqSaveState};
19
20/// A raw spin lock implementing [`lock_api::RawMutex`], whose critical-section
21/// guard behaviour is determined by the [`GuardState`] type parameter `G`.
22///
23/// On a single-core build (without the `smp` feature) the atomic flag is
24/// elided, but `G::acquire()`/`G::release()` are still run so preemption and
25/// IRQ state are managed correctly.
26pub struct BaseRawSpinLock<G: GuardState> {
27    _phantom: PhantomData<G>,
28
29    #[cfg(feature = "smp")]
30    locked: AtomicBool,
31
32    // Saved guard state from `G::acquire()`. Only the lock owner writes or
33    // reads this slot while the lock is held, so the lack of synchronisation
34    // is sound.
35    state: UnsafeCell<Option<G::State>>,
36}
37
38// The `UnsafeCell<Option<G::State>>` is only ever touched by the thread that
39// owns the lock, so the lock as a whole is `Sync`.
40unsafe impl<G: GuardState> Sync for BaseRawSpinLock<G> {}
41
42impl<G: GuardState> BaseRawSpinLock<G> {
43    /// Creates a new, unlocked raw spin lock.
44    pub const fn new() -> Self {
45        Self {
46            _phantom: PhantomData,
47            #[cfg(feature = "smp")]
48            locked: AtomicBool::new(false),
49            state: UnsafeCell::new(None),
50        }
51    }
52
53    #[inline]
54    fn save_state(&self, state: G::State) {
55        // SAFETY: called only by the thread that just acquired the lock.
56        unsafe {
57            *self.state.get() = Some(state);
58        }
59    }
60
61    #[inline]
62    fn take_state(&self) -> G::State {
63        // SAFETY: called only by the thread that currently holds the lock,
64        // which is the same thread that stored the state in `lock()`.
65        unsafe {
66            (*self.state.get())
67                .take()
68                .expect("raw spinlock unlocked without saved guard state")
69        }
70    }
71}
72
73impl<G: GuardState> Default for BaseRawSpinLock<G> {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79unsafe impl<G: GuardState + Send + Sync + 'static> lock_api::RawMutex for BaseRawSpinLock<G> {
80    const INIT: Self = Self::new();
81
82    type GuardMarker = lock_api::GuardNoSend;
83
84    fn lock(&self) {
85        let state = G::acquire();
86
87        #[cfg(feature = "smp")]
88        {
89            super::atomic::spin_acquire(&self.locked, || {
90                super::atomic::spin_try_acquire_weak(&self.locked)
91            });
92        }
93
94        self.save_state(state);
95    }
96
97    fn try_lock(&self) -> bool {
98        let state = G::acquire();
99
100        #[cfg(feature = "smp")]
101        {
102            if !super::atomic::spin_try_acquire_strong(&self.locked) {
103                G::release(state);
104                return false;
105            }
106        }
107
108        self.save_state(state);
109        true
110    }
111
112    unsafe fn unlock(&self) {
113        let state = self.take_state();
114
115        #[cfg(feature = "smp")]
116        super::atomic::spin_release(&self.locked);
117
118        G::release(state);
119    }
120}
121
122/// A raw spin lock that disables kernel preemption and local IRQs while held,
123/// mirroring [`SpinLock::lock_irqsave`](super::SpinLock::lock_irqsave) but exposed as a
124/// [`lock_api::RawMutex`] for use with foreign generic code.
125pub type RawIrqSaveMutex = BaseRawSpinLock<PreemptIrqSaveState>;