Jason Macnak | 09c3688 | 2020-04-01 16:22:56 +0000 | [diff] [blame] | 1 | //! A "mutex" which only supports `try_lock` |
| 2 | //! |
| 3 | //! As a futures library the eventual call to an event loop should be the only |
| 4 | //! thing that ever blocks, so this is assisted with a fast user-space |
| 5 | //! implementation of a lock that can only have a `try_lock` operation. |
| 6 | |
| 7 | use core::cell::UnsafeCell; |
| 8 | use core::ops::{Deref, DerefMut}; |
| 9 | use core::sync::atomic::Ordering::SeqCst; |
| 10 | use core::sync::atomic::AtomicBool; |
| 11 | |
| 12 | /// A "mutex" around a value, similar to `std::sync::Mutex<T>`. |
| 13 | /// |
| 14 | /// This lock only supports the `try_lock` operation, however, and does not |
| 15 | /// implement poisoning. |
| 16 | #[derive(Debug)] |
| 17 | pub(crate) struct Lock<T> { |
| 18 | locked: AtomicBool, |
| 19 | data: UnsafeCell<T>, |
| 20 | } |
| 21 | |
| 22 | /// Sentinel representing an acquired lock through which the data can be |
| 23 | /// accessed. |
| 24 | pub(crate) struct TryLock<'a, T> { |
| 25 | __ptr: &'a Lock<T>, |
| 26 | } |
| 27 | |
| 28 | // The `Lock` structure is basically just a `Mutex<T>`, and these two impls are |
| 29 | // intended to mirror the standard library's corresponding impls for `Mutex<T>`. |
| 30 | // |
| 31 | // If a `T` is sendable across threads, so is the lock, and `T` must be sendable |
| 32 | // across threads to be `Sync` because it allows mutable access from multiple |
| 33 | // threads. |
| 34 | unsafe impl<T: Send> Send for Lock<T> {} |
| 35 | unsafe impl<T: Send> Sync for Lock<T> {} |
| 36 | |
| 37 | impl<T> Lock<T> { |
| 38 | /// Creates a new lock around the given value. |
Haibo Huang | 09da603 | 2021-02-09 17:02:02 -0800 | [diff] [blame] | 39 | pub(crate) fn new(t: T) -> Self { |
| 40 | Self { |
Jason Macnak | 09c3688 | 2020-04-01 16:22:56 +0000 | [diff] [blame] | 41 | locked: AtomicBool::new(false), |
| 42 | data: UnsafeCell::new(t), |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | /// Attempts to acquire this lock, returning whether the lock was acquired or |
| 47 | /// not. |
| 48 | /// |
| 49 | /// If `Some` is returned then the data this lock protects can be accessed |
| 50 | /// through the sentinel. This sentinel allows both mutable and immutable |
| 51 | /// access. |
| 52 | /// |
| 53 | /// If `None` is returned then the lock is already locked, either elsewhere |
| 54 | /// on this thread or on another thread. |
| 55 | pub(crate) fn try_lock(&self) -> Option<TryLock<'_, T>> { |
| 56 | if !self.locked.swap(true, SeqCst) { |
| 57 | Some(TryLock { __ptr: self }) |
| 58 | } else { |
| 59 | None |
| 60 | } |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | impl<T> Deref for TryLock<'_, T> { |
| 65 | type Target = T; |
| 66 | fn deref(&self) -> &T { |
| 67 | // The existence of `TryLock` represents that we own the lock, so we |
| 68 | // can safely access the data here. |
| 69 | unsafe { &*self.__ptr.data.get() } |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | impl<T> DerefMut for TryLock<'_, T> { |
| 74 | fn deref_mut(&mut self) -> &mut T { |
| 75 | // The existence of `TryLock` represents that we own the lock, so we |
| 76 | // can safely access the data here. |
| 77 | // |
| 78 | // Additionally, we're the *only* `TryLock` in existence so mutable |
| 79 | // access should be ok. |
| 80 | unsafe { &mut *self.__ptr.data.get() } |
| 81 | } |
| 82 | } |
| 83 | |
| 84 | impl<T> Drop for TryLock<'_, T> { |
| 85 | fn drop(&mut self) { |
| 86 | self.__ptr.locked.store(false, SeqCst); |
| 87 | } |
| 88 | } |
| 89 | |
| 90 | #[cfg(test)] |
| 91 | mod tests { |
| 92 | use super::Lock; |
| 93 | |
| 94 | #[test] |
| 95 | fn smoke() { |
| 96 | let a = Lock::new(1); |
| 97 | let mut a1 = a.try_lock().unwrap(); |
| 98 | assert!(a.try_lock().is_none()); |
| 99 | assert_eq!(*a1, 1); |
| 100 | *a1 = 2; |
| 101 | drop(a1); |
| 102 | assert_eq!(*a.try_lock().unwrap(), 2); |
| 103 | assert_eq!(*a.try_lock().unwrap(), 2); |
| 104 | } |
| 105 | } |