synchrony/
waker_slot.rs

1//! A slot holds up to one waker for task wakeup.
2//!
3//! `sync` version is just [`futures_util::task::AtomicWaker`]; unsync version
4//! is a hand-rolled singlethreaded version with similar API.
5
6/// Multithreaded `WakerSlot` based on [`futures_util::task::AtomicWaker`].
7pub mod sync {
8    pub use futures_util::task::AtomicWaker as WakerSlot;
9
10    impl crate::AssertMt for WakerSlot {}
11}
12
13/// Singlethreaded `WakerSlot`
14pub mod unsync {
15    use std::{cell::RefCell, task::Waker};
16
17    /// A singlethreaded registry holds up to one waker for task wakeup.
18    #[derive(Debug, Default)]
19    pub struct WakerSlot {
20        waker: RefCell<Option<Waker>>,
21    }
22
23    impl WakerSlot {
24        /// Create a new [`WakerSlot`]
25        pub const fn new() -> Self {
26            Self {
27                waker: RefCell::new(None),
28            }
29        }
30
31        /// Register given waker
32        pub fn register(&self, waker: &Waker) {
33            let mut w = self.waker.borrow_mut();
34            // Avoid unnecessary clone if two wakers point to the same task
35            if w.as_ref().is_some_and(|x| x.will_wake(waker)) {
36                return;
37            }
38            *w = Some(waker.clone())
39        }
40
41        /// Try to take the stored waker
42        pub fn take(&self) -> Option<Waker> {
43            self.waker.borrow_mut().take()
44        }
45
46        /// Wake currently stored waker
47        pub fn wake(&self) {
48            if let Some(waker) = self.take() {
49                waker.wake()
50            }
51        }
52    }
53}