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        self.waker.replace(Some(waker.clone())).is_some()
46    }
47
48    #[inline]
49    /// Calls `wake` on the last `Waker` passed to `register`.
50    ///
51    /// The waker is consumed, so a later `wake` does nothing until `register`
52    /// is called again. If no waker is registered, this does nothing.
53    pub fn wake(&self) {
54        if let Some(waker) = self.take() {
55            waker.wake();
56        }
57    }
58
59    #[inline]
60    /// Calls `wake` on the last `Waker` passed to `register`.
61    ///
62    /// Returns `false` if no waker is registered. Like [`wake`](Self::wake),
63    /// this consumes the registered waker.
64    pub fn wake_checked(&self) -> bool {
65        if let Some(waker) = self.take() {
66            waker.wake();
67            true
68        } else {
69            false
70        }
71    }
72
73    /// Returns the last `Waker` passed to `register`, so that the user can wake it.
74    ///
75    /// If a waker has not been registered, this returns `None`.
76    pub fn take(&self) -> Option<Waker> {
77        self.waker.take()
78    }
79
80    #[doc(hidden)]
81    /// Returns `true` if a waker is registered.
82    pub fn is_set(&self) -> bool {
83        let waker = self.waker.take();
84        let set = waker.is_some();
85        self.waker.set(waker);
86        set
87    }
88}
89
90/// Cloning creates an empty `LocalWaker`; the registered waker is not copied.
91impl Clone for LocalWaker {
92    fn clone(&self) -> Self {
93        LocalWaker::new()
94    }
95}
96
97impl fmt::Debug for LocalWaker {
98    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
99        write!(f, "LocalWaker")
100    }
101}
102
103/// Yields execution back to the current runtime.
104pub async fn yield_to() {
105    use std::{future::Future, pin::Pin, task::Context, task::Poll};
106
107    struct Yield {
108        completed: bool,
109    }
110
111    impl Future for Yield {
112        type Output = ();
113
114        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
115            if self.completed {
116                return Poll::Ready(());
117            }
118
119            self.completed = true;
120            cx.waker().wake_by_ref();
121
122            Poll::Pending
123        }
124    }
125
126    Yield { completed: false }.await;
127}
128
129#[cfg(test)]
130mod test {
131    use super::*;
132
133    #[ntex::test]
134    async fn yield_test() {
135        yield_to().await;
136    }
137}
ntex_util/task.rs

ntex_util/
task.rs
