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 [`RawSpinLock::lock_irqsave`](super::RawSpinLock::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 [`RawSpinLock::lock_irqsave`](super::RawSpinLock::lock_irqsave) but exposed as a
124/// [`lock_api::RawMutex`] for use with foreign generic code.
125pub type RawIrqSaveMutex = BaseRawSpinLock<PreemptIrqSaveState>;