one_shot_mutex/
mutex.rs

1use core::sync::atomic::{AtomicBool, Ordering};
2
3use lock_api::{GuardSend, RawMutex, RawMutexFair};
4
5/// A one-shot mutex that panics instead of (dead)locking on contention.
6///
7/// This mutex allows no contention and panics instead of blocking on [`lock`] if it is already locked.
8/// This is useful in situations where contention would be a bug,
9/// such as in single-threaded programs that would deadlock on contention.
10///
11/// This mutex should be used through [`OneShotMutex`].
12///
13/// [`lock`]: Self::lock
14///
15/// # Examples
16///
17/// ```
18/// use one_shot_mutex::OneShotMutex;
19///
20/// static X: OneShotMutex<i32> = OneShotMutex::new(42);
21///
22/// // This is equivalent to `X.try_lock().unwrap()`.
23/// let x = X.lock();
24///
25/// // This panics instead of deadlocking.
26/// // let x2 = X.lock();
27///
28/// // Once we unlock the mutex, we can lock it again.
29/// drop(x);
30/// let x = X.lock();
31/// ```
32pub struct RawOneShotMutex {
33    lock: AtomicBool,
34}
35
36unsafe impl RawMutex for RawOneShotMutex {
37    #[allow(clippy::declare_interior_mutable_const)]
38    const INIT: Self = Self {
39        lock: AtomicBool::new(false),
40    };
41
42    type GuardMarker = GuardSend;
43
44    #[inline]
45    fn lock(&self) {
46        assert!(
47            self.try_lock(),
48            "called `lock` on a `RawOneShotMutex` that is already locked"
49        );
50    }
51
52    #[inline]
53    fn try_lock(&self) -> bool {
54        self.lock
55            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
56            .is_ok()
57    }
58
59    #[inline]
60    unsafe fn unlock(&self) {
61        self.lock.store(false, Ordering::Release);
62    }
63
64    #[inline]
65    fn is_locked(&self) -> bool {
66        self.lock.load(Ordering::Relaxed)
67    }
68}
69
70unsafe impl RawMutexFair for RawOneShotMutex {
71    #[inline]
72    unsafe fn unlock_fair(&self) {
73        unsafe { self.unlock() }
74    }
75
76    #[inline]
77    unsafe fn bump(&self) {}
78}
79
80/// A [`lock_api::Mutex`] based on [`RawOneShotMutex`].
81pub type OneShotMutex<T> = lock_api::Mutex<RawOneShotMutex, T>;
82
83/// A [`lock_api::MutexGuard`] based on [`RawOneShotMutex`].
84pub type OneShotMutexGuard<'a, T> = lock_api::MutexGuard<'a, RawOneShotMutex, T>;
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn lock() {
92        let mutex = OneShotMutex::new(42);
93        let mut guard = mutex.lock();
94        assert_eq!(*guard, 42);
95
96        *guard += 1;
97        drop(guard);
98        let guard = mutex.lock();
99        assert_eq!(*guard, 43);
100    }
101
102    #[test]
103    #[should_panic]
104    fn lock_panic() {
105        let mutex = OneShotMutex::new(42);
106        let _guard = mutex.lock();
107        let _guard2 = mutex.lock();
108    }
109
110    #[test]
111    fn try_lock() {
112        let mutex = OneShotMutex::new(42);
113        let mut guard = mutex.try_lock().unwrap();
114        assert_eq!(*guard, 42);
115        assert!(mutex.try_lock().is_none());
116
117        *guard += 1;
118        drop(guard);
119        let guard = mutex.try_lock().unwrap();
120        assert_eq!(*guard, 43);
121    }
122}