one_shot_mutex/sync/
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::sync::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
36impl RawOneShotMutex {
37    pub const fn new() -> Self {
38        Self::INIT
39    }
40}
41
42impl Default for RawOneShotMutex {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48unsafe impl RawMutex for RawOneShotMutex {
49    #[allow(clippy::declare_interior_mutable_const)]
50    const INIT: Self = Self {
51        lock: AtomicBool::new(false),
52    };
53
54    type GuardMarker = GuardSend;
55
56    #[inline]
57    fn lock(&self) {
58        assert!(
59            self.try_lock(),
60            "called `lock` on a `RawOneShotMutex` that is already locked"
61        );
62    }
63
64    #[inline]
65    fn try_lock(&self) -> bool {
66        self.lock
67            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
68            .is_ok()
69    }
70
71    #[inline]
72    unsafe fn unlock(&self) {
73        self.lock.store(false, Ordering::Release);
74    }
75
76    #[inline]
77    fn is_locked(&self) -> bool {
78        self.lock.load(Ordering::Relaxed)
79    }
80}
81
82unsafe impl RawMutexFair for RawOneShotMutex {
83    #[inline]
84    unsafe fn unlock_fair(&self) {
85        unsafe { self.unlock() }
86    }
87
88    #[inline]
89    unsafe fn bump(&self) {}
90}
91
92/// A [`lock_api::Mutex`] based on [`RawOneShotMutex`].
93pub type OneShotMutex<T> = lock_api::Mutex<RawOneShotMutex, T>;
94
95/// A [`lock_api::MutexGuard`] based on [`RawOneShotMutex`].
96pub type OneShotMutexGuard<'a, T> = lock_api::MutexGuard<'a, RawOneShotMutex, T>;
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn lock() {
104        let mutex = OneShotMutex::new(42);
105        let mut guard = mutex.lock();
106        assert_eq!(*guard, 42);
107
108        *guard += 1;
109        drop(guard);
110        let guard = mutex.lock();
111        assert_eq!(*guard, 43);
112    }
113
114    #[test]
115    #[should_panic]
116    fn lock_panic() {
117        let mutex = OneShotMutex::new(42);
118        let _guard = mutex.lock();
119        let _guard2 = mutex.lock();
120    }
121
122    #[test]
123    fn try_lock() {
124        let mutex = OneShotMutex::new(42);
125        let mut guard = mutex.try_lock().unwrap();
126        assert_eq!(*guard, 42);
127        assert!(mutex.try_lock().is_none());
128
129        *guard += 1;
130        drop(guard);
131        let guard = mutex.try_lock().unwrap();
132        assert_eq!(*guard, 43);
133    }
134}