Skip to main content

moirai_async/sync/
condvar.rs

1#![expect(
2    clippy::unwrap_used,
3    reason = "ratchet MOIRAI-UNWRAP-1: pre-existing debt"
4)]
5
6use std::future::Future;
7use std::pin::Pin;
8use std::task::{Context, Poll, Waker};
9
10use crate::sync::wait_queue::{WaitQueue, WaiterPoll};
11
12use super::mutex::{Mutex, MutexGuard};
13
14fn noop_waker() -> Waker {
15    Waker::noop().clone()
16}
17
18/// Async condition variable pairing with the async [`Mutex`].
19pub struct Condvar {
20    state: std::sync::Mutex<CondvarState>,
21}
22
23struct CondvarState {
24    waiters: WaitQueue<()>,
25}
26
27impl Condvar {
28    /// Create a condition variable with no waiters.
29    pub fn new() -> Self {
30        Self {
31            state: std::sync::Mutex::new(CondvarState {
32                waiters: WaitQueue::new(),
33            }),
34        }
35    }
36
37    /// Release the guard, wait for a notification, and reacquire.
38    pub async fn wait<'a, T>(&self, guard: MutexGuard<'a, T>) -> MutexGuard<'a, T> {
39        let mutex_ref: &'a Mutex<T> = guard.mutex;
40        // Register a pending waiter WHILE still holding the outer MutexGuard.
41        // This closes the lost-notification window: a concurrent notify_one/all
42        // after the guard is dropped will see this waiter (or will have already
43        // seen it and set Granted, which poll_waiter returns immediately).
44        let id = {
45            let mut state = self.state.lock().unwrap();
46            state.waiters.register(noop_waker())
47        };
48        drop(guard);
49        CondvarNotifyFuture {
50            condvar: self,
51            id: Some(id),
52        }
53        .await;
54        mutex_ref.lock().await
55    }
56
57    /// Wait until `condition` on the guarded value turns false.
58    pub async fn wait_while<'a, T, F>(
59        &self,
60        guard: MutexGuard<'a, T>,
61        mut condition: F,
62    ) -> MutexGuard<'a, T>
63    where
64        F: FnMut(&T) -> bool,
65    {
66        let mut guard = guard;
67        while condition(&guard) {
68            guard = self.wait(guard).await;
69        }
70        guard
71    }
72
73    /// Wake the oldest waiter, if any.
74    ///
75    /// The waker is taken under the state lock and woken after it is released:
76    /// `Waker::wake` may poll the task inline on this thread, and that poll
77    /// re-locks this state — waking under the lock would self-deadlock. Same
78    /// discipline as `notify_all` below and `hybrid::notify`.
79    pub fn notify_one(&self) {
80        let waker = {
81            let mut state = self.state.lock().unwrap();
82            state.waiters.grant_oldest(())
83        };
84        if let Some(waker) = waker {
85            waker.wake();
86        }
87    }
88
89    /// Wake every current waiter.
90    pub fn notify_all(&self) {
91        let mut state = self.state.lock().unwrap();
92        let wakers = state.waiters.grant_all(());
93        drop(state);
94        for waker in wakers {
95            waker.wake();
96        }
97    }
98}
99
100impl Default for Condvar {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106struct CondvarNotifyFuture<'a> {
107    condvar: &'a Condvar,
108    id: Option<u64>,
109}
110
111impl<'a> Future for CondvarNotifyFuture<'a> {
112    type Output = ();
113
114    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
115        let mut state = self.condvar.state.lock().unwrap();
116
117        if let Some(id) = self.id {
118            match state.waiters.poll_waiter(id, cx.waker()) {
119                WaiterPoll::Granted(()) => {
120                    self.id = None;
121                    return Poll::Ready(());
122                }
123                WaiterPoll::Pending => return Poll::Pending,
124                WaiterPoll::NotRegistered => {}
125            }
126        }
127
128        if self.id.is_none() {
129            self.id = Some(state.waiters.register(cx.waker().clone()));
130        }
131
132        Poll::Pending
133    }
134}
135
136impl<'a> Drop for CondvarNotifyFuture<'a> {
137    fn drop(&mut self) {
138        if let Some(id) = self.id
139            && let Ok(mut state) = self.condvar.state.lock()
140        {
141            state.waiters.deregister(id);
142        }
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    #[test]
151    fn test_condvar_notify_one() {
152        let cv = Condvar::new();
153        cv.notify_one();
154    }
155
156    #[test]
157    fn test_condvar_notify_all() {
158        let cv = Condvar::new();
159        cv.notify_all();
160    }
161}