1use std::{
2 cell::Cell,
3 future::Future,
4 pin::Pin,
5 rc::Rc,
6 sync::{
7 Arc,
8 Mutex,
9 atomic::AtomicBool,
10 },
11 task::{
12 Context,
13 Poll,
14 Waker,
15 },
16};
17
18#[derive(Clone, Default)]
19pub struct Notify {
20 state: Rc<State>,
21}
22
23#[derive(Default)]
24struct State {
25 flag: Cell<bool>,
26 waker: Cell<Option<Waker>>,
27}
28
29impl Notify {
30 pub fn new() -> Self {
31 Self {
32 state: Rc::new(State {
33 flag: Cell::new(false),
34 waker: Cell::new(None),
35 }),
36 }
37 }
38
39 pub fn notify(&self) {
40 self.state.flag.set(true);
41
42 if let Some(w) = self.state.waker.take() {
43 w.wake();
44 }
45 }
46
47 pub fn notified(&self) -> Notified {
48 Notified {
49 state: self.state.clone(),
50 }
51 }
52}
53
54pub struct Notified {
55 state: Rc<State>,
56}
57
58impl Future for Notified {
59 type Output = ();
60
61 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
62 if self.state.flag.replace(false) {
63 Poll::Ready(())
64 } else {
65 self.state.waker.set(Some(cx.waker().clone()));
66 Poll::Pending
67 }
68 }
69}
70
71#[derive(Clone, Default)]
72pub struct ArcNotify {
73 state: Arc<StateArc>,
74}
75
76#[derive(Default)]
77struct StateArc {
78 flag: AtomicBool,
79 waker: Mutex<Option<Waker>>,
80}
81
82impl ArcNotify {
83 pub fn new() -> Self {
84 Self::default()
85 }
86
87 pub fn notify(&self) {
88 self.state
89 .flag
90 .store(true, std::sync::atomic::Ordering::SeqCst);
91
92 if let Ok(mut w) = self.state.waker.lock()
93 && let Some(waker) = w.take()
94 {
95 waker.wake();
96 }
97 }
98
99 pub fn notified(&self) -> NotifiedArc {
100 NotifiedArc {
101 state: self.state.clone(),
102 }
103 }
104}
105
106pub struct NotifiedArc {
107 state: Arc<StateArc>,
108}
109
110impl Future for NotifiedArc {
111 type Output = ();
112
113 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
114 if self
115 .state
116 .flag
117 .swap(false, std::sync::atomic::Ordering::SeqCst)
118 {
119 Poll::Ready(())
120 } else {
121 if let Ok(mut w) = self.state.waker.lock() {
122 *w = Some(cx.waker().clone());
123 }
124 Poll::Pending
125 }
126 }
127}