Skip to main content

kernel_api/sync/
mutex.rs

1use core::mem::ManuallyDrop;
2use core::panic::Location;
3use core::sync::atomic::{AtomicPtr, AtomicU8, AtomicUsize, Ordering};
4use log::warn;
5
6/// A mutual exclusion primitive useful for protecting shared data.
7///
8/// This spinlock will block threads waiting for the lock to become available. The
9/// spinlock can also be statically initialized or created via a `new`
10/// constructor. Each spinlock has a type parameter which represents the data that
11/// it is protecting. The data can only be accessed through the RAII guards
12/// returned from `lock` and `try_lock`, which guarantees that the data is only
13/// ever accessed when the spinlock is locked.
14pub type Spinlock<T> = lock_api::Mutex<RawSpinlock, T>;
15
16/// An RAII implementation of a "scoped lock" of a spinlock. When this structure is
17/// dropped (falls out of scope), the lock will be unlocked.
18///
19/// The data protected by the spinlock can be accessed through this guard via its
20/// `Deref` and `DerefMut` implementations.
21pub type SpinlockGuard<'a, T> = lock_api::MutexGuard<'a, RawSpinlock, T>;
22
23/// An RAII spinlock guard returned by `SpinlockGuard::map`, which can point to a
24/// subfield of the protected data.
25pub type MappedSpinlockGuard<'a, T> = lock_api::MappedMutexGuard<'a, RawSpinlock, T>;
26
27/// Extension functions to [`SpinlockGuard`].
28pub trait SpinlockGuardExt {
29    /// Unlock the spinlock without enabling interrupts, regardless of whether interrupts were enabled
30    /// before the spinlock was locked.
31    fn unlock_no_interrupts(this: Self);
32}
33
34impl<T> SpinlockGuardExt for SpinlockGuard<'_, T> {
35    fn unlock_no_interrupts(this: Self) {
36        let this = ManuallyDrop::new(this);
37	    // SAFETY: We consume the guard, preventing a guard from existing after we unlock the mutex
38	    let spinlock = unsafe { Self::mutex(&this).raw() };
39	    // SAFETY: `self` consuming method on `SpinlockGuard` so this thread must own a guard
40	    unsafe { spinlock.unlock_no_interrupts() };
41    }
42}
43
44#[derive(Debug, Copy, Clone, Eq, PartialEq)]
45enum State {
46    Unlocked,
47    Locked,
48}
49
50impl State {
51    const fn const_into_u8(self) -> u8 {
52        match self {
53            State::Unlocked => 0,
54            State::Locked => 1,
55        }
56    }
57
58    const fn const_from_u8(value: u8) -> Result<Self, ()> {
59        match value {
60            0 => Ok(State::Unlocked),
61            1 => Ok(State::Locked),
62            _ => Err(())
63        }
64    }
65}
66
67impl From<State> for u8 {
68    fn from(value: State) -> Self {
69        value.const_into_u8()
70    }
71}
72
73impl TryFrom<u8> for State {
74    type Error = ();
75
76    fn try_from(value: u8) -> Result<Self, Self::Error> {
77        Self::const_from_u8(value)
78    }
79}
80
81// INVARIANT: `location` is always a valid 'static ref to a `Location<'static>` unless null
82#[derive(Debug)]
83#[doc(hidden)]
84pub struct RawSpinlock {
85    state: AtomicU8,
86    irq_state: AtomicUsize,
87    #[cfg(debug_assertions)] location: AtomicPtr<Location<'static>>,
88}
89
90impl RawSpinlock {
91	/// # Safety
92	///
93	/// This method may only be called if the mutex is held in the current context, i.e. it must
94	/// be paired with a successful call to [`lock`](`Self::lock`) or [`try_lock`](`Self::try_lock`).
95    unsafe fn unlock_no_interrupts(&self) {
96        let old_state = self.state.swap(State::Unlocked.into(), Ordering::Release);
97        let old_state = State::try_from(old_state).expect("Spinlock in undefined state");
98
99        match old_state {
100            State::Unlocked => unreachable!("Mutex was unlocked while unlocked"),
101            State::Locked => {},
102        }
103    }
104
105    #[cfg(debug_assertions)]
106    fn lock_location(&self) -> Option<&'static Location<'static>> {
107        let location = self.location.load(Ordering::Relaxed);
108	    // SAFETY: invariant of type that `location` is always valid unless null
109        unsafe { location.as_ref::<'static>() }
110    }
111}
112
113// SAFETY: locked state is only modified in `lock` and `unlock` functions which both check the existing state
114unsafe impl lock_api::RawMutex for RawSpinlock {
115    const INIT: Self = Self {
116        state: AtomicU8::new(State::Unlocked.const_into_u8()),
117        irq_state: AtomicUsize::new(0),
118        #[cfg(debug_assertions)] location: AtomicPtr::new(core::ptr::null_mut()),
119    };
120
121    type GuardMarker = lock_api::GuardNoSend; // Dropping guard on other core would cause interrupts to be enabled in the wrong place
122
123    #[track_caller]
124    fn lock(&self) {
125        let irq_state = crate::bridge::irq::disable();
126
127        let mut printed = true;
128        while self.state.compare_exchange_weak(
129            State::Unlocked.into(),
130            State::Locked.into(),
131            Ordering::Acquire,
132            Ordering::Relaxed
133        ).is_err() {
134            core::hint::spin_loop();
135            #[cfg(debug_assertions)] if printed {
136                printed = false;
137                warn!("locked at {:?}", self.lock_location());
138            }
139        }
140
141        self.irq_state.store(irq_state, Ordering::Relaxed);
142        #[cfg(debug_assertions)] self.location.store(core::ptr::from_ref(Location::caller()).cast_mut(), Ordering::Relaxed);
143    }
144
145    fn try_lock(&self) -> bool {
146        let irq_state = crate::bridge::irq::disable();
147        let success = self.state.compare_exchange(
148            State::Unlocked.into(),
149            State::Locked.into(),
150            Ordering::Acquire,
151            Ordering::Relaxed
152        ).is_ok();
153
154        if success { self.irq_state.store(irq_state, Ordering::Relaxed) }
155        else { crate::bridge::irq::set(irq_state) }
156
157        success
158    }
159
160    unsafe fn unlock(&self) {
161        let old_irq_state = self.irq_state.load(Ordering::Relaxed);
162        let old_state = self.state.swap(State::Unlocked.into(), Ordering::Release);
163        let old_state = State::try_from(old_state).expect("Spinlock in undefined state");
164
165        match old_state {
166            State::Unlocked => unreachable!("Mutex was unlocked while unlocked"),
167            State::Locked => crate::bridge::irq::set(old_irq_state),
168        }
169    }
170}