projects/linux·blob master · refs

217 lines · 5.4 KiB

1// SPDX-License-Identifier: Apache-2.0 OR MIT
2
3#![allow(clippy::undocumented_unsafe_blocks)]
4#![cfg_attr(feature = "alloc", feature(allocator_api))]
5#![allow(clippy::missing_safety_doc)]
6
7use core::{
8 cell::{Cell, UnsafeCell},
9 marker::PhantomPinned,
10 ops::{Deref, DerefMut},
11 pin::Pin,
12 sync::atomic::{AtomicBool, Ordering},
13};
14#[cfg(feature = "std")]
15use std::{
16 sync::Arc,
17 thread::{self, sleep, Builder, Thread},
18 time::Duration,
19};
20
21use pin_init::*;
22#[allow(unused_attributes)]
23#[path = "./linked_list.rs"]
24pub mod linked_list;
25use linked_list::*;
26
27pub struct SpinLock {
28 inner: AtomicBool,
29}
30
31impl SpinLock {
32 #[inline]
33 pub fn acquire(&self) -> SpinLockGuard<'_> {
34 while self
35 .inner
36 .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
37 .is_err()
38 {
39 #[cfg(feature = "std")]
40 while self.inner.load(Ordering::Relaxed) {
41 thread::yield_now();
42 }
43 }
44 SpinLockGuard(self)
45 }
46
47 #[inline]
48 #[allow(clippy::new_without_default)]
49 pub const fn new() -> Self {
50 Self {
51 inner: AtomicBool::new(false),
52 }
53 }
54}
55
56pub struct SpinLockGuard<'a>(&'a SpinLock);
57
58impl Drop for SpinLockGuard<'_> {
59 #[inline]
60 fn drop(&mut self) {
61 self.0.inner.store(false, Ordering::Release);
62 }
63}
64
65#[pin_data]
66pub struct CMutex<T> {
67 #[pin]
68 wait_list: ListHead,
69 spin_lock: SpinLock,
70 locked: Cell<bool>,
71 #[pin]
72 data: UnsafeCell<T>,
73}
74
75impl<T> CMutex<T> {
76 #[inline]
77 pub fn new(val: impl PinInit<T>) -> impl PinInit<Self> {
78 pin_init!(CMutex {
79 wait_list <- ListHead::new(),
80 spin_lock: SpinLock::new(),
81 locked: Cell::new(false),
82 data <- UnsafeCell::pin_init(val),
83 })
84 }
85
86 #[inline]
87 pub fn lock(&self) -> Pin<CMutexGuard<'_, T>> {
88 let mut sguard = self.spin_lock.acquire();
89 if self.locked.get() {
90 stack_pin_init!(let _wait_entry = WaitEntry::insert_new(&self.wait_list));
91 // println!("wait list length: {}", self.wait_list.size());
92 while self.locked.get() {
93 drop(sguard);
94 #[cfg(feature = "std")]
95 thread::park();
96 sguard = self.spin_lock.acquire();
97 }
98 }
99 self.locked.set(true);
100 unsafe {
101 Pin::new_unchecked(CMutexGuard {
102 mtx: self,
103 _pin: PhantomPinned,
104 })
105 }
106 }
107
108 #[allow(dead_code)]
109 pub fn get_data_mut(self: Pin<&mut Self>) -> &mut T {
110 // SAFETY: we have an exclusive reference and thus nobody has access to data.
111 unsafe { &mut *self.data.get() }
112 }
113}
114
115unsafe impl<T: Send> Send for CMutex<T> {}
116unsafe impl<T: Send> Sync for CMutex<T> {}
117
118pub struct CMutexGuard<'a, T> {
119 mtx: &'a CMutex<T>,
120 _pin: PhantomPinned,
121}
122
123impl<T> Drop for CMutexGuard<'_, T> {
124 #[inline]
125 fn drop(&mut self) {
126 let sguard = self.mtx.spin_lock.acquire();
127 self.mtx.locked.set(false);
128 if let Some(list_field) = self.mtx.wait_list.next() {
129 let _wait_entry = list_field.as_ptr().cast::<WaitEntry>();
130 #[cfg(feature = "std")]
131 unsafe {
132 (*_wait_entry).thread.unpark()
133 };
134 }
135 drop(sguard);
136 }
137}
138
139impl<T> Deref for CMutexGuard<'_, T> {
140 type Target = T;
141
142 #[inline]
143 fn deref(&self) -> &Self::Target {
144 unsafe { &*self.mtx.data.get() }
145 }
146}
147
148impl<T> DerefMut for CMutexGuard<'_, T> {
149 #[inline]
150 fn deref_mut(&mut self) -> &mut Self::Target {
151 unsafe { &mut *self.mtx.data.get() }
152 }
153}
154
155#[pin_data]
156#[repr(C)]
157struct WaitEntry {
158 #[pin]
159 wait_list: ListHead,
160 #[cfg(feature = "std")]
161 thread: Thread,
162}
163
164impl WaitEntry {
165 #[inline]
166 fn insert_new(list: &ListHead) -> impl PinInit<Self> + '_ {
167 #[cfg(feature = "std")]
168 {
169 pin_init!(Self {
170 thread: thread::current(),
171 wait_list <- ListHead::insert_prev(list),
172 })
173 }
174 #[cfg(not(feature = "std"))]
175 {
176 pin_init!(Self {
177 wait_list <- ListHead::insert_prev(list),
178 })
179 }
180 }
181}
182
183#[cfg_attr(test, test)]
184#[allow(dead_code)]
185fn main() {
186 #[cfg(feature = "std")]
187 {
188 let mtx: Pin<Arc<CMutex<usize>>> = Arc::pin_init(CMutex::new(0)).unwrap();
189 let mut handles = vec![];
190 let thread_count = 20;
191 let workload = if cfg!(miri) { 100 } else { 1_000 };
192 for i in 0..thread_count {
193 let mtx = mtx.clone();
194 handles.push(
195 Builder::new()
196 .name(format!("worker #{i}"))
197 .spawn(move || {
198 for _ in 0..workload {
199 *mtx.lock() += 1;
200 }
201 println!("{i} halfway");
202 sleep(Duration::from_millis((i as u64) * 10));
203 for _ in 0..workload {
204 *mtx.lock() += 1;
205 }
206 println!("{i} finished");
207 })
208 .expect("should not fail"),
209 );
210 }
211 for h in handles {
212 h.join().expect("thread panicked");
213 }
214 println!("{:?}", *mtx.lock());
215 assert_eq!(*mtx.lock(), workload * thread_count * 2);
216 }
217}