freya_core/
notify.rs

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}