Skip to main content

pollster/
lib.rs

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