1use 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 __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}