maniac_runtime/future/
block_on.rs

1use std::{
2    future::{Future, IntoFuture},
3    sync::Arc,
4    task::{Context, Poll, Wake, Waker},
5    thread,
6};
7
8thread_local! {
9    // A local reusable signal for each thread.
10    static LOCAL_THREAD_SIGNAL: Arc<Signal> = Arc::new(Signal {
11        owning_thread: thread::current(),
12    });
13}
14
15/// An extension trait that allows blocking on a future in suffix position.
16pub trait FutureExt: Future {
17    /// Block the thread until the future is ready.
18    ///
19    /// # Example
20    ///
21    /// ```
22    /// use pollster::FutureExt as _;
23    ///
24    /// let my_fut = async {};
25    ///
26    /// let result = my_fut.block_on();
27    /// ```
28    fn block_on(self) -> Self::Output
29    where
30        Self: Sized,
31    {
32        block_on(self)
33    }
34}
35
36impl<F: Future> FutureExt for F {}
37
38struct Signal {
39    /// The thread that owns the signal.
40    owning_thread: thread::Thread,
41}
42
43impl Wake for Signal {
44    fn wake(self: Arc<Self>) {
45        self.owning_thread.unpark();
46    }
47
48    fn wake_by_ref(self: &Arc<Self>) {
49        self.owning_thread.unpark();
50    }
51}
52
53/// Block the thread until the future is ready.
54///
55/// # Example
56///
57/// ```
58/// let my_fut = async {};
59/// let result = pollster::block_on(my_fut);
60/// ```
61pub fn block_on<F: IntoFuture>(fut: F) -> F::Output {
62    let mut fut = core::pin::pin!(fut.into_future());
63
64    // A signal used to wake up the thread for polling as the future moves to completion.
65    LOCAL_THREAD_SIGNAL.with(|signal| {
66        // Create a waker and a context to be passed to the future.
67        let waker = Waker::from(Arc::clone(signal));
68        let mut context = Context::from_waker(&waker);
69
70        // Poll the future to completion.
71        loop {
72            match fut.as_mut().poll(&mut context) {
73                Poll::Pending => thread::park(),
74                Poll::Ready(item) => break item,
75            }
76        }
77    })
78}