dtact_util/sync/
notify.rs

1//! Single-slot wake-up notification, matching `tokio::sync::Notify`'s
2//! semantics: a permit that survives a `notify_one()` arriving before
3//! anyone is waiting, but `notify_waiters()` only reaches tasks already
4//! parked at the moment it's called.
5
6use super::wait_queue::WaitQueue;
7use std::future::Future;
8use std::pin::Pin;
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::task::{Context, Poll};
11
12/// A notification primitive commonly used to hand-wake another async
13/// task/structure — e.g. signaling "state changed, go re-check" without a
14/// full channel.
15#[repr(align(64))]
16pub struct Notify {
17    /// A single-slot "permit": set by `notify_one()` when nothing was
18    /// waiting, consumed by the very next `notified().await`.
19    permit: AtomicBool,
20    wait: WaitQueue,
21}
22
23impl Default for Notify {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl Notify {
30    /// Create a `Notify` with no stored permit and no waiters.
31    #[must_use]
32    pub const fn new() -> Self {
33        Self {
34            permit: AtomicBool::new(false),
35            wait: WaitQueue::new(),
36        }
37    }
38
39    /// Wake one waiting `notified().await`, or — if none is currently
40    /// waiting — store a permit so the *next* `notified().await` returns
41    /// immediately without waiting at all.
42    #[inline(always)]
43    pub fn notify_one(&self) {
44        // Unconditionally store the permit, *then* wake one queued waiter
45        // — the same "set flag, then wake the (lock-free) queue" shape
46        // `Mutex`/`RwLock`/`Semaphore` all use, rather than the mutex-
47        // queue-only version this replaced, which conditionally stored
48        // the permit solely when a locked look at the queue found it
49        // empty. That conditional version needed the queue-emptiness
50        // check and the permit store to be atomic together (true under a
51        // shared mutex, false under `WaitQueue`'s lock-free stack: a
52        // `notified()` call's `register` could land in the gap between
53        // "found empty" and "store permit", registering a waker nothing
54        // would ever wake, and permanently stalling that task while a
55        // *different*, later `notified()` caller sees the (still-set)
56        // permit and gets a wakeup that was never meant for it).
57        //
58        // This version can't lose a wakeup: any waiter's `Notified::poll`
59        // rechecks the permit *after* registering (see below), so a store
60        // that lands at any point relative to that registration is either
61        // seen by the first check, the second check, or — if the register
62        // raced ahead of both checks — by this call's `wake_one()`
63        // finding that waiter already queued. If a waiter was already
64        // parked, it gets directly woken (by `wake_one`) *and* finds the
65        // permit set when it re-polls, consuming it immediately — so no
66        // permit is ever left stranded for an unrelated later waiter to
67        // pick up, matching `tokio::sync::Notify`'s "at most one buffered
68        // permit" semantics exactly, just via a different code path.
69        self.permit.store(true, Ordering::Release);
70        self.wait.wake_one();
71    }
72
73    /// Wake every task currently waiting in `notified().await`. Does
74    /// *not* store a permit — a task that calls `notified()` after this
75    /// returns will wait for a future notification, not this one.
76    #[inline(always)]
77    pub fn notify_waiters(&self) {
78        self.wait.wake_all();
79    }
80
81    /// A future that resolves once this `Notify` is notified — either a
82    /// stored permit is consumed immediately, or a future
83    /// [`notify_one`](Self::notify_one)/[`notify_waiters`](Self::notify_waiters)
84    /// wakes it.
85    pub const fn notified(&self) -> Notified<'_> {
86        Notified { notify: self }
87    }
88}
89
90/// Future returned by [`Notify::notified`].
91#[repr(align(64))]
92pub struct Notified<'a> {
93    notify: &'a Notify,
94}
95
96impl Future for Notified<'_> {
97    type Output = ();
98
99    #[inline(always)]
100    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
101        if self
102            .notify
103            .permit
104            .compare_exchange(true, false, Ordering::Acquire, Ordering::Relaxed)
105            .is_ok()
106        {
107            return Poll::Ready(());
108        }
109        let token = self.notify.wait.register(cx.waker());
110        if self
111            .notify
112            .permit
113            .compare_exchange(true, false, Ordering::Acquire, Ordering::Relaxed)
114            .is_ok()
115        {
116            self.notify.wait.cancel(token);
117            return Poll::Ready(());
118        }
119        Poll::Pending
120    }
121}
dtact_util/sync/notify.rs

dtact_util/sync/
notify.rs
