ax_task/future/
mod.rs

1//! Future support.
2
3use alloc::{sync::Arc, task::Wake};
4use core::{
5    future::poll_fn,
6    pin::pin,
7    task::{Context, Poll, Waker},
8};
9
10use crate::{
11    AxTaskRef, WeakAxTaskRef, current, current_run_queue, select_wake_run_queue,
12    sync::{PreemptIrqSaveState, SpinLock},
13};
14
15mod poll;
16pub use poll::*;
17
18pub(crate) mod time;
19pub use time::*;
20
21/// Errors owned by task waiting and notification operations.
22#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
23pub enum TaskError {
24    /// A signal or explicit task notification interrupted the wait.
25    #[error(transparent)]
26    Interrupted(#[from] Interrupted),
27    /// A task wait exceeded its deadline.
28    #[error(transparent)]
29    Elapsed(#[from] Elapsed),
30    /// A nonblocking task operation cannot currently make progress.
31    #[error("task operation would block")]
32    WouldBlock,
33    /// An IRQ operation used by a task-owned waker failed.
34    #[error(transparent)]
35    Irq(#[from] ax_hal::irq::IrqError),
36}
37
38/// A result returned by a task-domain operation.
39pub type TaskResult<T = ()> = Result<T, TaskError>;
40
41/// Error capability required by [`poll_io`].
42///
43/// The caller keeps ownership of its domain error while the task layer only
44/// asks how to recognize retryable I/O and how to publish an interruption.
45pub trait PollIoError {
46    /// Returns whether this error means the I/O operation should be retried.
47    fn is_would_block(&self) -> bool;
48
49    /// Creates the caller's domain error for an interrupted blocking wait.
50    fn interrupted(error: Interrupted) -> Self;
51}
52
53impl PollIoError for TaskError {
54    fn is_would_block(&self) -> bool {
55        matches!(self, Self::WouldBlock)
56    }
57
58    fn interrupted(error: Interrupted) -> Self {
59        error.into()
60    }
61}
62
63struct AxWaker {
64    task: WeakAxTaskRef,
65    woke: SpinLock<bool>,
66}
67
68impl AxWaker {
69    fn new(task: &AxTaskRef) -> Arc<Self> {
70        Arc::new(AxWaker {
71            task: Arc::downgrade(task),
72            woke: SpinLock::new(false),
73        })
74    }
75}
76
77impl Wake for AxWaker {
78    fn wake(self: Arc<Self>) {
79        self.wake_by_ref();
80    }
81
82    fn wake_by_ref(self: &Arc<Self>) {
83        if let Some(task) = self.task.upgrade() {
84            let mut rq = select_wake_run_queue::<PreemptIrqSaveState>(&task);
85            *self.woke.lock_irqsave() = true;
86            rq.unblock_task(task, true);
87        }
88    }
89}
90
91/// Blocks the current task until the given future is resolved or the task
92/// is interrupted by a signal.
93///
94/// When the task's `interrupted` flag is set (by `task.interrupt()`, typically
95/// from signal delivery), this function yields the CPU to allow signal
96/// processing on the return-to-userspace path. The future will be re-polled
97/// after the yield.
98#[track_caller]
99pub fn block_on<F: IntoFuture>(f: F) -> F::Output {
100    crate::api::might_sleep();
101
102    let mut fut = pin!(f.into_future());
103
104    let curr = current();
105    let task = curr.clone();
106
107    let axwaker = AxWaker::new(&task);
108    let waker = Waker::from(axwaker.clone());
109    let mut cx = Context::from_waker(&waker);
110
111    loop {
112        match fut.as_mut().poll(&mut cx) {
113            Poll::Pending => {
114                // Before sleeping, check if a signal has arrived. If so,
115                // yield instead of blocking so that the future's
116                // interruptible wrapper or poll_interrupt can observe
117                // the flag on the next poll. Use a non-consuming read
118                // to avoid stealing the flag from consumers that call
119                // poll_interrupt / take_interrupt themselves.
120                if task.interrupted() {
121                    crate::yield_now();
122                    continue;
123                }
124
125                let mut rq = current_run_queue::<PreemptIrqSaveState>();
126                let mut woke = axwaker.woke.lock_irqsave();
127                if !*woke {
128                    rq.future_blocked_resched(woke);
129                } else {
130                    *woke = false;
131                    drop(woke);
132                    drop(rq);
133                    crate::yield_now();
134                }
135            }
136            Poll::Ready(output) => break output,
137        }
138    }
139}
140
141/// Error returned by [`interruptible`].
142#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
143#[error("task wait was interrupted")]
144pub struct Interrupted;
145
146/// Makes a future interruptible.
147pub async fn interruptible<F: IntoFuture>(f: F) -> Result<F::Output, Interrupted> {
148    let mut f = pin!(f.into_future());
149    let curr = current();
150    poll_fn(|cx| {
151        if curr.poll_interrupt(cx).is_ready() {
152            return Poll::Ready(Err(Interrupted));
153        }
154        f.as_mut().poll(cx).map(Ok)
155    })
156    .await
157}
ax_task/future/mod.rs

ax_task/future/
mod.rs
