Skip to main content

tokio/sync/
notify.rs

1// Allow `unreachable_pub` warnings when sync is not enabled
2// due to the usage of `Notify` within the `rt` feature set.
3// When this module is compiled with `sync` enabled we will warn on
4// this lint. When `rt` is enabled we use `pub(crate)` which
5// triggers this warning but it is safe to ignore in this case.
6#![cfg_attr(not(feature = "sync"), allow(unreachable_pub, dead_code))]
7
8use crate::loom::cell::UnsafeCell;
9use crate::loom::sync::atomic::AtomicUsize;
10use crate::loom::sync::Mutex;
11use crate::util::linked_list::{self, GuardedLinkedList, LinkedList};
12use crate::util::WakeList;
13
14use std::future::Future;
15use std::marker::PhantomPinned;
16use std::panic::{RefUnwindSafe, UnwindSafe};
17use std::pin::Pin;
18use std::ptr::NonNull;
19use std::sync::atomic::Ordering::{self, Acquire, Relaxed, Release, SeqCst};
20use std::sync::Arc;
21use std::task::{Context, Poll, Waker};
22
23/// Notifies a single task to wake up.
24///
25/// `Notify` provides a basic mechanism to notify a single task of an event.
26/// `Notify` itself does not carry any data. Instead, it is to be used to signal
27/// another task to perform an operation.
28///
29/// A `Notify` can be thought of as a [`Semaphore`] starting with 0 permits. The
30/// [`notified().await`] method waits for a permit to become available, and
31/// [`notify_one()`] sets a permit **if there currently are no available
32/// permits**.
33///
34/// The synchronization details of `Notify` are similar to
35/// [`thread::park`][park] and [`Thread::unpark`][unpark] from std. A [`Notify`]
36/// value contains a single permit. [`notified().await`] waits for the permit to
37/// be made available, consumes the permit, and resumes.  [`notify_one()`] sets
38/// the permit, waking a pending task if there is one.
39///
40/// If `notify_one()` is called **before** `notified().await`, then the next
41/// call to `notified().await` will complete immediately, consuming the permit.
42/// Any subsequent calls to `notified().await` will wait for a new permit.
43///
44/// If `notify_one()` is called **multiple** times before `notified().await`,
45/// only a **single** permit is stored. The next call to `notified().await` will
46/// complete immediately, but the one after will wait for a new permit.
47///
48/// # Examples
49///
50/// Basic usage.
51///
52/// ```
53/// use tokio::sync::Notify;
54/// use std::sync::Arc;
55///
56/// # #[tokio::main(flavor = "current_thread")]
57/// # async fn main() {
58/// let notify = Arc::new(Notify::new());
59/// let notify2 = notify.clone();
60///
61/// let handle = tokio::spawn(async move {
62///     notify2.notified().await;
63///     println!("received notification");
64/// });
65///
66/// println!("sending notification");
67/// notify.notify_one();
68///
69/// // Wait for task to receive notification.
70/// handle.await.unwrap();
71/// # }
72/// ```
73///
74/// Unbound multi-producer single-consumer (mpsc) channel.
75///
76/// No wakeups can be lost when using this channel because the call to
77/// `notify_one()` will store a permit in the `Notify`, which the following call
78/// to `notified()` will consume.
79///
80/// ```
81/// use tokio::sync::Notify;
82///
83/// use std::collections::VecDeque;
84/// use std::sync::Mutex;
85///
86/// struct Channel<T> {
87///     values: Mutex<VecDeque<T>>,
88///     notify: Notify,
89/// }
90///
91/// impl<T> Channel<T> {
92///     pub fn send(&self, value: T) {
93///         self.values.lock().unwrap()
94///             .push_back(value);
95///
96///         // Notify the consumer a value is available
97///         self.notify.notify_one();
98///     }
99///
100///     // This is a single-consumer channel, so several concurrent calls to
101///     // `recv` are not allowed.
102///     pub async fn recv(&self) -> T {
103///         loop {
104///             // Drain values
105///             if let Some(value) = self.values.lock().unwrap().pop_front() {
106///                 return value;
107///             }
108///
109///             // Wait for values to be available
110///             self.notify.notified().await;
111///         }
112///     }
113/// }
114/// ```
115///
116/// Unbound multi-producer multi-consumer (mpmc) channel.
117///
118/// The call to [`enable`] is important because otherwise if you have two
119/// calls to `recv` and two calls to `send` in parallel, the following could
120/// happen:
121///
122///  1. Both calls to `try_recv` return `None`.
123///  2. Both new elements are added to the vector.
124///  3. The `notify_one` method is called twice, adding only a single
125///     permit to the `Notify`.
126///  4. Both calls to `recv` reach the `Notified` future. One of them
127///     consumes the permit, and the other sleeps forever.
128///
129/// By adding the `Notified` futures to the list by calling `enable` before
130/// `try_recv`, the `notify_one` calls in step three would remove the
131/// futures from the list and mark them notified instead of adding a permit
132/// to the `Notify`. This ensures that both futures are woken.
133///
134/// Notice that this failure can only happen if there are two concurrent calls
135/// to `recv`. This is why the mpsc example above does not require a call to
136/// `enable`.
137///
138/// ```
139/// use tokio::sync::Notify;
140///
141/// use std::collections::VecDeque;
142/// use std::sync::Mutex;
143///
144/// struct Channel<T> {
145///     messages: Mutex<VecDeque<T>>,
146///     notify_on_sent: Notify,
147/// }
148///
149/// impl<T> Channel<T> {
150///     pub fn send(&self, msg: T) {
151///         let mut locked_queue = self.messages.lock().unwrap();
152///         locked_queue.push_back(msg);
153///         drop(locked_queue);
154///
155///         // Send a notification to one of the calls currently
156///         // waiting in a call to `recv`.
157///         self.notify_on_sent.notify_one();
158///     }
159///
160///     pub fn try_recv(&self) -> Option<T> {
161///         let mut locked_queue = self.messages.lock().unwrap();
162///         locked_queue.pop_front()
163///     }
164///
165///     pub async fn recv(&self) -> T {
166///         let future = self.notify_on_sent.notified();
167///         tokio::pin!(future);
168///
169///         loop {
170///             // Make sure that no wakeup is lost if we get
171///             // `None` from `try_recv`.
172///             future.as_mut().enable();
173///
174///             if let Some(msg) = self.try_recv() {
175///                 return msg;
176///             }
177///
178///             // Wait for a call to `notify_one`.
179///             //
180///             // This uses `.as_mut()` to avoid consuming the future,
181///             // which lets us call `Pin::set` below.
182///             future.as_mut().await;
183///
184///             // Reset the future in case another call to
185///             // `try_recv` got the message before us.
186///             future.set(self.notify_on_sent.notified());
187///         }
188///     }
189/// }
190/// ```
191///
192/// [park]: std::thread::park
193/// [unpark]: std::thread::Thread::unpark
194/// [`notified().await`]: Notify::notified()
195/// [`notify_one()`]: Notify::notify_one()
196/// [`enable`]: Notified::enable()
197/// [`Semaphore`]: crate::sync::Semaphore
198#[derive(Debug)]
199pub struct Notify {
200    // `state` uses 2 bits to store one of `EMPTY`,
201    // `WAITING` or `NOTIFIED`. The rest of the bits
202    // are used to store the number of times `notify_waiters`
203    // was called.
204    //
205    // Throughout the code there are two assumptions:
206    // - state can be transitioned *from* `WAITING` only if
207    //   `waiters` lock is held
208    // - number of times `notify_waiters` was called can
209    //   be modified only if `waiters` lock is held
210    state: AtomicUsize,
211    waiters: Mutex<LinkedList<Waiter>>,
212}
213
214#[derive(Debug)]
215struct Waiter {
216    /// Intrusive linked-list pointers.
217    pointers: linked_list::Pointers<Waiter>,
218
219    /// Waiting task's waker. Depending on the value of `notification`,
220    /// this field is either protected by the `waiters` lock in
221    /// `Notify`, or it is exclusively owned by the enclosing `Waiter`.
222    waker: UnsafeCell<Option<Waker>>,
223
224    /// Notification for this waiter. Uses 2 bits to store if and how was
225    /// notified, 1 bit for storing if it was woken up using FIFO or LIFO, and
226    /// the rest of it is unused.
227    /// * if it's `None`, then `waker` is protected by the `waiters` lock.
228    /// * if it's `Some`, then `waker` is exclusively owned by the
229    ///   enclosing `Waiter` and can be accessed without locking.
230    notification: AtomicNotification,
231
232    /// Should not be `Unpin`.
233    _p: PhantomPinned,
234}
235
236impl Waiter {
237    fn new() -> Waiter {
238        Waiter {
239            pointers: linked_list::Pointers::new(),
240            waker: UnsafeCell::new(None),
241            notification: AtomicNotification::none(),
242            _p: PhantomPinned,
243        }
244    }
245}
246
247generate_addr_of_methods! {
248    impl<> Waiter {
249        unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> {
250            &self.pointers
251        }
252    }
253}
254
255// No notification.
256const NOTIFICATION_NONE: usize = 0b000;
257
258// Notification type used by `notify_one`.
259const NOTIFICATION_ONE: usize = 0b001;
260
261// Notification type used by `notify_last`.
262const NOTIFICATION_LAST: usize = 0b101;
263
264// Notification type used by `notify_waiters`.
265const NOTIFICATION_ALL: usize = 0b010;
266
267/// Notification for a `Waiter`.
268/// This struct is equivalent to `Option<Notification>`, but uses
269/// `AtomicUsize` inside for atomic operations.
270#[derive(Debug)]
271struct AtomicNotification(AtomicUsize);
272
273impl AtomicNotification {
274    fn none() -> Self {
275        AtomicNotification(AtomicUsize::new(NOTIFICATION_NONE))
276    }
277
278    /// Store-release a notification.
279    /// This method should be called exactly once.
280    fn store_release(&self, notification: Notification) {
281        let data: usize = match notification {
282            Notification::All => NOTIFICATION_ALL,
283            Notification::One(NotifyOneStrategy::Fifo) => NOTIFICATION_ONE,
284            Notification::One(NotifyOneStrategy::Lifo) => NOTIFICATION_LAST,
285        };
286        self.0.store(data, Release);
287    }
288
289    fn load(&self, ordering: Ordering) -> Option<Notification> {
290        let data = self.0.load(ordering);
291        match data {
292            NOTIFICATION_NONE => None,
293            NOTIFICATION_ONE => Some(Notification::One(NotifyOneStrategy::Fifo)),
294            NOTIFICATION_LAST => Some(Notification::One(NotifyOneStrategy::Lifo)),
295            NOTIFICATION_ALL => Some(Notification::All),
296            _ => unreachable!(),
297        }
298    }
299
300    /// Clears the notification.
301    /// This method is used by a `Notified` future to consume the
302    /// notification. It uses relaxed ordering and should be only
303    /// used once the atomic notification is no longer shared.
304    fn clear(&self) {
305        self.0.store(NOTIFICATION_NONE, Relaxed);
306    }
307}
308
309#[derive(Debug, PartialEq, Eq)]
310#[repr(usize)]
311enum NotifyOneStrategy {
312    Fifo,
313    Lifo,
314}
315
316#[derive(Debug, PartialEq, Eq)]
317#[repr(usize)]
318enum Notification {
319    One(NotifyOneStrategy),
320    All,
321}
322
323/// List used in `Notify::notify_waiters`. It wraps a guarded linked list
324/// and gates the access to it on `notify.waiters` mutex. It also empties
325/// the list on drop.
326struct NotifyWaitersList<'a> {
327    list: GuardedLinkedList<Waiter>,
328    is_empty: bool,
329    notify: &'a Notify,
330}
331
332impl<'a> NotifyWaitersList<'a> {
333    fn new(
334        unguarded_list: LinkedList<Waiter>,
335        guard: Pin<&'a Waiter>,
336        notify: &'a Notify,
337    ) -> NotifyWaitersList<'a> {
338        let guard_ptr = NonNull::from(guard.get_ref());
339        let list = unguarded_list.into_guarded(guard_ptr);
340        NotifyWaitersList {
341            list,
342            is_empty: false,
343            notify,
344        }
345    }
346
347    /// Removes the last element from the guarded list. Modifying this list
348    /// requires an exclusive access to the main list in `Notify`.
349    fn pop_back_locked(&mut self, _waiters: &mut LinkedList<Waiter>) -> Option<NonNull<Waiter>> {
350        let result = self.list.pop_back();
351        if result.is_none() {
352            // Save information about emptiness to avoid waiting for lock
353            // in the destructor.
354            self.is_empty = true;
355        }
356        result
357    }
358}
359
360impl Drop for NotifyWaitersList<'_> {
361    fn drop(&mut self) {
362        // If the list is not empty, we unlink all waiters from it.
363        // We do not wake the waiters to avoid double panics.
364        if !self.is_empty {
365            let _lock_guard = self.notify.waiters.lock();
366            while let Some(waiter) = self.list.pop_back() {
367                // Safety: we never make mutable references to waiters.
368                let waiter = unsafe { waiter.as_ref() };
369                waiter.notification.store_release(Notification::All);
370            }
371        }
372    }
373}
374
375/// Future returned from [`Notify::notified()`].
376///
377/// This future is fused, so once it has completed, any future calls to poll
378/// will immediately return `Poll::Ready`.
379#[derive(Debug)]
380#[must_use = "futures do nothing unless you `.await` or poll them"]
381pub struct Notified<'a> {
382    /// The `Notify` being received on.
383    notify: &'a Notify,
384
385    /// The current state of the receiving process.
386    state: State,
387
388    /// Number of calls to `notify_waiters` at the time of creation.
389    notify_waiters_calls: usize,
390
391    /// Entry in the waiter `LinkedList`.
392    waiter: Waiter,
393}
394
395unsafe impl<'a> Send for Notified<'a> {}
396unsafe impl<'a> Sync for Notified<'a> {}
397
398/// Future returned from [`Notify::notified_owned()`].
399///
400/// This future is fused, so once it has completed, any future calls to poll
401/// will immediately return `Poll::Ready`.
402#[derive(Debug)]
403#[must_use = "futures do nothing unless you `.await` or poll them"]
404pub struct OwnedNotified {
405    /// The `Notify` being received on.
406    notify: Arc<Notify>,
407
408    /// The current state of the receiving process.
409    state: State,
410
411    /// Number of calls to `notify_waiters` at the time of creation.
412    notify_waiters_calls: usize,
413
414    /// Entry in the waiter `LinkedList`.
415    waiter: Waiter,
416}
417
418unsafe impl Sync for OwnedNotified {}
419
420/// A custom `project` implementation is used in place of `pin-project-lite`
421/// as a custom drop for [`Notified`] and [`OwnedNotified`] implementation
422/// is needed.
423struct NotifiedProject<'a> {
424    notify: &'a Notify,
425    state: &'a mut State,
426    notify_waiters_calls: &'a usize,
427    waiter: &'a Waiter,
428}
429
430#[derive(Debug)]
431enum State {
432    Init,
433    Waiting,
434    Done,
435}
436
437const NOTIFY_WAITERS_SHIFT: usize = 2;
438const STATE_MASK: usize = (1 << NOTIFY_WAITERS_SHIFT) - 1;
439const NOTIFY_WAITERS_CALLS_MASK: usize = !STATE_MASK;
440
441/// Initial "idle" state.
442const EMPTY: usize = 0;
443
444/// One or more threads are currently waiting to be notified.
445const WAITING: usize = 1;
446
447/// Pending notification.
448const NOTIFIED: usize = 2;
449
450fn set_state(data: usize, state: usize) -> usize {
451    (data & NOTIFY_WAITERS_CALLS_MASK) | (state & STATE_MASK)
452}
453
454fn get_state(data: usize) -> usize {
455    data & STATE_MASK
456}
457
458fn get_num_notify_waiters_calls(data: usize) -> usize {
459    (data & NOTIFY_WAITERS_CALLS_MASK) >> NOTIFY_WAITERS_SHIFT
460}
461
462fn inc_num_notify_waiters_calls(data: usize) -> usize {
463    data + (1 << NOTIFY_WAITERS_SHIFT)
464}
465
466fn atomic_inc_num_notify_waiters_calls(data: &AtomicUsize) {
467    data.fetch_add(1 << NOTIFY_WAITERS_SHIFT, SeqCst);
468}
469
470impl Notify {
471    /// Create a new `Notify`, initialized without a permit.
472    ///
473    /// # Examples
474    ///
475    /// ```
476    /// use tokio::sync::Notify;
477    ///
478    /// let notify = Notify::new();
479    /// ```
480    pub fn new() -> Notify {
481        Notify {
482            state: AtomicUsize::new(0),
483            waiters: Mutex::new(LinkedList::new()),
484        }
485    }
486
487    /// Create a new `Notify`, initialized without a permit.
488    ///
489    /// When using the `tracing` [unstable feature], a `Notify` created with
490    /// `const_new` will not be instrumented. As such, it will not be visible
491    /// in [`tokio-console`]. Instead, [`Notify::new`] should be used to create
492    /// an instrumented object if that is needed.
493    ///
494    /// # Examples
495    ///
496    /// ```
497    /// use tokio::sync::Notify;
498    ///
499    /// static NOTIFY: Notify = Notify::const_new();
500    /// ```
501    ///
502    /// [`tokio-console`]: https://github.com/tokio-rs/console
503    /// [unstable feature]: crate#unstable-features
504    #[cfg(not(all(loom, test)))]
505    pub const fn const_new() -> Notify {
506        Notify {
507            state: AtomicUsize::new(0),
508            waiters: Mutex::const_new(LinkedList::new()),
509        }
510    }
511
512    /// Wait for a notification.
513    ///
514    /// Equivalent to:
515    ///
516    /// ```ignore
517    /// async fn notified(&self);
518    /// ```
519    ///
520    /// Each `Notify` value holds a single permit. If a permit is available from
521    /// an earlier call to [`notify_one()`], then `notified().await` will complete
522    /// immediately, consuming that permit. Otherwise, `notified().await` waits
523    /// for a permit to be made available by the next call to `notify_one()`.
524    ///
525    /// The `Notified` future is not guaranteed to receive wakeups from calls to
526    /// `notify_one()` if it has not yet been polled. See the documentation for
527    /// [`Notified::enable()`] for more details.
528    ///
529    /// The `Notified` future is guaranteed to receive wakeups from
530    /// `notify_waiters()` as soon as it has been created, even if it has not
531    /// yet been polled.
532    ///
533    /// [`notify_one()`]: Notify::notify_one
534    /// [`Notified::enable()`]: Notified::enable
535    ///
536    /// # Cancel safety
537    ///
538    /// This method uses a queue to fairly distribute notifications in the order
539    /// they were requested. Cancelling a call to `notified` makes you lose your
540    /// place in the queue.
541    ///
542    /// # Examples
543    ///
544    /// ```
545    /// use tokio::sync::Notify;
546    /// use std::sync::Arc;
547    ///
548    /// # #[tokio::main(flavor = "current_thread")]
549    /// # async fn main() {
550    /// let notify = Arc::new(Notify::new());
551    /// let notify2 = notify.clone();
552    ///
553    /// tokio::spawn(async move {
554    ///     notify2.notified().await;
555    ///     println!("received notification");
556    /// });
557    ///
558    /// println!("sending notification");
559    /// notify.notify_one();
560    /// # }
561    /// ```
562    pub fn notified(&self) -> Notified<'_> {
563        // we load the number of times notify_waiters
564        // was called and store that in the future.
565        let state = self.state.load(SeqCst);
566        Notified {
567            notify: self,
568            state: State::Init,
569            notify_waiters_calls: get_num_notify_waiters_calls(state),
570            waiter: Waiter::new(),
571        }
572    }
573
574    /// Wait for a notification with an owned `Future`.
575    ///
576    /// Unlike [`Self::notified`] which returns a future tied to the `Notify`'s
577    /// lifetime, `notified_owned` creates a self-contained future that owns its
578    /// notification state, making it safe to move between threads.
579    ///
580    /// See [`Self::notified`] for more details.
581    ///
582    /// # Cancel safety
583    ///
584    /// This method uses a queue to fairly distribute notifications in the order
585    /// they were requested. Cancelling a call to `notified_owned` makes you lose your
586    /// place in the queue.
587    ///
588    /// # Examples
589    ///
590    /// ```
591    /// use std::sync::Arc;
592    /// use tokio::sync::Notify;
593    ///
594    /// # #[tokio::main(flavor = "current_thread")]
595    /// # async fn main() {
596    /// let notify = Arc::new(Notify::new());
597    ///
598    /// for _ in 0..10 {
599    ///     let notified = notify.clone().notified_owned();
600    ///     tokio::spawn(async move {
601    ///         notified.await;
602    ///         println!("received notification");
603    ///     });
604    /// }
605    ///
606    /// println!("sending notification");
607    /// notify.notify_waiters();
608    /// # }
609    /// ```
610    pub fn notified_owned(self: Arc<Self>) -> OwnedNotified {
611        // we load the number of times notify_waiters
612        // was called and store that in the future.
613        let state = self.state.load(SeqCst);
614        OwnedNotified {
615            notify: self,
616            state: State::Init,
617            notify_waiters_calls: get_num_notify_waiters_calls(state),
618            waiter: Waiter::new(),
619        }
620    }
621    /// Notifies the first waiting task.
622    ///
623    /// If a task is currently waiting, that task is notified. Otherwise, a
624    /// permit is stored in this `Notify` value and the **next** call to
625    /// [`notified().await`] will complete immediately consuming the permit made
626    /// available by this call to `notify_one()`.
627    ///
628    /// At most one permit may be stored by `Notify`. Many sequential calls to
629    /// `notify_one` will result in a single permit being stored. The next call to
630    /// `notified().await` will complete immediately, but the one after that
631    /// will wait.
632    ///
633    /// [`notified().await`]: Notify::notified()
634    ///
635    /// # Examples
636    ///
637    /// ```
638    /// use tokio::sync::Notify;
639    /// use std::sync::Arc;
640    ///
641    /// # #[tokio::main(flavor = "current_thread")]
642    /// # async fn main() {
643    /// let notify = Arc::new(Notify::new());
644    /// let notify2 = notify.clone();
645    ///
646    /// tokio::spawn(async move {
647    ///     notify2.notified().await;
648    ///     println!("received notification");
649    /// });
650    ///
651    /// println!("sending notification");
652    /// notify.notify_one();
653    /// # }
654    /// ```
655    // Alias for old name in 0.x
656    #[cfg_attr(docsrs, doc(alias = "notify"))]
657    pub fn notify_one(&self) {
658        self.notify_with_strategy(NotifyOneStrategy::Fifo);
659    }
660
661    /// Notifies the last waiting task.
662    ///
663    /// This function behaves similar to `notify_one`. The only difference is that it wakes
664    /// the most recently added waiter instead of the oldest waiter.
665    ///
666    /// Check the [`notify_one()`] documentation for more info and
667    /// examples.
668    ///
669    /// [`notify_one()`]: Notify::notify_one
670    pub fn notify_last(&self) {
671        self.notify_with_strategy(NotifyOneStrategy::Lifo);
672    }
673
674    fn notify_with_strategy(&self, strategy: NotifyOneStrategy) {
675        // Load the current state
676        let mut curr = self.state.load(SeqCst);
677
678        // If the state is `EMPTY`, transition to `NOTIFIED` and return.
679        while let EMPTY | NOTIFIED = get_state(curr) {
680            // The compare-exchange from `NOTIFIED` -> `NOTIFIED` is intended. A
681            // happens-before synchronization must happen between this atomic
682            // operation and a task calling `notified().await`.
683            let new = set_state(curr, NOTIFIED);
684            let res = self.state.compare_exchange(curr, new, SeqCst, SeqCst);
685
686            match res {
687                // No waiters, no further work to do
688                Ok(_) => return,
689                Err(actual) => {
690                    curr = actual;
691                }
692            }
693        }
694
695        // There are waiters, the lock must be acquired to notify.
696        let mut waiters = self.waiters.lock();
697
698        // The state must be reloaded while the lock is held. The state may only
699        // transition out of WAITING while the lock is held.
700        curr = self.state.load(SeqCst);
701
702        if let Some(waker) = notify_locked(&mut waiters, &self.state, curr, strategy) {
703            drop(waiters);
704            waker.wake();
705        }
706    }
707
708    /// Notifies all waiting tasks.
709    ///
710    /// If a task is currently waiting, that task is notified. Unlike with
711    /// `notify_one()`, no permit is stored to be used by the next call to
712    /// `notified().await`. The purpose of this method is to notify all
713    /// already registered waiters. Registering for notification is done by
714    /// acquiring an instance of the `Notified` future via calling `notified()`.
715    ///
716    /// # Examples
717    ///
718    /// ```
719    /// use tokio::sync::Notify;
720    /// use std::sync::Arc;
721    ///
722    /// # #[tokio::main(flavor = "current_thread")]
723    /// # async fn main() {
724    /// let notify = Arc::new(Notify::new());
725    /// let notify2 = notify.clone();
726    ///
727    /// let notified1 = notify.notified();
728    /// let notified2 = notify.notified();
729    ///
730    /// let handle = tokio::spawn(async move {
731    ///     println!("sending notifications");
732    ///     notify2.notify_waiters();
733    /// });
734    ///
735    /// notified1.await;
736    /// notified2.await;
737    /// println!("received notifications");
738    /// # }
739    /// ```
740    pub fn notify_waiters(&self) {
741        self.lock_waiter_list().notify_waiters();
742    }
743
744    fn inner_notify_waiters<'a>(
745        &'a self,
746        curr: usize,
747        mut waiters: crate::loom::sync::MutexGuard<'a, LinkedList<Waiter>>,
748    ) {
749        if matches!(get_state(curr), EMPTY | NOTIFIED) {
750            // There are no waiting tasks. All we need to do is increment the
751            // number of times this method was called.
752            atomic_inc_num_notify_waiters_calls(&self.state);
753            return;
754        }
755
756        // Increment the number of times this method was called
757        // and transition to empty.
758        let new_state = set_state(inc_num_notify_waiters_calls(curr), EMPTY);
759        self.state.store(new_state, SeqCst);
760
761        // It is critical for `GuardedLinkedList` safety that the guard node is
762        // pinned in memory and is not dropped until the guarded list is dropped.
763        let guard = Waiter::new();
764        pin!(guard);
765
766        // We move all waiters to a secondary list. It uses a `GuardedLinkedList`
767        // underneath to allow every waiter to safely remove itself from it.
768        //
769        // * This list will be still guarded by the `waiters` lock.
770        //   `NotifyWaitersList` wrapper makes sure we hold the lock to modify it.
771        // * This wrapper will empty the list on drop. It is critical for safety
772        //   that we will not leave any list entry with a pointer to the local
773        //   guard node after this function returns / panics.
774        let mut list = NotifyWaitersList::new(std::mem::take(&mut *waiters), guard.as_ref(), self);
775
776        let mut wakers = WakeList::new();
777        'outer: loop {
778            while wakers.can_push() {
779                match list.pop_back_locked(&mut waiters) {
780                    Some(waiter) => {
781                        // Safety: we never make mutable references to waiters.
782                        let waiter = unsafe { waiter.as_ref() };
783
784                        // Safety: we hold the lock, so we can access the waker.
785                        if let Some(waker) =
786                            unsafe { waiter.waker.with_mut(|waker| (*waker).take()) }
787                        {
788                            wakers.push(waker);
789                        }
790
791                        // This waiter is unlinked and will not be shared ever again, release it.
792                        waiter.notification.store_release(Notification::All);
793                    }
794                    None => {
795                        break 'outer;
796                    }
797                }
798            }
799
800            // Release the lock before notifying.
801            drop(waiters);
802
803            // One of the wakers may panic, but the remaining waiters will still
804            // be unlinked from the list in `NotifyWaitersList` destructor.
805            wakers.wake_all();
806
807            // Acquire the lock again.
808            waiters = self.waiters.lock();
809        }
810
811        // Release the lock before notifying
812        drop(waiters);
813
814        wakers.wake_all();
815    }
816
817    pub(crate) fn lock_waiter_list(&self) -> NotifyGuard<'_> {
818        let guarded_waiters = self.waiters.lock();
819
820        // The state must be loaded while the lock is held. The state may only
821        // transition out of WAITING while the lock is held.
822        let current_state = self.state.load(SeqCst);
823
824        NotifyGuard {
825            guarded_notify: self,
826            guarded_waiters,
827            current_state,
828        }
829    }
830}
831
832impl Default for Notify {
833    fn default() -> Notify {
834        Notify::new()
835    }
836}
837
838impl UnwindSafe for Notify {}
839impl RefUnwindSafe for Notify {}
840
841fn notify_locked(
842    waiters: &mut LinkedList<Waiter>,
843    state: &AtomicUsize,
844    curr: usize,
845    strategy: NotifyOneStrategy,
846) -> Option<Waker> {
847    match get_state(curr) {
848        EMPTY | NOTIFIED => {
849            let res = state.compare_exchange(curr, set_state(curr, NOTIFIED), SeqCst, SeqCst);
850
851            match res {
852                Ok(_) => None,
853                Err(actual) => {
854                    let actual_state = get_state(actual);
855                    assert!(actual_state == EMPTY || actual_state == NOTIFIED);
856                    state.store(set_state(actual, NOTIFIED), SeqCst);
857                    None
858                }
859            }
860        }
861        WAITING => {
862            // At this point, it is guaranteed that the state will not
863            // concurrently change as holding the lock is required to
864            // transition **out** of `WAITING`.
865            //
866            // Get a pending waiter using one of the available dequeue strategies.
867            let waiter = match strategy {
868                NotifyOneStrategy::Fifo => waiters.pop_back().unwrap(),
869                NotifyOneStrategy::Lifo => waiters.pop_front().unwrap(),
870            };
871
872            // Safety: we never make mutable references to waiters.
873            let waiter = unsafe { waiter.as_ref() };
874
875            // Safety: we hold the lock, so we can access the waker.
876            let waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) };
877
878            // This waiter is unlinked and will not be shared ever again, release it.
879            waiter
880                .notification
881                .store_release(Notification::One(strategy));
882
883            if waiters.is_empty() {
884                // As this the **final** waiter in the list, the state
885                // must be transitioned to `EMPTY`. As transitioning
886                // **from** `WAITING` requires the lock to be held, a
887                // `store` is sufficient.
888                state.store(set_state(curr, EMPTY), SeqCst);
889            }
890            waker
891        }
892        _ => unreachable!(),
893    }
894}
895
896// ===== impl Notified =====
897
898impl Notified<'_> {
899    /// Adds this future to the list of futures that are ready to receive
900    /// wakeups from calls to [`notify_one`].
901    ///
902    /// Polling the future also adds it to the list, so this method should only
903    /// be used if you want to add the future to the list before the first call
904    /// to `poll`. (In fact, this method is equivalent to calling `poll` except
905    /// that no `Waker` is registered.)
906    ///
907    /// This has no effect on notifications sent using [`notify_waiters`], which
908    /// are received as long as they happen after the creation of the `Notified`
909    /// regardless of whether `enable` or `poll` has been called.
910    ///
911    /// This method returns true if the `Notified` is ready. This happens in the
912    /// following situations:
913    ///
914    ///  1. The `notify_waiters` method was called between the creation of the
915    ///     `Notified` and the call to this method.
916    ///  2. This is the first call to `enable` or `poll` on this future, and the
917    ///     `Notify` was holding a permit from a previous call to `notify_one`.
918    ///     The call consumes the permit in that case.
919    ///  3. The future has previously been enabled or polled, and it has since
920    ///     then been marked ready by either consuming a permit from the
921    ///     `Notify`, or by a call to `notify_one` or `notify_waiters` that
922    ///     removed it from the list of futures ready to receive wakeups.
923    ///
924    /// If this method returns true, any future calls to poll on the same future
925    /// will immediately return `Poll::Ready`.
926    ///
927    /// # Examples
928    ///
929    /// Unbound multi-producer multi-consumer (mpmc) channel.
930    ///
931    /// The call to `enable` is important because otherwise if you have two
932    /// calls to `recv` and two calls to `send` in parallel, the following could
933    /// happen:
934    ///
935    ///  1. Both calls to `try_recv` return `None`.
936    ///  2. Both new elements are added to the vector.
937    ///  3. The `notify_one` method is called twice, adding only a single
938    ///     permit to the `Notify`.
939    ///  4. Both calls to `recv` reach the `Notified` future. One of them
940    ///     consumes the permit, and the other sleeps forever.
941    ///
942    /// By adding the `Notified` futures to the list by calling `enable` before
943    /// `try_recv`, the `notify_one` calls in step three would remove the
944    /// futures from the list and mark them notified instead of adding a permit
945    /// to the `Notify`. This ensures that both futures are woken.
946    ///
947    /// ```
948    /// use tokio::sync::Notify;
949    ///
950    /// use std::collections::VecDeque;
951    /// use std::sync::Mutex;
952    ///
953    /// struct Channel<T> {
954    ///     messages: Mutex<VecDeque<T>>,
955    ///     notify_on_sent: Notify,
956    /// }
957    ///
958    /// impl<T> Channel<T> {
959    ///     pub fn send(&self, msg: T) {
960    ///         let mut locked_queue = self.messages.lock().unwrap();
961    ///         locked_queue.push_back(msg);
962    ///         drop(locked_queue);
963    ///
964    ///         // Send a notification to one of the calls currently
965    ///         // waiting in a call to `recv`.
966    ///         self.notify_on_sent.notify_one();
967    ///     }
968    ///
969    ///     pub fn try_recv(&self) -> Option<T> {
970    ///         let mut locked_queue = self.messages.lock().unwrap();
971    ///         locked_queue.pop_front()
972    ///     }
973    ///
974    ///     pub async fn recv(&self) -> T {
975    ///         let future = self.notify_on_sent.notified();
976    ///         tokio::pin!(future);
977    ///
978    ///         loop {
979    ///             // Make sure that no wakeup is lost if we get
980    ///             // `None` from `try_recv`.
981    ///             future.as_mut().enable();
982    ///
983    ///             if let Some(msg) = self.try_recv() {
984    ///                 return msg;
985    ///             }
986    ///
987    ///             // Wait for a call to `notify_one`.
988    ///             //
989    ///             // This uses `.as_mut()` to avoid consuming the future,
990    ///             // which lets us call `Pin::set` below.
991    ///             future.as_mut().await;
992    ///
993    ///             // Reset the future in case another call to
994    ///             // `try_recv` got the message before us.
995    ///             future.set(self.notify_on_sent.notified());
996    ///         }
997    ///     }
998    /// }
999    /// ```
1000    ///
1001    /// [`notify_one`]: Notify::notify_one()
1002    /// [`notify_waiters`]: Notify::notify_waiters()
1003    pub fn enable(self: Pin<&mut Self>) -> bool {
1004        self.poll_notified(None).is_ready()
1005    }
1006
1007    fn project(self: Pin<&mut Self>) -> NotifiedProject<'_> {
1008        unsafe {
1009            // Safety: `notify`, `state` and `notify_waiters_calls` are `Unpin`.
1010
1011            is_unpin::<&Notify>();
1012            is_unpin::<State>();
1013            is_unpin::<usize>();
1014
1015            let me = self.get_unchecked_mut();
1016            NotifiedProject {
1017                notify: me.notify,
1018                state: &mut me.state,
1019                notify_waiters_calls: &me.notify_waiters_calls,
1020                waiter: &me.waiter,
1021            }
1022        }
1023    }
1024
1025    fn poll_notified(self: Pin<&mut Self>, waker: Option<&Waker>) -> Poll<()> {
1026        self.project().poll_notified(waker)
1027    }
1028}
1029
1030impl Future for Notified<'_> {
1031    type Output = ();
1032
1033    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
1034        self.poll_notified(Some(cx.waker()))
1035    }
1036}
1037
1038impl Drop for Notified<'_> {
1039    fn drop(&mut self) {
1040        // Safety: The type only transitions to a "Waiting" state when pinned.
1041        unsafe { Pin::new_unchecked(self) }
1042            .project()
1043            .drop_notified();
1044    }
1045}
1046
1047// ===== impl OwnedNotified =====
1048
1049impl OwnedNotified {
1050    /// Adds this future to the list of futures that are ready to receive
1051    /// wakeups from calls to [`notify_one`].
1052    ///
1053    /// See [`Notified::enable`] for more details.
1054    ///
1055    /// [`notify_one`]: Notify::notify_one()
1056    pub fn enable(self: Pin<&mut Self>) -> bool {
1057        self.poll_notified(None).is_ready()
1058    }
1059
1060    /// A custom `project` implementation is used in place of `pin-project-lite`
1061    /// as a custom drop implementation is needed.
1062    fn project(self: Pin<&mut Self>) -> NotifiedProject<'_> {
1063        unsafe {
1064            // Safety: `notify`, `state` and `notify_waiters_calls` are `Unpin`.
1065
1066            is_unpin::<&Notify>();
1067            is_unpin::<State>();
1068            is_unpin::<usize>();
1069
1070            let me = self.get_unchecked_mut();
1071            NotifiedProject {
1072                notify: &me.notify,
1073                state: &mut me.state,
1074                notify_waiters_calls: &me.notify_waiters_calls,
1075                waiter: &me.waiter,
1076            }
1077        }
1078    }
1079
1080    fn poll_notified(self: Pin<&mut Self>, waker: Option<&Waker>) -> Poll<()> {
1081        self.project().poll_notified(waker)
1082    }
1083}
1084
1085impl Future for OwnedNotified {
1086    type Output = ();
1087
1088    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
1089        self.poll_notified(Some(cx.waker()))
1090    }
1091}
1092
1093impl Drop for OwnedNotified {
1094    fn drop(&mut self) {
1095        // Safety: The type only transitions to a "Waiting" state when pinned.
1096        unsafe { Pin::new_unchecked(self) }
1097            .project()
1098            .drop_notified();
1099    }
1100}
1101
1102// ===== impl NotifiedProject =====
1103
1104impl NotifiedProject<'_> {
1105    fn poll_notified(self, waker: Option<&Waker>) -> Poll<()> {
1106        let NotifiedProject {
1107            notify,
1108            state,
1109            notify_waiters_calls,
1110            waiter,
1111        } = self;
1112
1113        'outer_loop: loop {
1114            match *state {
1115                State::Init => {
1116                    let curr = notify.state.load(SeqCst);
1117
1118                    // Check if `notify_waiters` was called before attempting to acquire
1119                    // the `NOTIFIED` state. If a broadcast occurred, we will be woken by it,
1120                    // leaving the `notify_one` permit for other waiters.
1121                    if get_num_notify_waiters_calls(curr) != *notify_waiters_calls {
1122                        *state = State::Done;
1123                        continue 'outer_loop;
1124                    }
1125
1126                    // Optimistically try acquiring a pending notification
1127                    let res = notify.state.compare_exchange(
1128                        set_state(curr, NOTIFIED),
1129                        set_state(curr, EMPTY),
1130                        SeqCst,
1131                        SeqCst,
1132                    );
1133
1134                    if res.is_ok() {
1135                        // Acquired the notification
1136                        *state = State::Done;
1137                        continue 'outer_loop;
1138                    }
1139
1140                    // Clone the waker before locking, a waker clone can be
1141                    // triggering arbitrary code.
1142                    let waker = waker.cloned();
1143
1144                    // Acquire the lock and attempt to transition to the waiting
1145                    // state.
1146                    let mut waiters = notify.waiters.lock();
1147
1148                    // Reload the state with the lock held
1149                    let mut curr = notify.state.load(SeqCst);
1150
1151                    // if notify_waiters has been called after the future
1152                    // was created, then we are done
1153                    if get_num_notify_waiters_calls(curr) != *notify_waiters_calls {
1154                        *state = State::Done;
1155                        continue 'outer_loop;
1156                    }
1157
1158                    // Transition the state to WAITING.
1159                    loop {
1160                        match get_state(curr) {
1161                            EMPTY => {
1162                                // Transition to WAITING
1163                                let res = notify.state.compare_exchange(
1164                                    set_state(curr, EMPTY),
1165                                    set_state(curr, WAITING),
1166                                    SeqCst,
1167                                    SeqCst,
1168                                );
1169
1170                                if let Err(actual) = res {
1171                                    assert_eq!(get_state(actual), NOTIFIED);
1172                                    curr = actual;
1173                                } else {
1174                                    break;
1175                                }
1176                            }
1177                            WAITING => break,
1178                            NOTIFIED => {
1179                                // Try consuming the notification
1180                                let res = notify.state.compare_exchange(
1181                                    set_state(curr, NOTIFIED),
1182                                    set_state(curr, EMPTY),
1183                                    SeqCst,
1184                                    SeqCst,
1185                                );
1186
1187                                match res {
1188                                    Ok(_) => {
1189                                        // Acquired the notification
1190                                        *state = State::Done;
1191                                        continue 'outer_loop;
1192                                    }
1193                                    Err(actual) => {
1194                                        assert_eq!(get_state(actual), EMPTY);
1195                                        curr = actual;
1196                                    }
1197                                }
1198                            }
1199                            _ => unreachable!(),
1200                        }
1201                    }
1202
1203                    let mut old_waker = None;
1204                    if waker.is_some() {
1205                        // Safety: called while locked.
1206                        //
1207                        // The use of `old_waiter` here is not necessary, as the field is always
1208                        // None when we reach this line.
1209                        unsafe {
1210                            old_waker =
1211                                waiter.waker.with_mut(|v| std::mem::replace(&mut *v, waker));
1212                        }
1213                    }
1214
1215                    // Insert the waiter into the linked list
1216                    waiters.push_front(NonNull::from(waiter));
1217
1218                    *state = State::Waiting;
1219
1220                    drop(waiters);
1221                    drop(old_waker);
1222
1223                    return Poll::Pending;
1224                }
1225                State::Waiting => {
1226                    #[cfg(feature = "taskdump")]
1227                    if let Some(_waker) = waker {
1228                        std::task::ready!(crate::trace::trace_leaf());
1229                    }
1230
1231                    if waiter.notification.load(Acquire).is_some() {
1232                        // Safety: waiter is already unlinked and will not be shared again,
1233                        // so we have an exclusive access to `waker`.
1234                        drop(unsafe { waiter.waker.with_mut(|waker| (*waker).take()) });
1235
1236                        waiter.notification.clear();
1237                        *state = State::Done;
1238                        return Poll::Ready(());
1239                    }
1240
1241                    // Our waiter was not notified, implying it is still stored in a waiter
1242                    // list (guarded by `notify.waiters`). In order to access the waker
1243                    // fields, we must acquire the lock.
1244
1245                    let mut old_waker = None;
1246                    let mut waiters = notify.waiters.lock();
1247
1248                    // We hold the lock and notifications are set only with the lock held,
1249                    // so this can be relaxed, because the happens-before relationship is
1250                    // established through the mutex.
1251                    if waiter.notification.load(Relaxed).is_some() {
1252                        // Safety: waiter is already unlinked and will not be shared again,
1253                        // so we have an exclusive access to `waker`.
1254                        old_waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) };
1255
1256                        waiter.notification.clear();
1257
1258                        // Drop the old waker after releasing the lock.
1259                        drop(waiters);
1260                        drop(old_waker);
1261
1262                        *state = State::Done;
1263                        return Poll::Ready(());
1264                    }
1265
1266                    // Load the state with the lock held.
1267                    let curr = notify.state.load(SeqCst);
1268
1269                    if get_num_notify_waiters_calls(curr) != *notify_waiters_calls {
1270                        // Before we add a waiter to the list we check if these numbers are
1271                        // different while holding the lock. If these numbers are different now,
1272                        // it means that there is a call to `notify_waiters` in progress and this
1273                        // waiter must be contained by a guarded list used in `notify_waiters`.
1274                        // We can treat the waiter as notified and remove it from the list, as
1275                        // it would have been notified in the `notify_waiters` call anyways.
1276
1277                        // Safety: we hold the lock, so we can modify the waker.
1278                        old_waker = unsafe { waiter.waker.with_mut(|waker| (*waker).take()) };
1279
1280                        // Safety: we hold the lock, so we have an exclusive access to the list.
1281                        // The list is used in `notify_waiters`, so it must be guarded.
1282                        unsafe { waiters.remove(NonNull::from(waiter)) };
1283
1284                        *state = State::Done;
1285                    } else {
1286                        // Safety: we hold the lock, so we can modify the waker.
1287                        unsafe {
1288                            waiter.waker.with_mut(|v| {
1289                                if let Some(waker) = waker {
1290                                    let should_update = match &*v {
1291                                        Some(current_waker) => !current_waker.will_wake(waker),
1292                                        None => true,
1293                                    };
1294                                    if should_update {
1295                                        old_waker = (*v).replace(waker.clone());
1296                                    }
1297                                }
1298                            });
1299                        }
1300
1301                        // Drop the old waker after releasing the lock.
1302                        drop(waiters);
1303                        drop(old_waker);
1304
1305                        return Poll::Pending;
1306                    }
1307
1308                    // Explicit drop of the lock to indicate the scope that the
1309                    // lock is held. Because holding the lock is required to
1310                    // ensure safe access to fields not held within the lock, it
1311                    // is helpful to visualize the scope of the critical
1312                    // section.
1313                    drop(waiters);
1314
1315                    // Drop the old waker after releasing the lock.
1316                    drop(old_waker);
1317                }
1318                State::Done => {
1319                    #[cfg(feature = "taskdump")]
1320                    if let Some(_waker) = waker {
1321                        std::task::ready!(crate::trace::trace_leaf());
1322                    }
1323                    return Poll::Ready(());
1324                }
1325            }
1326        }
1327    }
1328
1329    fn drop_notified(self) {
1330        let NotifiedProject {
1331            notify,
1332            state,
1333            waiter,
1334            ..
1335        } = self;
1336
1337        // This is where we ensure safety. The `Notified` value is being
1338        // dropped, which means we must ensure that the waiter entry is no
1339        // longer stored in the linked list.
1340        if matches!(*state, State::Waiting) {
1341            let mut waiters = notify.waiters.lock();
1342            let mut notify_state = notify.state.load(SeqCst);
1343
1344            // We hold the lock, so this field is not concurrently accessed by
1345            // `notify_*` functions and we can use the relaxed ordering.
1346            let notification = waiter.notification.load(Relaxed);
1347
1348            // remove the entry from the list (if not already removed)
1349            //
1350            // Safety: we hold the lock, so we have an exclusive access to every list the
1351            // waiter may be contained in. If the node is not contained in the `waiters`
1352            // list, then it is contained by a guarded list used by `notify_waiters`.
1353            unsafe { waiters.remove(NonNull::from(waiter)) };
1354
1355            if waiters.is_empty() && get_state(notify_state) == WAITING {
1356                notify_state = set_state(notify_state, EMPTY);
1357                notify.state.store(notify_state, SeqCst);
1358            }
1359
1360            // See if the node was notified but not received. In this case, if
1361            // the notification was triggered via `notify_one`, it must be sent
1362            // to the next waiter.
1363            if let Some(Notification::One(strategy)) = notification {
1364                if let Some(waker) =
1365                    notify_locked(&mut waiters, &notify.state, notify_state, strategy)
1366                {
1367                    drop(waiters);
1368                    waker.wake();
1369                }
1370            }
1371        }
1372    }
1373}
1374
1375/// # Safety
1376///
1377/// `Waiter` is forced to be !Unpin.
1378unsafe impl linked_list::Link for Waiter {
1379    type Handle = NonNull<Waiter>;
1380    type Target = Waiter;
1381
1382    fn as_raw(handle: &NonNull<Waiter>) -> NonNull<Waiter> {
1383        *handle
1384    }
1385
1386    unsafe fn from_raw(ptr: NonNull<Waiter>) -> NonNull<Waiter> {
1387        ptr
1388    }
1389
1390    unsafe fn pointers(target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {
1391        unsafe { Waiter::addr_of_pointers(target) }
1392    }
1393}
1394
1395fn is_unpin<T: Unpin>() {}
1396
1397/// A guard that provides exclusive access to a `Notify`'s internal
1398/// waiters list.
1399///
1400/// While this guard is held, the `Notify` instance's waiter list is locked.
1401pub(crate) struct NotifyGuard<'a> {
1402    guarded_notify: &'a Notify,
1403    guarded_waiters: crate::loom::sync::MutexGuard<'a, LinkedList<Waiter>>,
1404    current_state: usize,
1405}
1406
1407impl NotifyGuard<'_> {
1408    pub(crate) fn notify_waiters(self) {
1409        self.guarded_notify
1410            .inner_notify_waiters(self.current_state, self.guarded_waiters);
1411    }
1412}