diatomic_waker/arc_waker.rs
1use alloc::sync::Arc;
2use core::task::Waker;
3
4use crate::DiatomicWaker;
5use crate::WaitUntil;
6
7/// An owned object that can await notifications from one or several
8/// [`WakeSource`]s.
9///
10/// See the [crate-level documentation](crate) for usage.
11#[derive(Debug, Default)]
12pub struct WakeSink {
13 /// The shared data.
14 inner: Arc<DiatomicWaker>,
15}
16
17impl WakeSink {
18 /// Creates a new sink.
19 pub fn new() -> Self {
20 Self {
21 inner: Arc::new(DiatomicWaker::new()),
22 }
23 }
24
25 /// Creates an owned source.
26 #[inline]
27 pub fn source(&self) -> WakeSource {
28 WakeSource {
29 inner: self.inner.clone(),
30 }
31 }
32
33 /// Registers a new waker.
34 ///
35 /// Registration is lazy: the waker is cloned only if it differs from the
36 /// last registered waker (note that the last registered waker is cached
37 /// even if it was unregistered).
38 #[inline]
39 pub fn register(&mut self, waker: &Waker) {
40 // Safety: `DiatomicWaker::register`, `DiatomicWaker::unregister` and
41 // `DiatomicWaker::wait_until` cannot be used concurrently from multiple
42 // thread since `WakeSink` does not implement `Clone` and the wrappers
43 // of the above methods require exclusive ownership to `WakeSink`.
44 unsafe { self.inner.register(waker) };
45 }
46
47 /// Unregisters the waker.
48 ///
49 /// After the waker is unregistered, subsequent calls to
50 /// `WakeSource::notify` will be ignored.
51 ///
52 /// Note that the previously-registered waker (if any) remains cached.
53 #[inline]
54 pub fn unregister(&mut self) {
55 // Safety: `DiatomicWaker::register`, `DiatomicWaker::unregister` and
56 // `DiatomicWaker::wait_until` cannot be used concurrently from multiple
57 // thread since `WakeSink` does not implement `Clone` and the wrappers
58 // of the above methods require exclusive ownership to `WakeSink`.
59 unsafe { self.inner.unregister() };
60 }
61
62 /// Returns a future that can be `await`ed until the provided predicate
63 /// returns a value.
64 ///
65 /// The predicate is checked each time a notification is received.
66 #[inline]
67 pub fn wait_until<P, T>(&mut self, predicate: P) -> WaitUntil<'_, P, T>
68 where
69 P: FnMut() -> Option<T>,
70 {
71 // Safety: `DiatomicWaker::register`, `DiatomicWaker::unregister` and
72 // `DiatomicWaker::wait_until` cannot be used concurrently from multiple
73 // thread since `WakeSink` does not implement `Clone` and the wrappers
74 // of the above methods require exclusive ownership to `WakeSink`.
75 unsafe { self.inner.wait_until(predicate) }
76 }
77}
78
79/// An owned object that can send notifications to a [`WakeSink`].
80///
81/// See the [crate-level documentation](crate) for usage.
82#[derive(Clone, Debug)]
83pub struct WakeSource {
84 /// The shared data.
85 inner: Arc<DiatomicWaker>,
86}
87
88impl WakeSource {
89 /// Notifies the sink if a waker is registered.
90 ///
91 /// This automatically unregisters any waker that may have been previously
92 /// registered.
93 #[inline]
94 pub fn notify(&self) {
95 self.inner.notify();
96 }
97}