local_executor/
executor.rs

1use alloc::collections::VecDeque;
2use core::cell::RefCell;
3use core::pin::pin;
4use core::task::{Context, Poll, Waker};
5
6use crate::task::*;
7
8pub fn spawn<Ex, T, F>(handle: Ex, fut: F) -> Task<T>
9where
10    T: 'static,
11    F: 'static + Future<Output = T>,
12    Ex: ExecutorHandle,
13{
14    Executor::spawn(handle, fut)
15}
16
17pub async fn yield_now() {
18    let mut polled = false;
19    core::future::poll_fn(|cx| {
20        if polled {
21            Poll::Ready(())
22        } else {
23            polled = true;
24            cx.local_waker().wake_by_ref();
25            Poll::Pending
26        }
27    })
28    .await
29}
30
31pub struct Executor {
32    queue: RefCell<VecDeque<TaskRef>>,
33}
34
35impl Executor {
36    pub fn new() -> Self {
37        Self {
38            queue: RefCell::new(VecDeque::new()),
39        }
40    }
41
42    pub(crate) fn wake(&self, task: TaskRef) {
43        self.queue.borrow_mut().push_back(task);
44    }
45
46    pub fn block_on<T>(&self, fut: impl Future<Output = T>) -> T {
47        // let _guard = ExecutorHandle::enter(&self.0);
48        let Self { queue } = self;
49        let mut cx = Context::from_waker(Waker::noop());
50        let mut fut = pin!(fut);
51        loop {
52            if let Poll::Ready(output) = fut.as_mut().poll(&mut cx) {
53                return output;
54            }
55
56            // Newly waked tasks are deferred to the next loop.
57            let count = queue.borrow().len();
58            for _ in 0..count {
59                let task = queue.borrow_mut().pop_front().unwrap();
60                _ = task.poll_wakeable();
61            }
62        }
63    }
64
65    pub fn spawn<T, F, Ex>(handle: Ex, fut: F) -> Task<T>
66    where
67        T: 'static,
68        F: 'static + Future<Output = T>,
69        Ex: ExecutorHandle,
70    {
71        let ex = handle.get();
72        let task = Task::new(handle, fut);
73        ex.wake(task.inner());
74        task
75    }
76}
77
78impl Default for Executor {
79    fn default() -> Self {
80        Self::new()
81    }
82}
83
84pub trait ExecutorHandle: 'static + Unpin {
85    type Ref: core::ops::Deref<Target = Executor>;
86
87    fn get(&self) -> Self::Ref;
88}
89impl ExecutorHandle for alloc::rc::Weak<Executor> {
90    type Ref = alloc::rc::Rc<Executor>;
91    fn get(&self) -> Self::Ref {
92        self.upgrade().expect("not inside a valid executor")
93    }
94}