scrappy_utils/
task.rs

1use std::cell::UnsafeCell;
2use std::marker::PhantomData;
3use std::task::Waker;
4use std::{fmt, rc};
5
6/// A synchronization primitive for task wakeup.
7///
8/// Sometimes the task interested in a given event will change over time.
9/// An `LocalWaker` can coordinate concurrent notifications with the consumer
10/// potentially "updating" the underlying task to wake up. This is useful in
11/// scenarios where a computation completes in another task and wants to
12/// notify the consumer, but the consumer is in the process of being migrated to
13/// a new logical task.
14///
15/// Consumers should call `register` before checking the result of a computation
16/// and producers should call `wake` after producing the computation (this
17/// differs from the usual `thread::park` pattern). It is also permitted for
18/// `wake` to be called **before** `register`. This results in a no-op.
19///
20/// A single `AtomicWaker` may be reused for any number of calls to `register` or
21/// `wake`.
22#[derive(Default)]
23pub struct LocalWaker {
24    pub(crate) waker: UnsafeCell<Option<Waker>>,
25    _t: PhantomData<rc::Rc<()>>,
26}
27
28impl LocalWaker {
29    /// Create an `LocalWaker`.
30    pub fn new() -> Self {
31        LocalWaker {
32            waker: UnsafeCell::new(None),
33            _t: PhantomData,
34        }
35    }
36
37    #[inline]
38    /// Check if waker has been registered.
39    pub fn is_registed(&self) -> bool {
40        unsafe { (*self.waker.get()).is_some() }
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        unsafe {
49            let w = self.waker.get();
50            let is_registered = (*w).is_some();
51            *w = Some(waker.clone());
52            is_registered
53        }
54    }
55
56    #[inline]
57    /// Calls `wake` on the last `Waker` passed to `register`.
58    ///
59    /// If `register` has not been called yet, then this does nothing.
60    pub fn wake(&self) {
61        if let Some(waker) = self.take() {
62            waker.wake();
63        }
64    }
65
66    /// Returns the last `Waker` passed to `register`, so that the user can wake it.
67    ///
68    /// If a waker has not been registered, this returns `None`.
69    pub fn take(&self) -> Option<Waker> {
70        unsafe { (*self.waker.get()).take() }
71    }
72}
73
74impl fmt::Debug for LocalWaker {
75    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
76        write!(f, "LocalWaker")
77    }
78}
scrappy_utils/task.rs

scrappy_utils/
task.rs
