ntex_util/
task.rs

1//! A synchronization primitive for task wakeup.
2use std::{cell::Cell, fmt, marker::PhantomData, rc, task::Waker};
3
4/// A synchronization primitive for task wakeup.
5///
6/// Sometimes the task interested in a given event will change over time.
7/// An `LocalWaker` can coordinate concurrent 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    /// Create an `LocalWaker`.
28    pub fn new() -> Self {
29        LocalWaker {
30            waker: Cell::new(None),
31            _t: PhantomData,
32        }
33    }
34
35    /// Create an `LocalWaker`.
36    pub fn with(waker: Option<Waker>) -> Self {
37        LocalWaker {
38            waker: Cell::new(waker),
39            _t: PhantomData,
40        }
41    }
42
43    #[inline]
44    /// Registers the waker to be notified on calls to `wake`.
45    ///
46    /// Returns `true` if waker was registered before.
47    pub fn register(&self, waker: &Waker) -> bool {
48        self.waker.replace(Some(waker.clone())).is_some()
49    }
50
51    #[inline]
52    /// Calls `wake` on the last `Waker` passed to `register`.
53    ///
54    /// If `register` has not been called yet, then this does nothing.
55    pub fn wake(&self) {
56        if let Some(waker) = self.take() {
57            waker.wake();
58        }
59    }
60
61    #[inline]
62    /// Calls `wake` on the last `Waker` passed to `register`.
63    ///
64    /// If `register` has not been called yet, then this returns `false`.
65    pub fn wake_checked(&self) -> bool {
66        if let Some(waker) = self.take() {
67            waker.wake();
68            true
69        } else {
70            false
71        }
72    }
73
74    /// Returns the last `Waker` passed to `register`, so that the user can wake it.
75    ///
76    /// If a waker has not been registered, this returns `None`.
77    pub fn take(&self) -> Option<Waker> {
78        self.waker.take()
79    }
80}
81
82impl Clone for LocalWaker {
83    fn clone(&self) -> Self {
84        LocalWaker::new()
85    }
86}
87
88impl fmt::Debug for LocalWaker {
89    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
90        write!(f, "LocalWaker")
91    }
92}
93
94/// Yields execution back to the current runtime.
95pub async fn yield_to() {
96    use std::{future::Future, pin::Pin, task::Context, task::Poll};
97
98    struct Yield {
99        completed: bool,
100    }
101
102    impl Future for Yield {
103        type Output = ();
104
105        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
106            if self.completed {
107                return Poll::Ready(());
108            }
109
110            self.completed = true;
111            cx.waker().wake_by_ref();
112
113            Poll::Pending
114        }
115    }
116
117    Yield { completed: false }.await;
118}
119
120#[cfg(test)]
121mod test {
122    use super::*;
123
124    #[ntex_macros::rt_test2]
125    async fn yield_test() {
126        yield_to().await;
127    }
128}
ntex_util/task.rs

ntex_util/
task.rs
