Skip to main content

ax_sync/spin/
mod.rs

1//! OS-independent non-sleeping lock wrappers.
2
3#[cfg(feature = "lock-api")]
4mod raw;
5
6use core::{
7    cell::UnsafeCell,
8    fmt,
9    marker::PhantomData,
10    ops::{Deref, DerefMut},
11    panic::Location,
12    sync::atomic::{AtomicBool, AtomicUsize},
13};
14
15#[cfg(feature = "lock-api")]
16pub use self::raw::*;
17use crate::interface::{
18    CONTEXT_PREEMPT, CONTEXT_PREEMPT_IRQSAVE, CONTEXT_RAW, LOCK_MODE_READ, LOCK_MODE_WRITE,
19    LockMetadata,
20};
21
22/// A non-sleeping mutual-exclusion lock.
23#[repr(C)]
24pub struct SpinLock<T: ?Sized> {
25    locked: AtomicBool,
26    metadata: LockMetadata,
27    data: UnsafeCell<T>,
28}
29
30/// A guard returned by any [`SpinLock`] acquisition method.
31///
32/// ```compile_fail
33/// fn require_send<T: Send>() {}
34/// require_send::<ax_sync::SpinLockGuard<'static, ()>>();
35/// ```
36pub struct SpinLockGuard<'a, T: ?Sized> {
37    lock: &'a SpinLock<T>,
38    context: u8,
39    context_state: usize,
40    _not_send: PhantomData<*mut ()>,
41}
42
43/// A guard returned by [`SpinLock::lock_irqsave`].
44pub type SpinLockIrqSaveGuard<'a, T> = SpinLockGuard<'a, T>;
45/// A guard returned by [`SpinLock::lock_raw`].
46pub type RawSpinLockGuard<'a, T> = SpinLockGuard<'a, T>;
47
48unsafe impl<T: ?Sized + Send> Send for SpinLock<T> {}
49unsafe impl<T: ?Sized + Send> Sync for SpinLock<T> {}
50
51impl<T> SpinLock<T> {
52    /// Creates an unlocked spin lock.
53    #[track_caller]
54    pub const fn new(data: T) -> Self {
55        Self {
56            locked: AtomicBool::new(false),
57            metadata: LockMetadata::new(),
58            data: UnsafeCell::new(data),
59        }
60    }
61
62    /// Consumes the lock and returns the protected value.
63    pub fn into_inner(self) -> T {
64        self.data.into_inner()
65    }
66}
67
68impl<T: ?Sized> SpinLock<T> {
69    #[inline(always)]
70    #[track_caller]
71    fn acquire(&self, context: u8, subclass: u32, is_try: bool) -> Option<SpinLockGuard<'_, T>> {
72        let result = crate::interface::spin_acquire(
73            &self.locked,
74            &self.metadata,
75            self as *const Self as *const () as usize,
76            context,
77            subclass,
78            is_try,
79            Location::caller(),
80        );
81        result.acquired().then(|| SpinLockGuard {
82            lock: self,
83            context,
84            context_state: result.context_state(),
85            _not_send: PhantomData,
86        })
87    }
88
89    /// Acquires the lock after disabling kernel preemption.
90    #[track_caller]
91    pub fn lock(&self) -> SpinLockGuard<'_, T> {
92        self.lock_nested(0)
93    }
94
95    /// Acquires the lock with a lockdep subclass.
96    #[track_caller]
97    pub fn lock_nested(&self, subclass: u32) -> SpinLockGuard<'_, T> {
98        self.acquire(CONTEXT_PREEMPT, subclass, false)
99            .expect("blocking spin acquisition returned failure")
100    }
101
102    /// Attempts to acquire the lock after disabling preemption.
103    #[track_caller]
104    pub fn try_lock(&self) -> Option<SpinLockGuard<'_, T>> {
105        self.acquire(CONTEXT_PREEMPT, 0, true)
106    }
107
108    /// Acquires after disabling preemption and saving/disabling IRQs.
109    #[track_caller]
110    pub fn lock_irqsave(&self) -> SpinLockIrqSaveGuard<'_, T> {
111        self.lock_irqsave_nested(0)
112    }
113
114    /// Acquires in IRQ-save mode with a lockdep subclass.
115    #[track_caller]
116    pub fn lock_irqsave_nested(&self, subclass: u32) -> SpinLockIrqSaveGuard<'_, T> {
117        self.acquire(CONTEXT_PREEMPT_IRQSAVE, subclass, false)
118            .expect("blocking IRQ-save spin acquisition returned failure")
119    }
120
121    /// Attempts an IRQ-save acquisition.
122    #[track_caller]
123    pub fn try_lock_irqsave(&self) -> Option<SpinLockIrqSaveGuard<'_, T>> {
124        self.acquire(CONTEXT_PREEMPT_IRQSAVE, 0, true)
125    }
126
127    /// Acquires without changing execution context.
128    ///
129    /// # Safety
130    ///
131    /// The caller must prevent same-CPU re-entry and concurrent access which
132    /// could violate exclusive ownership.
133    #[track_caller]
134    pub unsafe fn lock_raw(&self) -> RawSpinLockGuard<'_, T> {
135        self.acquire(CONTEXT_RAW, 0, false)
136            .expect("blocking raw spin acquisition returned failure")
137    }
138
139    /// Attempts a raw acquisition.
140    ///
141    /// # Safety
142    ///
143    /// The caller must uphold the same exclusion contract as
144    /// [`Self::lock_raw`].
145    #[track_caller]
146    pub unsafe fn try_lock_raw(&self) -> Option<RawSpinLockGuard<'_, T>> {
147        self.acquire(CONTEXT_RAW, 0, true)
148    }
149
150    /// Returns whether the lock appears held.
151    pub fn is_locked(&self) -> bool {
152        crate::interface::spin_is_locked(&self.locked)
153    }
154
155    /// Returns exclusive access without locking.
156    pub fn get_mut(&mut self) -> &mut T {
157        self.data.get_mut()
158    }
159
160    /// Releases a deliberately leaked preemption-mode guard.
161    ///
162    /// # Safety
163    ///
164    /// The caller must own exactly one forgotten guard and prove no reference
165    /// derived from it remains live.
166    #[doc(hidden)]
167    pub unsafe fn force_unlock(&self) {
168        crate::interface::spin_force_release(
169            &self.locked,
170            self as *const Self as *const () as usize,
171            CONTEXT_PREEMPT,
172        );
173    }
174}
175
176impl<T: Default> Default for SpinLock<T> {
177    fn default() -> Self {
178        Self::new(T::default())
179    }
180}
181
182impl<T: fmt::Debug> fmt::Debug for SpinLock<T> {
183    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
184        match self.try_lock() {
185            Some(guard) => f.debug_struct("SpinLock").field("data", &&*guard).finish(),
186            None => f
187                .debug_struct("SpinLock")
188                .field("data", &"<locked>")
189                .finish(),
190        }
191    }
192}
193
194impl<T: ?Sized> Deref for SpinLockGuard<'_, T> {
195    type Target = T;
196
197    fn deref(&self) -> &Self::Target {
198        // SAFETY: the provider granted this guard shared access under the
199        // exclusive lock acquisition.
200        unsafe { &*self.lock.data.get() }
201    }
202}
203
204impl<T: ?Sized> DerefMut for SpinLockGuard<'_, T> {
205    fn deref_mut(&mut self) -> &mut Self::Target {
206        // SAFETY: this guard uniquely represents the exclusive acquisition.
207        unsafe { &mut *self.lock.data.get() }
208    }
209}
210
211impl<T: ?Sized> Drop for SpinLockGuard<'_, T> {
212    fn drop(&mut self) {
213        crate::interface::spin_release(
214            &self.lock.locked,
215            self.lock as *const SpinLock<T> as *const () as usize,
216            self.context,
217            self.context_state,
218        );
219    }
220}
221
222/// A non-sleeping read-write lock.
223#[repr(C)]
224pub struct SpinRwLock<T: ?Sized> {
225    state: AtomicUsize,
226    metadata: LockMetadata,
227    data: UnsafeCell<T>,
228}
229
230/// A read guard returned by [`SpinRwLock`].
231///
232/// ```compile_fail
233/// fn require_send<T: Send>() {}
234/// require_send::<ax_sync::SpinRwLockReadGuard<'static, ()>>();
235/// ```
236pub struct SpinRwLockReadGuard<'a, T: ?Sized> {
237    lock: &'a SpinRwLock<T>,
238    context: u8,
239    context_state: usize,
240    _not_send: PhantomData<*mut ()>,
241}
242
243/// A write guard returned by [`SpinRwLock`].
244///
245/// ```compile_fail
246/// fn require_send<T: Send>() {}
247/// require_send::<ax_sync::SpinRwLockWriteGuard<'static, ()>>();
248/// ```
249pub struct SpinRwLockWriteGuard<'a, T: ?Sized> {
250    lock: &'a SpinRwLock<T>,
251    context: u8,
252    context_state: usize,
253    _not_send: PhantomData<*mut ()>,
254}
255
256/// An IRQ-save read guard.
257pub type SpinRwLockIrqSaveReadGuard<'a, T> = SpinRwLockReadGuard<'a, T>;
258/// An IRQ-save write guard.
259pub type SpinRwLockIrqSaveWriteGuard<'a, T> = SpinRwLockWriteGuard<'a, T>;
260/// A raw read guard.
261pub type RawSpinRwLockReadGuard<'a, T> = SpinRwLockReadGuard<'a, T>;
262/// A raw write guard.
263pub type RawSpinRwLockWriteGuard<'a, T> = SpinRwLockWriteGuard<'a, T>;
264
265unsafe impl<T: ?Sized + Send + Sync> Send for SpinRwLock<T> {}
266unsafe impl<T: ?Sized + Send + Sync> Sync for SpinRwLock<T> {}
267
268impl<T> SpinRwLock<T> {
269    /// Creates an unlocked spin read-write lock.
270    #[track_caller]
271    pub const fn new(data: T) -> Self {
272        Self {
273            state: AtomicUsize::new(0),
274            metadata: LockMetadata::new(),
275            data: UnsafeCell::new(data),
276        }
277    }
278
279    /// Consumes the lock and returns the protected value.
280    pub fn into_inner(self) -> T {
281        self.data.into_inner()
282    }
283}
284
285impl<T: ?Sized> SpinRwLock<T> {
286    #[track_caller]
287    fn acquire(&self, context: u8, mode: u8, is_try: bool) -> Option<usize> {
288        let result = crate::interface::rwlock_acquire(
289            &self.state,
290            &self.metadata,
291            self as *const Self as *const () as usize,
292            context,
293            mode,
294            is_try,
295            Location::caller(),
296        );
297        result.acquired().then(|| result.context_state())
298    }
299
300    #[track_caller]
301    fn read_with(&self, context: u8, is_try: bool) -> Option<SpinRwLockReadGuard<'_, T>> {
302        self.acquire(context, LOCK_MODE_READ, is_try)
303            .map(|context_state| SpinRwLockReadGuard {
304                lock: self,
305                context,
306                context_state,
307                _not_send: PhantomData,
308            })
309    }
310
311    #[track_caller]
312    fn write_with(&self, context: u8, is_try: bool) -> Option<SpinRwLockWriteGuard<'_, T>> {
313        self.acquire(context, LOCK_MODE_WRITE, is_try)
314            .map(|context_state| SpinRwLockWriteGuard {
315                lock: self,
316                context,
317                context_state,
318                _not_send: PhantomData,
319            })
320    }
321
322    /// Acquires a read guard after disabling preemption.
323    #[track_caller]
324    pub fn read(&self) -> SpinRwLockReadGuard<'_, T> {
325        self.read_with(CONTEXT_PREEMPT, false)
326            .expect("blocking spin read acquisition returned failure")
327    }
328
329    /// Attempts a read acquisition after disabling preemption.
330    #[track_caller]
331    pub fn try_read(&self) -> Option<SpinRwLockReadGuard<'_, T>> {
332        self.read_with(CONTEXT_PREEMPT, true)
333    }
334
335    /// Acquires a write guard after disabling preemption.
336    #[track_caller]
337    pub fn write(&self) -> SpinRwLockWriteGuard<'_, T> {
338        self.write_with(CONTEXT_PREEMPT, false)
339            .expect("blocking spin write acquisition returned failure")
340    }
341
342    /// Attempts a write acquisition after disabling preemption.
343    #[track_caller]
344    pub fn try_write(&self) -> Option<SpinRwLockWriteGuard<'_, T>> {
345        self.write_with(CONTEXT_PREEMPT, true)
346    }
347
348    /// Acquires an IRQ-save read guard.
349    #[track_caller]
350    pub fn read_irqsave(&self) -> SpinRwLockIrqSaveReadGuard<'_, T> {
351        self.read_with(CONTEXT_PREEMPT_IRQSAVE, false)
352            .expect("blocking IRQ-save read acquisition returned failure")
353    }
354
355    /// Attempts an IRQ-save read acquisition.
356    #[track_caller]
357    pub fn try_read_irqsave(&self) -> Option<SpinRwLockIrqSaveReadGuard<'_, T>> {
358        self.read_with(CONTEXT_PREEMPT_IRQSAVE, true)
359    }
360
361    /// Acquires an IRQ-save write guard.
362    #[track_caller]
363    pub fn write_irqsave(&self) -> SpinRwLockIrqSaveWriteGuard<'_, T> {
364        self.write_with(CONTEXT_PREEMPT_IRQSAVE, false)
365            .expect("blocking IRQ-save write acquisition returned failure")
366    }
367
368    /// Attempts an IRQ-save write acquisition.
369    #[track_caller]
370    pub fn try_write_irqsave(&self) -> Option<SpinRwLockIrqSaveWriteGuard<'_, T>> {
371        self.write_with(CONTEXT_PREEMPT_IRQSAVE, true)
372    }
373
374    /// Acquires a raw read guard.
375    ///
376    /// # Safety
377    ///
378    /// The caller must prevent re-entry and uphold shared exclusion.
379    #[track_caller]
380    pub unsafe fn read_raw(&self) -> RawSpinRwLockReadGuard<'_, T> {
381        self.read_with(CONTEXT_RAW, false)
382            .expect("blocking raw read acquisition returned failure")
383    }
384
385    /// Attempts a raw read acquisition.
386    ///
387    /// # Safety
388    ///
389    /// The caller must uphold the contract of [`Self::read_raw`].
390    #[track_caller]
391    pub unsafe fn try_read_raw(&self) -> Option<RawSpinRwLockReadGuard<'_, T>> {
392        self.read_with(CONTEXT_RAW, true)
393    }
394
395    /// Acquires a raw write guard.
396    ///
397    /// # Safety
398    ///
399    /// The caller must prevent re-entry and concurrent readers or writers.
400    #[track_caller]
401    pub unsafe fn write_raw(&self) -> RawSpinRwLockWriteGuard<'_, T> {
402        self.write_with(CONTEXT_RAW, false)
403            .expect("blocking raw write acquisition returned failure")
404    }
405
406    /// Attempts a raw write acquisition.
407    ///
408    /// # Safety
409    ///
410    /// The caller must uphold the contract of [`Self::write_raw`].
411    #[track_caller]
412    pub unsafe fn try_write_raw(&self) -> Option<RawSpinRwLockWriteGuard<'_, T>> {
413        self.write_with(CONTEXT_RAW, true)
414    }
415
416    /// Returns exclusive access without locking.
417    pub fn get_mut(&mut self) -> &mut T {
418        self.data.get_mut()
419    }
420
421    /// Removes one deliberately leaked raw read guard.
422    ///
423    /// # Safety
424    ///
425    /// The caller must own one forgotten raw read guard and prove that no live
426    /// reference derived from it remains.
427    #[doc(hidden)]
428    pub unsafe fn force_read_decrement_raw(&self) {
429        crate::interface::rwlock_force_read_decrement(
430            &self.state,
431            self as *const Self as *const () as usize,
432            CONTEXT_RAW,
433        );
434    }
435}
436
437impl<T: Default> Default for SpinRwLock<T> {
438    fn default() -> Self {
439        Self::new(T::default())
440    }
441}
442
443impl<T> From<T> for SpinRwLock<T> {
444    fn from(value: T) -> Self {
445        Self::new(value)
446    }
447}
448
449impl<T: fmt::Debug> fmt::Debug for SpinRwLock<T> {
450    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
451        match self.try_read() {
452            Some(guard) => f
453                .debug_struct("SpinRwLock")
454                .field("data", &&*guard)
455                .finish(),
456            None => f
457                .debug_struct("SpinRwLock")
458                .field("data", &"<write locked>")
459                .finish(),
460        }
461    }
462}
463
464impl<T: ?Sized> Deref for SpinRwLockReadGuard<'_, T> {
465    type Target = T;
466
467    fn deref(&self) -> &Self::Target {
468        // SAFETY: the provider granted this guard shared read access.
469        unsafe { &*self.lock.data.get() }
470    }
471}
472
473impl<T: ?Sized> Drop for SpinRwLockReadGuard<'_, T> {
474    fn drop(&mut self) {
475        crate::interface::rwlock_release(
476            &self.lock.state,
477            self.lock as *const SpinRwLock<T> as *const () as usize,
478            self.context,
479            self.context_state,
480            LOCK_MODE_READ,
481        );
482    }
483}
484
485impl<T: ?Sized> Deref for SpinRwLockWriteGuard<'_, T> {
486    type Target = T;
487
488    fn deref(&self) -> &Self::Target {
489        // SAFETY: the provider granted this guard exclusive write access.
490        unsafe { &*self.lock.data.get() }
491    }
492}
493
494impl<T: ?Sized> DerefMut for SpinRwLockWriteGuard<'_, T> {
495    fn deref_mut(&mut self) -> &mut Self::Target {
496        // SAFETY: this guard uniquely represents the write acquisition.
497        unsafe { &mut *self.lock.data.get() }
498    }
499}
500
501impl<T: ?Sized> Drop for SpinRwLockWriteGuard<'_, T> {
502    fn drop(&mut self) {
503        crate::interface::rwlock_release(
504            &self.lock.state,
505            self.lock as *const SpinRwLock<T> as *const () as usize,
506            self.context,
507            self.context_state,
508            LOCK_MODE_WRITE,
509        );
510    }
511}