Skip to main content

ostd/sync/
spin.rs

1// SPDX-License-Identifier: MPL-2.0
2
3use core::{
4    cell::UnsafeCell,
5    fmt,
6    marker::PhantomData,
7    ops::{Deref, DerefMut},
8    sync::atomic::{AtomicBool, Ordering},
9};
10
11use super::{LocalIrqDisabled, PreemptDisabled, guard::SpinGuardian};
12use crate::task::atomic_mode::AsAtomicModeGuard;
13
14/// A spin lock.
15///
16/// # Guard behavior
17///
18/// The type `G' specifies the guard behavior of the spin lock. While holding the lock,
19/// - if `G` is [`PreemptDisabled`], preemption is disabled;
20/// - if `G` is [`LocalIrqDisabled`], local IRQs are disabled.
21///
22/// The `G` can also be provided by other crates other than ostd,
23/// if it behaves similar like [`PreemptDisabled`] or [`LocalIrqDisabled`].
24///
25/// The guard behavior can be temporarily upgraded from [`PreemptDisabled`] to
26/// [`LocalIrqDisabled`] using the [`disable_irq`] method.
27///
28/// [`disable_irq`]: Self::disable_irq
29#[repr(transparent)]
30pub struct SpinLock<T: ?Sized, G = PreemptDisabled> {
31    phantom: PhantomData<G>,
32    /// Only the last field of a struct may have a dynamically sized type.
33    /// That's why `SpinLockInner` is put in the last field.
34    inner: SpinLockInner<T>,
35}
36
37struct SpinLockInner<T: ?Sized> {
38    lock: AtomicBool,
39    val: UnsafeCell<T>,
40}
41
42impl<T, G> SpinLock<T, G> {
43    /// Creates a new spin lock.
44    pub const fn new(val: T) -> Self {
45        let lock_inner = SpinLockInner {
46            lock: AtomicBool::new(false),
47            val: UnsafeCell::new(val),
48        };
49        Self {
50            phantom: PhantomData,
51            inner: lock_inner,
52        }
53    }
54}
55
56impl<T: ?Sized> SpinLock<T, PreemptDisabled> {
57    /// Converts the guard behavior from disabling preemption to disabling IRQs.
58    pub fn disable_irq(&self) -> &SpinLock<T, LocalIrqDisabled> {
59        let ptr = self as *const SpinLock<T, PreemptDisabled>;
60        let ptr = ptr as *const SpinLock<T, LocalIrqDisabled>;
61        // SAFETY:
62        // 1. The types `SpinLock<T, PreemptDisabled>`, `SpinLockInner<T>` and `SpinLock<T,
63        //    IrqDisabled>` have the same memory layout guaranteed by `#[repr(transparent)]`.
64        // 2. The specified memory location can be borrowed as an immutable reference for the
65        //    specified lifetime.
66        unsafe { &*ptr }
67    }
68}
69
70impl<T: ?Sized, G: SpinGuardian> SpinLock<T, G> {
71    /// Acquires the spin lock.
72    pub fn lock(&self) -> SpinLockGuard<'_, T, G> {
73        // Notice the guard must be created before acquiring the lock.
74        let inner_guard = G::guard();
75        self.acquire_lock();
76        SpinLockGuard {
77            guard: inner_guard,
78            lock: self,
79        }
80    }
81
82    /// Tries acquiring the spin lock immediately.
83    pub fn try_lock(&self) -> Option<SpinLockGuard<'_, T, G>> {
84        let inner_guard = G::guard();
85        if self.try_acquire_lock() {
86            let lock_guard = SpinLockGuard {
87                guard: inner_guard,
88                lock: self,
89            };
90            return Some(lock_guard);
91        }
92        None
93    }
94
95    /// Returns a mutable reference to the underlying data.
96    ///
97    /// This method is zero-cost: By holding a mutable reference to the lock, the compiler has
98    /// already statically guaranteed that access to the data is exclusive.
99    pub fn get_mut(&mut self) -> &mut T {
100        self.inner.val.get_mut()
101    }
102
103    fn acquire_lock(&self) {
104        while !self.try_acquire_lock() {
105            core::hint::spin_loop();
106        }
107    }
108
109    fn try_acquire_lock(&self) -> bool {
110        self.inner
111            .lock
112            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
113            .is_ok()
114    }
115
116    /// # Safety
117    ///
118    /// The caller must ensure that it has acquired the lock.
119    unsafe fn release_lock(&self) {
120        self.inner.lock.store(false, Ordering::Release);
121    }
122}
123
124impl<T: ?Sized + fmt::Debug, G> fmt::Debug for SpinLock<T, G> {
125    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
126        fmt::Debug::fmt(&self.inner.val, f)
127    }
128}
129
130// SAFETY: Only a single lock holder is permitted to access the inner data of `Spinlock`.
131unsafe impl<T: ?Sized + Send, G> Send for SpinLock<T, G> {}
132unsafe impl<T: ?Sized + Send, G> Sync for SpinLock<T, G> {}
133
134/// A guard that provides exclusive access to the data protected by a [`SpinLock`].
135#[clippy::has_significant_drop]
136#[must_use]
137pub struct SpinLockGuard<'a, T: ?Sized, G: SpinGuardian> {
138    guard: G::Guard,
139    lock: &'a SpinLock<T, G>,
140}
141
142impl<T: ?Sized, G: SpinGuardian> AsAtomicModeGuard for SpinLockGuard<'_, T, G> {
143    fn as_atomic_mode_guard(&self) -> &dyn crate::task::atomic_mode::InAtomicMode {
144        self.guard.as_atomic_mode_guard()
145    }
146}
147
148impl<T: ?Sized, G: SpinGuardian> Deref for SpinLockGuard<'_, T, G> {
149    type Target = T;
150
151    fn deref(&self) -> &T {
152        // SAFETY: The spin lock is locked, which provides exclusive access to the data.
153        unsafe { &*self.lock.inner.val.get() }
154    }
155}
156
157impl<T: ?Sized, G: SpinGuardian> DerefMut for SpinLockGuard<'_, T, G> {
158    fn deref_mut(&mut self) -> &mut Self::Target {
159        // SAFETY: The spin lock is locked, which provides exclusive access to the data.
160        unsafe { &mut *self.lock.inner.val.get() }
161    }
162}
163
164impl<T: ?Sized, G: SpinGuardian> Drop for SpinLockGuard<'_, T, G> {
165    fn drop(&mut self) {
166        // SAFETY: The spin lock is locked.
167        unsafe { self.lock.release_lock() };
168    }
169}
170
171impl<T: ?Sized + fmt::Debug, G: SpinGuardian> fmt::Debug for SpinLockGuard<'_, T, G> {
172    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
173        fmt::Debug::fmt(&**self, f)
174    }
175}
176
177// `G::Guard` is `!Send` and `!Sync`, so is `SpinLockGuard`.
178impl<T: ?Sized, G: SpinGuardian> !Send for SpinLockGuard<'_, T, G> {}
179impl<T: ?Sized, G: SpinGuardian> !Sync for SpinLockGuard<'_, T, G> {}