Skip to main content

shuttle_engine/future/
mod.rs

1pub mod batch_semaphore;
2
3use crate::runtime::execution::ExecutionState;
4use crate::runtime::thread;
5use std::future::Future;
6use std::pin::Pin;
7use std::task::{Context, Poll};
8
9/// Run a future to completion on the current thread.
10pub fn block_on<F: Future>(future: F) -> F::Output {
11    let mut future = Box::pin(future);
12    let waker = ExecutionState::with(|state| state.current_mut().waker());
13    let cx = &mut Context::from_waker(&waker);
14
15    loop {
16        match future.as_mut().poll(cx) {
17            Poll::Ready(result) => break result,
18            Poll::Pending => {
19                ExecutionState::with(|state| state.current_mut().sleep_unless_woken());
20                thread::switch();
21            }
22        }
23    }
24}
25
26/// Yields execution back to the scheduler.
27pub async fn yield_now() {
28    struct YieldNow {
29        yielded: bool,
30    }
31
32    impl Future for YieldNow {
33        type Output = ();
34
35        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
36            if self.yielded {
37                return Poll::Ready(());
38            }
39
40            self.yielded = true;
41            cx.waker().wake_by_ref();
42            ExecutionState::request_yield();
43            Poll::Pending
44        }
45    }
46
47    YieldNow { yielded: false }.await
48}