nio_future/
lib.rs

1use std::{
2    future::Future,
3    mem,
4    pin::{Pin, pin},
5    sync::{Arc, Condvar, Mutex},
6    task::{Context, Poll, Wake, Waker},
7};
8
9pub fn yield_now() -> impl Future<Output = ()> {
10    struct YieldNow {
11        yielded: bool,
12    }
13    impl Future for YieldNow {
14        type Output = ();
15        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
16            if self.yielded {
17                return Poll::Ready(());
18            }
19            self.yielded = true;
20            cx.waker().wake_by_ref();
21            Poll::Pending
22        }
23    }
24    YieldNow { yielded: false }
25}
26
27pub fn block_on<Fut>(fut: Fut) -> Fut::Output
28where
29    Fut: Future,
30{
31    let mut fut = pin!(fut);
32
33    let signal = Arc::new(Signal::default());
34    let waker = Waker::from(signal.clone());
35    let mut cx = Context::from_waker(&waker);
36    loop {
37        match fut.as_mut().poll(&mut cx) {
38            Poll::Ready(val) => return val,
39            Poll::Pending => signal.wait_for_wakeup(),
40        }
41    }
42}
43
44const RUNNING: u8 = 0;
45const NOTIFIED: u8 = 1;
46const SLEEP: u8 = 2;
47
48#[derive(Default)]
49struct Signal {
50    state: Mutex<u8>,
51    signal: Condvar,
52}
53
54impl Signal {
55    /// State transitions:
56    ///
57    /// ```text
58    /// RUNNING -> (SLEEP? -> NOTIFIED -> RUNNING)* -> Complete?
59    /// ```
60    fn wait_for_wakeup(&self) {
61        let mut state = self.state.lock().unwrap();
62        if *state == NOTIFIED {
63            *state = RUNNING;
64        } else {
65            *state = SLEEP;
66            'spurious_wakeups: loop {
67                state = self.signal.wait(state).unwrap();
68                if *state == NOTIFIED {
69                    *state = RUNNING;
70                    break 'spurious_wakeups;
71                }
72            }
73        }
74    }
75}
76
77impl Wake for Signal {
78    fn wake_by_ref(self: &Arc<Self>) {
79        let state = {
80            let mut state = self.state.lock().unwrap();
81            mem::replace(&mut *state, NOTIFIED)
82        };
83        if state == SLEEP {
84            self.signal.notify_one();
85        }
86    }
87
88    fn wake(self: Arc<Self>) {
89        self.wake_by_ref();
90    }
91}
nio_future/lib.rs

nio_future/
lib.rs
