task_stream/
lib.rs

1#![no_std]
2
3mod local;
4
5use ach_mpmc::Mpmc;
6pub use async_task::Runnable;
7use core::future::Future;
8use core::pin::Pin;
9use core::task::{Context, Poll};
10use futures_util::task::AtomicWaker;
11use futures_util::Stream;
12pub use local::*;
13
14const TASK_LEN: usize = 64;
15static EXEC: Executor<TASK_LEN> = Executor::new();
16
17pub struct Executor<const N: usize> {
18    queue: Mpmc<Runnable, N>,
19    waker: AtomicWaker,
20}
21impl<const N: usize> Executor<N> {
22    pub const fn new() -> Self {
23        Self {
24            queue: Mpmc::new(),
25            waker: AtomicWaker::new(),
26        }
27    }
28    pub const fn into_local(self) -> LocalExecutor<N> {
29        LocalExecutor::new(self)
30    }
31
32    /// task is waked
33    fn schedule(&self, runnable: Runnable) {
34        self.queue.push(runnable).unwrap();
35        self.waker.wake();
36    }
37
38    pub fn spawn<F>(&'static self, future: F)
39    where
40        F: Future + Send + 'static,
41        F::Output: Send + 'static,
42    {
43        let (runnable, task) = async_task::spawn(future, move |r| self.schedule(r));
44        runnable.schedule();
45        task.detach();
46    }
47
48    /// # Safety
49    ///
50    /// - TaskStream must be used  on the original thread.
51    unsafe fn spawn_local<F>(&self, future: F)
52    where
53        F: Future + 'static,
54        F::Output: 'static,
55    {
56        let (runnable, task) = async_task::spawn_unchecked(future, move |r| self.schedule(r));
57        runnable.schedule();
58        task.detach();
59    }
60
61    pub fn stream(&self) -> TaskStream<N> {
62        TaskStream { exec: self }
63    }
64}
65
66pub fn spawn<F>(future: F)
67where
68    F: Future + Send + 'static,
69    F::Output: Send + 'static,
70{
71    EXEC.spawn(future)
72}
73
74pub fn stream() -> TaskStream<'static, TASK_LEN> {
75    EXEC.stream()
76}
77
78pub struct TaskStream<'a, const N: usize> {
79    exec: &'a Executor<N>,
80}
81impl<'a, const N: usize> TaskStream<'a, N> {
82    pub fn get_task(&self) -> Option<Runnable> {
83        self.exec.queue.pop().ok()
84    }
85}
86impl<'a, const N: usize> Stream for TaskStream<'a, N> {
87    type Item = Runnable;
88    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
89        if let Some(task) = self.get_task() {
90            Poll::Ready(Some(task))
91        } else {
92            self.exec.waker.register(cx.waker());
93            Poll::Pending
94        }
95    }
96}
task_stream/lib.rs

task_stream/
lib.rs
