one_shot_mutex/unsync/
rwlock.rs

Help
1use core::cell::Cell;
2
3use lock_api::{
4    GuardSend, RawRwLock, RawRwLockDowngrade, RawRwLockRecursive, RawRwLockUpgrade,
5    RawRwLockUpgradeDowngrade,
6};
7
8/// A one-shot readers-writer lock that panics instead of (dead)locking on contention.
9///
10/// This lock allows no contention and panics on [`lock_shared`], [`lock_exclusive`], [`lock_upgradable`], and [`upgrade`] if it is already locked conflictingly.
11/// This is useful in situations where contention would be a bug,
12/// such as in single-threaded programs that would deadlock on contention.
13///
14/// This lock does not implement `Sync`, which permits a slightly more efficient implementation.
15/// For a variant that does implement `Sync`, see [`sync::RawOneShotRwLock`](crate::sync::RawOneShotRwLock).
16///
17/// [`lock_shared`]: RawOneShotRwLock::lock_shared
18/// [`lock_exclusive`]: RawOneShotRwLock::lock_exclusive
19/// [`lock_upgradable`]: RawOneShotRwLock::lock_upgradable
20/// [`upgrade`]: RawOneShotRwLock::upgrade
21///
22/// # Examples
23///
24/// ```
25/// use one_shot_mutex::unsync::OneShotRwLock;
26///
27/// let m: OneShotRwLock<i32> = OneShotRwLock::new(42);
28///
29/// // This is equivalent to `X.try_write().unwrap()`.
30/// let x = m.write();
31///
32/// // This panics instead of deadlocking.
33/// // let x2 = m.write();
34///
35/// // Once we unlock the mutex, we can lock it again.
36/// drop(x);
37/// let x = m.write();
38/// ```
39pub struct RawOneShotRwLock {
40    lock: Cell<usize>,
41}
42
43/// Normal shared lock counter
44const SHARED: usize = 1 << 2;
45/// Special upgradable shared lock flag
46const UPGRADABLE: usize = 1 << 1;
47/// Exclusive lock flag
48const EXCLUSIVE: usize = 1;
49
50impl RawOneShotRwLock {
51    pub const fn new() -> Self {
52        Self::INIT
53    }
54
55    #[inline]
56    fn over_state(&self, f: impl FnOnce(usize) -> usize) -> usize {
57        let old = self.lock.get();
58        self.lock.set(f(old));
59        old
60    }
61
62    #[inline]
63    fn is_locked_shared(&self) -> bool {
64        self.lock.get() & !(EXCLUSIVE | UPGRADABLE) != 0
65    }
66
67    #[inline]
68    fn is_locked_upgradable(&self) -> bool {
69        self.lock.get() & UPGRADABLE == UPGRADABLE
70    }
71
72    /// Acquire a shared lock, returning the new lock value.
73    #[inline]
74    fn acquire_shared(&self) -> usize {
75        let value = self.over_state(|state| state + SHARED);
76
77        // An arbitrary cap that allows us to catch overflows long before they happen
78        if value > usize::MAX / 2 {
79            self.over_state(|state| state - SHARED);
80            panic!("Too many shared locks, cannot safely proceed");
81        }
82
83        value
84    }
85}
86
87impl Default for RawOneShotRwLock {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93unsafe impl RawRwLock for RawOneShotRwLock {
94    #[allow(clippy::declare_interior_mutable_const)]
95    const INIT: Self = Self { lock: Cell::new(0) };
96
97    type GuardMarker = GuardSend;
98
99    #[inline]
100    fn lock_shared(&self) {
101        assert!(
102            self.try_lock_shared(),
103            "called `lock_shared` on a `RawOneShotRwLock` that is already locked exclusively"
104        );
105    }
106
107    #[inline]
108    fn try_lock_shared(&self) -> bool {
109        let value = self.acquire_shared();
110
111        let acquired = value & EXCLUSIVE != EXCLUSIVE;
112
113        if !acquired {
114            unsafe {
115                self.unlock_shared();
116            }
117        }
118
119        acquired
120    }
121
122    #[inline]
123    unsafe fn unlock_shared(&self) {
124        debug_assert!(self.is_locked_shared());
125
126        self.over_state(|state| state - SHARED);
127    }
128
129    #[inline]
130    fn lock_exclusive(&self) {
131        assert!(
132            self.try_lock_exclusive(),
133            "called `lock_exclusive` on a `RawOneShotRwLock` that is already locked"
134        );
135    }
136
137    #[inline]
138    fn try_lock_exclusive(&self) -> bool {
139        let ok = self.lock.get() == 0;
140        if ok {
141            self.lock.set(EXCLUSIVE);
142        }
143        ok
144    }
145
146    #[inline]
147    unsafe fn unlock_exclusive(&self) {
148        debug_assert!(self.is_locked_exclusive());
149
150        self.over_state(|state| state & !EXCLUSIVE);
151    }
152
153    #[inline]
154    fn is_locked(&self) -> bool {
155        self.lock.get() != 0
156    }
157
158    #[inline]
159    fn is_locked_exclusive(&self) -> bool {
160        self.lock.get() & EXCLUSIVE == EXCLUSIVE
161    }
162}
163
164unsafe impl RawRwLockRecursive for RawOneShotRwLock {
165    #[inline]
166    fn lock_shared_recursive(&self) {
167        self.lock_shared();
168    }
169
170    #[inline]
171    fn try_lock_shared_recursive(&self) -> bool {
172        self.try_lock_shared()
173    }
174}
175
176unsafe impl RawRwLockDowngrade for RawOneShotRwLock {
177    #[inline]
178    unsafe fn downgrade(&self) {
179        // Reserve the shared guard for ourselves
180        self.acquire_shared();
181
182        unsafe {
183            self.unlock_exclusive();
184        }
185    }
186}
187
188unsafe impl RawRwLockUpgrade for RawOneShotRwLock {
189    #[inline]
190    fn lock_upgradable(&self) {
191        assert!(
192            self.try_lock_upgradable(),
193            "called `lock_upgradable` on a `RawOneShotRwLock` that is already locked upgradably or exclusively"
194        );
195    }
196
197    #[inline]
198    fn try_lock_upgradable(&self) -> bool {
199        let value = self.over_state(|state| state | UPGRADABLE);
200
201        let acquired = value & (UPGRADABLE | EXCLUSIVE) == 0;
202
203        if !acquired && value & UPGRADABLE == 0 {
204            unsafe {
205                self.unlock_upgradable();
206            }
207        }
208
209        acquired
210    }
211
212    #[inline]
213    unsafe fn unlock_upgradable(&self) {
214        debug_assert!(self.is_locked_upgradable());
215
216        self.over_state(|state| state & !UPGRADABLE);
217    }
218
219    #[inline]
220    unsafe fn upgrade(&self) {
221        assert!(
222            self.try_upgrade(),
223            "called `upgrade` on a `RawOneShotRwLock` that is also locked shared by others"
224        );
225    }
226
227    #[inline]
228    unsafe fn try_upgrade(&self) -> bool {
229        let ok = self.lock.get() == UPGRADABLE;
230        if ok {
231            self.lock.set(EXCLUSIVE);
232        }
233        ok
234    }
235}
236
237unsafe impl RawRwLockUpgradeDowngrade for RawOneShotRwLock {
238    #[inline]
239    unsafe fn downgrade_upgradable(&self) {
240        self.acquire_shared();
241
242        unsafe {
243            self.unlock_upgradable();
244        }
245    }
246
247    #[inline]
248    unsafe fn downgrade_to_upgradable(&self) {
249        debug_assert!(self.is_locked_exclusive());
250
251        self.over_state(|state| state ^ (UPGRADABLE | EXCLUSIVE));
252    }
253}
254
255/// A [`lock_api::RwLock`] based on [`RawOneShotRwLock`].
256pub type OneShotRwLock<T> = lock_api::RwLock<RawOneShotRwLock, T>;
257
258/// A [`lock_api::RwLockReadGuard`] based on [`RawOneShotRwLock`].
259pub type OneShotRwLockReadGuard<'a, T> = lock_api::RwLockReadGuard<'a, RawOneShotRwLock, T>;
260
261/// A [`lock_api::RwLockUpgradableReadGuard`] based on [`RawOneShotRwLock`].
262pub type OneShotRwLockUpgradableReadGuard<'a, T> =
263    lock_api::RwLockUpgradableReadGuard<'a, RawOneShotRwLock, T>;
264
265/// A [`lock_api::RwLockWriteGuard`] based on [`RawOneShotRwLock`].
266pub type OneShotRwLockWriteGuard<'a, T> = lock_api::RwLockWriteGuard<'a, RawOneShotRwLock, T>;
267
268#[cfg(test)]
269mod tests {
270    use lock_api::RwLockUpgradableReadGuard;
271
272    use super::*;
273
274    #[test]
275    fn lock_exclusive() {
276        let lock = OneShotRwLock::new(42);
277        let mut guard = lock.write();
278        assert_eq!(*guard, 42);
279
280        *guard += 1;
281        drop(guard);
282        let guard = lock.write();
283        assert_eq!(*guard, 43);
284    }
285
286    #[test]
287    #[should_panic]
288    fn lock_exclusive_panic() {
289        let lock = OneShotRwLock::new(42);
290        let _guard = lock.write();
291        let _guard2 = lock.write();
292    }
293
294    #[test]
295    #[should_panic]
296    fn lock_exclusive_shared_panic() {
297        let lock = OneShotRwLock::new(42);
298        let _guard = lock.write();
299        let _guard2 = lock.read();
300    }
301
302    #[test]
303    fn try_lock_exclusive() {
304        let lock = OneShotRwLock::new(42);
305        let mut guard = lock.try_write().unwrap();
306        assert_eq!(*guard, 42);
307        assert!(lock.try_write().is_none());
308
309        *guard += 1;
310        drop(guard);
311        let guard = lock.try_write().unwrap();
312        assert_eq!(*guard, 43);
313    }
314
315    #[test]
316    fn lock_shared() {
317        let lock = OneShotRwLock::new(42);
318        let guard = lock.read();
319        assert_eq!(*guard, 42);
320        let guard2 = lock.read();
321        assert_eq!(*guard2, 42);
322    }
323
324    #[test]
325    #[should_panic]
326    fn lock_shared_panic() {
327        let lock = OneShotRwLock::new(42);
328        let _guard = lock.write();
329        let _guard2 = lock.read();
330    }
331
332    #[test]
333    fn try_lock_shared() {
334        let lock = OneShotRwLock::new(42);
335        let guard = lock.try_read().unwrap();
336        assert_eq!(*guard, 42);
337        assert!(lock.try_write().is_none());
338
339        let guard2 = lock.try_read().unwrap();
340        assert_eq!(*guard2, 42);
341    }
342
343    #[test]
344    fn lock_upgradable() {
345        let lock = OneShotRwLock::new(42);
346        let guard = lock.upgradable_read();
347        assert_eq!(*guard, 42);
348        assert!(lock.try_write().is_none());
349
350        let mut upgraded = RwLockUpgradableReadGuard::upgrade(guard);
351        *upgraded += 1;
352        drop(upgraded);
353        let guard2 = lock.upgradable_read();
354        assert_eq!(*guard2, 43);
355    }
356
357    #[test]
358    #[should_panic]
359    fn lock_upgradable_panic() {
360        let lock = OneShotRwLock::new(42);
361        let _guard = lock.upgradable_read();
362        let _guard2 = lock.upgradable_read();
363    }
364
365    #[test]
366    #[should_panic]
367    fn lock_upgradable_write_panic() {
368        let lock = OneShotRwLock::new(42);
369        let _guard = lock.write();
370        let _guard2 = lock.upgradable_read();
371    }
372
373    #[test]
374    fn try_lock_upgradable() {
375        let lock = OneShotRwLock::new(42);
376        let guard = lock.try_upgradable_read().unwrap();
377        assert_eq!(*guard, 42);
378        assert!(lock.try_write().is_none());
379
380        let mut upgraded = RwLockUpgradableReadGuard::try_upgrade(guard).unwrap();
381        *upgraded += 1;
382        drop(upgraded);
383        let guard2 = lock.try_upgradable_read().unwrap();
384        assert_eq!(*guard2, 43);
385    }
386
387    #[test]
388    #[should_panic]
389    fn upgrade_panic() {
390        let lock = OneShotRwLock::new(42);
391        let guard = lock.upgradable_read();
392        let _guard2 = lock.read();
393        let _guard3 = RwLockUpgradableReadGuard::upgrade(guard);
394    }
395}
one_shot_mutex/unsync/rwlock.rs

one_shot_mutex/unsync/
rwlock.rs
