Skip to main content

skim_common/
spinlock.rs

1//! SpinLock implemented using AtomicBool
2//! Just like Mutex except:
3//!
4//! 1. It uses CAS for locking, more efficient in low contention
5//! 2. Use `.lock()` instead of `.lock().unwrap()` to retrieve the guard.
6//! 3. It doesn't handle poison so data is still available on thread panic.
7use std::cell::UnsafeCell;
8use std::ops::Deref;
9use std::ops::DerefMut;
10use std::sync::atomic::AtomicBool;
11use std::sync::atomic::Ordering;
12
13pub struct SpinLock<T: ?Sized> {
14    locked: AtomicBool,
15    data: UnsafeCell<T>,
16}
17
18unsafe impl<T: ?Sized + Send> Send for SpinLock<T> {}
19unsafe impl<T: ?Sized + Send> Sync for SpinLock<T> {}
20
21pub struct SpinLockGuard<'a, T: ?Sized + 'a> {
22    // funny underscores due to how Deref/DerefMut currently work (they
23    // disregard field privacy).
24    __lock: &'a SpinLock<T>,
25}
26
27impl<'a, T: ?Sized + 'a> SpinLockGuard<'a, T> {
28    pub fn new(pool: &'a SpinLock<T>) -> SpinLockGuard<'a, T> {
29        Self { __lock: pool }
30    }
31}
32
33unsafe impl<T: ?Sized + Sync> Sync for SpinLockGuard<'_, T> {}
34
35impl<T> SpinLock<T> {
36    pub fn new(t: T) -> SpinLock<T> {
37        Self {
38            locked: AtomicBool::new(false),
39            data: UnsafeCell::new(t),
40        }
41    }
42}
43
44impl<T: ?Sized> SpinLock<T> {
45    pub fn lock(&self) -> SpinLockGuard<T> {
46        while self
47            .locked
48            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
49            .is_err()
50        {}
51        SpinLockGuard::new(self)
52    }
53}
54
55impl<T: ?Sized> Deref for SpinLockGuard<'_, T> {
56    type Target = T;
57
58    fn deref(&self) -> &T {
59        unsafe { &*self.__lock.data.get() }
60    }
61}
62
63impl<T: ?Sized> DerefMut for SpinLockGuard<'_, T> {
64    fn deref_mut(&mut self) -> &mut T {
65        unsafe { &mut *self.__lock.data.get() }
66    }
67}
68
69impl<T: ?Sized> Drop for SpinLockGuard<'_, T> {
70    #[inline]
71    fn drop(&mut self) {
72        while self
73            .__lock
74            .locked
75            .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst)
76            .is_err()
77        {}
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use std::sync::Arc;
85    use std::sync::mpsc::channel;
86    use std::thread;
87
88    #[test]
89    fn smoke() {
90        let m = SpinLock::new(());
91        drop(m.lock());
92        drop(m.lock());
93    }
94
95    #[test]
96    fn lots_and_lots() {
97        const J: u32 = 1000;
98        const K: u32 = 3;
99
100        let m = Arc::new(SpinLock::new(0));
101
102        fn inc(m: &SpinLock<u32>) {
103            for _ in 0..J {
104                *m.lock() += 1;
105            }
106        }
107
108        let (tx, rx) = channel();
109        for _ in 0..K {
110            let tx2 = tx.clone();
111            let m2 = m.clone();
112            thread::spawn(move || {
113                inc(&m2);
114                tx2.send(()).unwrap();
115            });
116            let tx2 = tx.clone();
117            let m2 = m.clone();
118            thread::spawn(move || {
119                inc(&m2);
120                tx2.send(()).unwrap();
121            });
122        }
123
124        drop(tx);
125        for _ in 0..2 * K {
126            rx.recv().unwrap();
127        }
128        assert_eq!(*m.lock(), J * K * 2);
129    }
130
131    #[test]
132    fn test_mutex_unsized() {
133        let mutex: &SpinLock<[i32]> = &SpinLock::new([1, 2, 3]);
134        {
135            let b = &mut *mutex.lock();
136            b[0] = 4;
137            b[2] = 5;
138        }
139        let comp: &[i32] = &[4, 2, 5];
140        assert_eq!(&*mutex.lock(), comp);
141    }
142}