1use alloc::{collections::VecDeque, sync::Arc};
2
3use event_listener::{Event, EventListener};
4use parking_lot::Mutex;
5
6pub(crate) struct Monitor {
7 state: Mutex<MonitorState>,
8 entry_event: Event,
9}
10
11struct MonitorState {
12 owner: Option<u64>,
13 depth: usize,
14 next_waiter_id: u64,
15 waiters: VecDeque<MonitorWaiter>,
16 active_waits: usize,
17}
18
19struct MonitorWaiter {
20 id: u64,
21 event: Arc<Event>,
22}
23
24pub struct MonitorWait {
25 monitor: Arc<Monitor>,
26 listener: EventListener,
27 depth: usize,
28 thread_id: u64,
29}
30
31#[derive(Clone)]
32pub struct MonitorWaitTimeout {
33 monitor: Arc<Monitor>,
34 waiter_id: u64,
35 event: Arc<Event>,
36}
37
38#[derive(Debug)]
39pub(crate) enum MonitorError {
40 NotOwner,
41}
42
43impl Monitor {
44 pub(crate) fn new() -> Self {
45 Self {
46 state: Mutex::new(MonitorState {
47 owner: None,
48 depth: 0,
49 next_waiter_id: 0,
50 waiters: VecDeque::new(),
51 active_waits: 0,
52 }),
53 entry_event: Event::new(),
54 }
55 }
56
57 pub(crate) fn is_active(&self) -> bool {
58 let state = self.state.lock();
59 state.owner.is_some() || state.active_waits != 0
60 }
61
62 pub(crate) async fn enter(self: &Arc<Self>, thread_id: u64) {
63 loop {
64 let listener = self.entry_event.listen();
65 {
66 let mut state = self.state.lock();
67 match state.owner {
68 None => {
69 state.owner = Some(thread_id);
70 state.depth = 1;
71 return;
72 }
73 Some(owner) if owner == thread_id => {
74 state.depth += 1;
75 return;
76 }
77 Some(_) => {}
78 }
79 }
80 listener.await;
81 }
82 }
83
84 pub(crate) fn exit(&self, thread_id: u64) -> core::result::Result<(), MonitorError> {
85 let released = {
86 let mut state = self.state.lock();
87 if state.owner != Some(thread_id) {
88 return Err(MonitorError::NotOwner);
89 }
90
91 state.depth -= 1;
92 if state.depth == 0 {
93 state.owner = None;
94 true
95 } else {
96 false
97 }
98 };
99
100 if released {
101 self.entry_event.notify(1);
102 }
103 Ok(())
104 }
105
106 pub(crate) fn prepare_wait(self: &Arc<Self>, thread_id: u64) -> core::result::Result<(MonitorWait, MonitorWaitTimeout), MonitorError> {
107 let event = Arc::new(Event::new());
108 let listener = event.listen();
109
110 let (waiter_id, depth) = {
111 let mut state = self.state.lock();
112 if state.owner != Some(thread_id) {
113 return Err(MonitorError::NotOwner);
114 }
115
116 let depth = state.depth;
117 let waiter_id = state.next_waiter_id;
118 state.next_waiter_id = state.next_waiter_id.wrapping_add(1);
119 state.waiters.push_back(MonitorWaiter {
120 id: waiter_id,
121 event: event.clone(),
122 });
123 state.active_waits += 1;
124 state.owner = None;
125 state.depth = 0;
126
127 (waiter_id, depth)
128 };
129
130 self.entry_event.notify(1);
131
132 Ok((
133 MonitorWait {
134 monitor: self.clone(),
135 listener,
136 depth,
137 thread_id,
138 },
139 MonitorWaitTimeout {
140 monitor: self.clone(),
141 waiter_id,
142 event,
143 },
144 ))
145 }
146
147 pub(crate) fn notify(&self, thread_id: u64, count: usize) -> core::result::Result<(), MonitorError> {
148 let events = {
149 let mut state = self.state.lock();
150 if state.owner != Some(thread_id) {
151 return Err(MonitorError::NotOwner);
152 }
153
154 let count = count.min(state.waiters.len());
155 (0..count)
156 .filter_map(|_| state.waiters.pop_front())
157 .map(|waiter| waiter.event)
158 .collect::<alloc::vec::Vec<_>>()
159 };
160
161 for event in events {
162 event.notify(1);
163 }
164 Ok(())
165 }
166}
167
168impl MonitorWait {
169 pub(crate) async fn wait(mut self) {
170 (&mut self.listener).await;
171 self.monitor.enter(self.thread_id).await;
172 self.monitor.state.lock().depth = self.depth;
173 }
174}
175
176impl Drop for MonitorWait {
177 fn drop(&mut self) {
178 self.monitor.state.lock().active_waits -= 1;
179 }
180}
181
182impl MonitorWaitTimeout {
183 pub fn notify(self) {
184 let event = {
185 let mut state = self.monitor.state.lock();
186 state
187 .waiters
188 .iter()
189 .position(|waiter| waiter.id == self.waiter_id)
190 .and_then(|position| state.waiters.remove(position))
191 .map(|waiter| waiter.event)
192 };
193
194 if let Some(event) = event {
195 debug_assert!(Arc::ptr_eq(&event, &self.event));
196 event.notify(1);
197 }
198 }
199}
200
201#[cfg(test)]
202mod tests {
203 use alloc::sync::Arc;
204 use core::{
205 sync::atomic::{AtomicBool, Ordering},
206 time::Duration,
207 };
208
209 use super::Monitor;
210
211 #[tokio::test]
212 async fn monitor_is_reentrant_and_excludes_other_threads() {
213 let monitor = Arc::new(Monitor::new());
214 monitor.enter(1).await;
215 monitor.enter(1).await;
216
217 let entered = Arc::new(AtomicBool::new(false));
218 let contender = {
219 let monitor = monitor.clone();
220 let entered = entered.clone();
221 tokio::spawn(async move {
222 monitor.enter(2).await;
223 entered.store(true, Ordering::SeqCst);
224 monitor.exit(2).unwrap();
225 })
226 };
227
228 tokio::time::sleep(Duration::from_millis(10)).await;
229 assert!(!entered.load(Ordering::SeqCst));
230 monitor.exit(1).unwrap();
231 tokio::time::sleep(Duration::from_millis(10)).await;
232 assert!(!entered.load(Ordering::SeqCst));
233 monitor.exit(1).unwrap();
234
235 tokio::time::timeout(Duration::from_secs(1), contender).await.unwrap().unwrap();
236 assert!(entered.load(Ordering::SeqCst));
237 }
238
239 #[tokio::test]
240 async fn wait_releases_and_restores_the_full_reentrancy_depth() {
241 let monitor = Arc::new(Monitor::new());
242 monitor.enter(1).await;
243 monitor.enter(1).await;
244 let (wait, _) = monitor.prepare_wait(1).unwrap();
245
246 monitor.enter(2).await;
247 monitor.notify(2, 1).unwrap();
248 monitor.exit(2).unwrap();
249 wait.wait().await;
250
251 monitor.exit(1).unwrap();
252 let entered = Arc::new(AtomicBool::new(false));
253 let contender = {
254 let monitor = monitor.clone();
255 let entered = entered.clone();
256 tokio::spawn(async move {
257 monitor.enter(3).await;
258 entered.store(true, Ordering::SeqCst);
259 monitor.exit(3).unwrap();
260 })
261 };
262 tokio::time::sleep(Duration::from_millis(10)).await;
263 assert!(!entered.load(Ordering::SeqCst));
264
265 monitor.exit(1).unwrap();
266 tokio::time::timeout(Duration::from_secs(1), contender).await.unwrap().unwrap();
267 assert!(entered.load(Ordering::SeqCst));
268 }
269
270 #[tokio::test]
271 async fn notify_one_and_notify_all_remove_the_expected_waiters() {
272 let monitor = Arc::new(Monitor::new());
273 monitor.enter(1).await;
274 let (first_wait, _) = monitor.prepare_wait(1).unwrap();
275 monitor.enter(2).await;
276 let (second_wait, _) = monitor.prepare_wait(2).unwrap();
277
278 monitor.enter(3).await;
279 monitor.notify(3, 1).unwrap();
280 assert_eq!(monitor.state.lock().waiters.len(), 1);
281 monitor.exit(3).unwrap();
282 first_wait.wait().await;
283 monitor.exit(1).unwrap();
284
285 monitor.enter(3).await;
286 monitor.notify(3, usize::MAX).unwrap();
287 assert!(monitor.state.lock().waiters.is_empty());
288 monitor.exit(3).unwrap();
289 second_wait.wait().await;
290 monitor.exit(2).unwrap();
291 }
292
293 #[tokio::test]
294 async fn a_stale_timeout_cannot_consume_a_later_notification() {
295 let monitor = Arc::new(Monitor::new());
296 monitor.enter(1).await;
297 let (first_wait, first_timeout) = monitor.prepare_wait(1).unwrap();
298 first_timeout.clone().notify();
299 first_wait.wait().await;
300 monitor.exit(1).unwrap();
301
302 monitor.enter(2).await;
303 let (second_wait, _) = monitor.prepare_wait(2).unwrap();
304 first_timeout.notify();
305 assert_eq!(monitor.state.lock().waiters.len(), 1);
306
307 monitor.enter(3).await;
308 monitor.notify(3, 1).unwrap();
309 monitor.exit(3).unwrap();
310 second_wait.wait().await;
311 monitor.exit(2).unwrap();
312 }
313}