ax_task/executor/
block_on.rs

1//! Scheduler-backed local future execution.
2
3use core::{
4    future::{Future, IntoFuture, poll_fn},
5    pin::{Pin, pin},
6    task::Poll,
7    time::Duration,
8};
9
10use super::LocalExecutor;
11use crate::{sync::WaitQueue, thread::current::current_thread_handle, time::MonotonicDeadline};
12
13/// Error returned when a runtime-local future misses its monotonic deadline.
14#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
15pub enum BlockOnError {
16    /// The requested relative timeout elapsed before the future completed.
17    #[error("future polling deadline elapsed")]
18    TimedOut,
19}
20
21/// Polls a future to completion on the calling runtime scheduler thread.
22///
23/// The executor is local to this call. Its waker may be invoked from another
24/// CPU or hard IRQ, while all future polling and destruction remain on the
25/// owner thread.
26#[track_caller]
27pub fn block_on<F: IntoFuture>(future: F) -> F::Output {
28    let thread = current_thread_handle()
29        .unwrap_or_else(|error| panic!("future polling requires a scheduler thread: {error}"));
30    let wait = WaitQueue::new();
31    let executor = LocalExecutor::new(thread.wake_handle())
32        .unwrap_or_else(|error| panic!("future executor requires its owner thread: {error}"));
33    let output = executor.run(future.into_future(), |condition| {
34        wait.wait_until(|| condition.should_abort());
35    });
36    drop(executor);
37    output
38}
39
40/// Polls a future until completion or a relative monotonic timeout.
41#[track_caller]
42pub fn block_on_timeout<F: IntoFuture>(
43    timeout: Duration,
44    future: F,
45) -> Result<F::Output, BlockOnError> {
46    let thread = current_thread_handle()
47        .unwrap_or_else(|error| panic!("future polling requires a scheduler thread: {error}"));
48    let wait = WaitQueue::new();
49    let executor = LocalExecutor::new(thread.wake_handle())
50        .unwrap_or_else(|error| panic!("future executor requires its owner thread: {error}"));
51    let deadline = crate::runtime::task_runtime::monotonic_now().deadline_after(timeout);
52    let mut future = pin!(future.into_future());
53    let timed = poll_fn(|context| poll_until_deadline(future.as_mut(), context, deadline));
54    let output = executor.run(timed, |condition| {
55        let _timed_out = wait.wait_until_deadline(deadline, || condition.should_abort());
56    });
57    drop(executor);
58    output
59}
60
61fn poll_until_deadline<F: Future>(
62    future: Pin<&mut F>,
63    context: &mut core::task::Context<'_>,
64    deadline: MonotonicDeadline,
65) -> Poll<Result<F::Output, BlockOnError>> {
66    if let Poll::Ready(output) = Future::poll(future, context) {
67        return Poll::Ready(Ok(output));
68    }
69    if crate::runtime::task_runtime::monotonic_now().reached(deadline) {
70        return Poll::Ready(Err(BlockOnError::TimedOut));
71    }
72    Poll::Pending
73}
ax_task/executor/block_on.rs

ax_task/executor/
block_on.rs
