pollster/
lib.rs

1#![doc = include_str!("../README.md")]
2#![cfg_attr(docsrs, feature(doc_auto_cfg))]
3
4use std::{
5    future::{Future, IntoFuture},
6    sync::{Arc, Condvar, Mutex},
7    task::{Context, Poll, Wake, Waker},
8};
9
10#[cfg(feature = "macro")]
11pub use pollster_macro::{main, test};
12
13/// An extension trait that allows blocking on a future in suffix position.
14pub trait FutureExt: Future {
15    /// Block the thread until the future is ready.
16    ///
17    /// # Example
18    ///
19    /// ```
20    /// use pollster::FutureExt as _;
21    ///
22    /// let my_fut = async {};
23    ///
24    /// let result = my_fut.block_on();
25    /// ```
26    fn block_on(self) -> Self::Output where Self: Sized { block_on(self) }
27}
28
29impl<F: Future> FutureExt for F {}
30
31enum SignalState {
32    Empty,
33    Waiting,
34    Notified,
35}
36
37struct Signal {
38    state: Mutex<SignalState>,
39    cond: Condvar,
40}
41
42impl Signal {
43    fn new() -> Self {
44        Self {
45            state: Mutex::new(SignalState::Empty),
46            cond: Condvar::new(),
47        }
48    }
49
50    fn wait(&self) {
51        let mut state = self.state.lock().unwrap();
52        match *state {
53            // Notify() was called before we got here, consume it here without waiting and return immediately.
54            SignalState::Notified => *state = SignalState::Empty,
55            // This should not be possible because our signal is created within a function and never handed out to any
56            // other threads. If this is the case, we have a serious problem so we panic immediately to avoid anything
57            // more problematic happening.
58            SignalState::Waiting => {
59                unreachable!("Multiple threads waiting on the same signal: Open a bug report!");
60            }
61            SignalState::Empty => {
62                // Nothing has happened yet, and we're the only thread waiting (as should be the case!). Set the state
63                // accordingly and begin polling the condvar in a loop until it's no longer telling us to wait. The
64                // loop prevents incorrect spurious wakeups.
65                *state = SignalState::Waiting;
66                while let SignalState::Waiting = *state {
67                    state = self.cond.wait(state).unwrap();
68                }
69            }
70        }
71    }
72
73    fn notify(&self) {
74        let mut state = self.state.lock().unwrap();
75        match *state {
76            // The signal was already notified, no need to do anything because the thread will be waking up anyway
77            SignalState::Notified => {}
78            // The signal wasn't notified but a thread isn't waiting on it, so we can avoid doing unnecessary work by
79            // skipping the condvar and leaving behind a message telling the thread that a notification has already
80            // occurred should it come along in the future.
81            SignalState::Empty => *state = SignalState::Notified,
82            // The signal wasn't notified and there's a waiting thread. Reset the signal so it can be wait()'ed on again
83            // and wake up the thread. Because there should only be a single thread waiting, `notify_all` would also be
84            // valid.
85            SignalState::Waiting => {
86                *state = SignalState::Empty;
87                self.cond.notify_one();
88            }
89        }
90    }
91}
92
93impl Wake for Signal {
94    fn wake(self: Arc<Self>) {
95        self.notify();
96    }
97
98    fn wake_by_ref(self: &Arc<Self>) {
99        self.notify();
100    }
101}
102
103/// Block the thread until the future is ready.
104///
105/// # Example
106///
107/// ```
108/// let my_fut = async {};
109/// let result = pollster::block_on(my_fut);
110/// ```
111pub fn block_on<F: IntoFuture>(fut: F) -> F::Output {
112    let mut fut = core::pin::pin!(fut.into_future());
113
114    // Signal used to wake up the thread for polling as the future moves to completion. We need to use an `Arc`
115    // because, although the lifetime of `fut` is limited to this function, the underlying IO abstraction might keep
116    // the signal alive for far longer. `Arc` is a thread-safe way to allow this to happen.
117    // TODO: Investigate ways to reuse this `Arc<Signal>`... perhaps via a `static`?
118    let signal = Arc::new(Signal::new());
119
120    // Create a context that will be passed to the future.
121    let waker = Waker::from(Arc::clone(&signal));
122    let mut context = Context::from_waker(&waker);
123
124    // Poll the future to completion
125    loop {
126        match fut.as_mut().poll(&mut context) {
127            Poll::Pending => signal.wait(),
128            Poll::Ready(item) => break item,
129        }
130    }
131}
pollster/lib.rs

pollster/
lib.rs
