ntex_util/
task.rs

1//! Task wake-up and scheduling utilities.
2use std::{cell::Cell, fmt, marker::PhantomData, rc, task::Waker};
3
4/// A single-threaded synchronization primitive for task wake-up.
5///
6/// Sometimes the task interested in a given event will change over time.
7/// A `LocalWaker` can coordinate notifications with the consumer
8/// potentially "updating" the underlying task to wake up. This is useful in
9/// scenarios where a computation completes in another task and wants to
10/// notify the consumer, but the consumer is in the process of being migrated to
11/// a new logical task.
12///
13/// Consumers should call `register` before checking the result of a computation
14/// and producers should call `wake` after producing the computation (this
15/// differs from the usual `thread::park` pattern). It is also permitted for
16/// `wake` to be called **before** `register`. This results in a no-op.
17///
18/// A single `LocalWaker` may be reused for any number of calls to `register` or
19/// `wake`.
20#[derive(Default)]
21pub struct LocalWaker {
22    waker: Cell<Option<Waker>>,
23    _t: PhantomData<rc::Rc<()>>,
24}
25
26impl LocalWaker {
27    /// Creates an empty `LocalWaker`.
28    pub fn new() -> Self {
29        LocalWaker::with(None)
30    }
31
32    /// Creates a `LocalWaker` with an optional initial waker.
33    pub fn with(waker: Option<Waker>) -> Self {
34        LocalWaker {
35            waker: Cell::new(waker),
36            _t: PhantomData,
37        }
38    }
39
40    #[inline]
41    /// Registers the waker to be notified on calls to `wake`.
42    ///
43    /// Returns `true` if a waker was already registered.
44    pub fn register(&self, waker: &Waker) -> bool {
45        match self.waker.take() {
46            Some(prev) if prev.will_wake(waker) => {
47                self.waker.set(Some(prev));
48                true
49            }
50            prev => {
51                self.waker.set(Some(waker.clone()));
52                prev.is_some()
53            }
54        }
55    }
56
57    #[inline]
58    /// Calls `wake` on the last `Waker` passed to `register`.
59    ///
60    /// The waker is consumed, so a later `wake` does nothing until `register`
61    /// is called again. If no waker is registered, this does nothing.
62    pub fn wake(&self) {
63        if let Some(waker) = self.take() {
64            waker.wake();
65        }
66    }
67
68    #[inline]
69    /// Calls `wake` on the last `Waker` passed to `register`.
70    ///
71    /// Returns `false` if no waker is registered. Like [`wake`](Self::wake),
72    /// this consumes the registered waker.
73    pub fn wake_checked(&self) -> bool {
74        if let Some(waker) = self.take() {
75            waker.wake();
76            true
77        } else {
78            false
79        }
80    }
81
82    /// Returns the last `Waker` passed to `register`, so that the user can wake it.
83    ///
84    /// If a waker has not been registered, this returns `None`.
85    pub fn take(&self) -> Option<Waker> {
86        self.waker.take()
87    }
88
89    #[doc(hidden)]
90    /// Returns `true` if a waker is registered.
91    pub fn is_set(&self) -> bool {
92        let waker = self.waker.take();
93        let set = waker.is_some();
94        self.waker.set(waker);
95        set
96    }
97}
98
99/// Cloning creates an empty `LocalWaker`; the registered waker is not copied.
100impl Clone for LocalWaker {
101    fn clone(&self) -> Self {
102        LocalWaker::new()
103    }
104}
105
106impl fmt::Debug for LocalWaker {
107    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
108        write!(f, "LocalWaker")
109    }
110}
111
112/// Yields execution back to the current runtime.
113pub async fn yield_to() {
114    use std::{future::Future, pin::Pin, task::Context, task::Poll};
115
116    struct Yield {
117        completed: bool,
118    }
119
120    impl Future for Yield {
121        type Output = ();
122
123        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
124            if self.completed {
125                return Poll::Ready(());
126            }
127
128            self.completed = true;
129            cx.waker().wake_by_ref();
130
131            Poll::Pending
132        }
133    }
134
135    Yield { completed: false }.await;
136}
137
138#[cfg(test)]
139mod test {
140    use super::*;
141
142    #[ntex::test]
143    async fn yield_test() {
144        yield_to().await;
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use std::sync::atomic::{AtomicUsize, Ordering};
151    use std::task::{RawWaker, RawWakerVTable};
152
153    use super::*;
154
155    static CLONES: AtomicUsize = AtomicUsize::new(0);
156
157    static VTABLE: RawWakerVTable = RawWakerVTable::new(
158        |p| {
159            CLONES.fetch_add(1, Ordering::Relaxed);
160            RawWaker::new(p, &VTABLE)
161        },
162        |_| {},
163        |_| {},
164        |_| {},
165    );
166
167    #[test]
168    fn test_register_same_waker() {
169        static A: u8 = 0;
170        static B: u8 = 0;
171        let a = unsafe { Waker::from_raw(RawWaker::new((&raw const A).cast(), &VTABLE)) };
172        let b = unsafe { Waker::from_raw(RawWaker::new((&raw const B).cast(), &VTABLE)) };
173
174        let w = LocalWaker::new();
175        assert!(!w.register(&a));
176        assert!(w.register(&a));
177        assert!(w.register(&a));
178        assert_eq!(CLONES.load(Ordering::Relaxed), 1);
179
180        assert!(w.register(&b));
181        assert_eq!(CLONES.load(Ordering::Relaxed), 2);
182        assert!(w.take().unwrap().will_wake(&b));
183    }
184
185    #[test]
186    fn local_waker_clone_is_empty() {
187        let waker = LocalWaker::new();
188        waker.register(std::task::Waker::noop());
189        assert!(waker.is_set());
190        let cloned = waker.clone();
191        assert!(!cloned.is_set());
192        assert!(!cloned.wake_checked());
193        assert_eq!(format!("{waker:?}"), "LocalWaker");
194    }
195}
ntex_util/task.rs

ntex_util/
task.rs
