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