tokio/util/
try_lock.rs

1use crate::loom::sync::atomic::AtomicBool;
2
3use std::cell::UnsafeCell;
4use std::marker::PhantomData;
5use std::ops::{Deref, DerefMut};
6use std::sync::atomic::Ordering::SeqCst;
7
8pub(crate) struct TryLock<T> {
9    locked: AtomicBool,
10    data: UnsafeCell<T>,
11}
12
13pub(crate) struct LockGuard<'a, T> {
14    lock: &'a TryLock<T>,
15    _p: PhantomData<std::rc::Rc<()>>,
16}
17
18unsafe impl<T: Send> Send for TryLock<T> {}
19unsafe impl<T: Send> Sync for TryLock<T> {}
20
21unsafe impl<T: Sync> Sync for LockGuard<'_, T> {}
22
23macro_rules! new {
24    ($data:ident) => {
25        TryLock {
26            locked: AtomicBool::new(false),
27            data: UnsafeCell::new($data),
28        }
29    };
30}
31
32impl<T> TryLock<T> {
33    #[cfg(not(loom))]
34    /// Create a new `TryLock`
35    pub(crate) const fn new(data: T) -> TryLock<T> {
36        new!(data)
37    }
38
39    #[cfg(loom)]
40    /// Create a new `TryLock`
41    pub(crate) fn new(data: T) -> TryLock<T> {
42        new!(data)
43    }
44
45    /// Attempt to acquire lock
46    pub(crate) fn try_lock(&self) -> Option<LockGuard<'_, T>> {
47        if self
48            .locked
49            .compare_exchange(false, true, SeqCst, SeqCst)
50            .is_err()
51        {
52            return None;
53        }
54
55        Some(LockGuard {
56            lock: self,
57            _p: PhantomData,
58        })
59    }
60}
61
62impl<T> Deref for LockGuard<'_, T> {
63    type Target = T;
64
65    fn deref(&self) -> &T {
66        unsafe { &*self.lock.data.get() }
67    }
68}
69
70impl<T> DerefMut for LockGuard<'_, T> {
71    fn deref_mut(&mut self) -> &mut T {
72        unsafe { &mut *self.lock.data.get() }
73    }
74}
75
76impl<T> Drop for LockGuard<'_, T> {
77    fn drop(&mut self) {
78        self.lock.locked.store(false, SeqCst);
79    }
80}
tokio/util/try_lock.rs

tokio/util/
try_lock.rs
