forked from AbleOS/ableos
102 lines
2.3 KiB
Rust
102 lines
2.3 KiB
Rust
//! Async task and executor
|
|
|
|
use alloc::{boxed::Box, collections::BTreeMap, sync::Arc, task::Wake};
|
|
use core::{
|
|
future::Future,
|
|
pin::Pin,
|
|
task::{Context, Poll, Waker},
|
|
};
|
|
use crossbeam_queue::SegQueue;
|
|
use slab::Slab;
|
|
|
|
type TaskQueue = Arc<SegQueue<TaskId>>;
|
|
|
|
/// Tasks executor
|
|
#[derive(Default)]
|
|
pub struct Executor {
|
|
/// All spawned tasks' stash
|
|
tasks: Slab<Task>,
|
|
|
|
/// Awake tasks' queue
|
|
queue: TaskQueue,
|
|
|
|
/// Wakers
|
|
wakers: BTreeMap<TaskId, Waker>,
|
|
}
|
|
|
|
impl Executor {
|
|
/// Spawn a future
|
|
pub fn spawn(&mut self, future: impl Future<Output = ()> + 'static) {
|
|
self.queue
|
|
.push(TaskId(self.tasks.insert(Task::new(future))));
|
|
}
|
|
|
|
/// Run tasks
|
|
pub fn run(&mut self) -> ! {
|
|
loop {
|
|
while let Some(id) = self.queue.pop() {
|
|
let task = match self.tasks.get_mut(id.0) {
|
|
Some(t) => t,
|
|
None => continue,
|
|
};
|
|
|
|
let mut cx = Context::from_waker(
|
|
self.wakers
|
|
.entry(id)
|
|
.or_insert_with(|| TaskWaker::new(id, Arc::clone(&self.queue))),
|
|
);
|
|
|
|
match task.poll(&mut cx) {
|
|
Poll::Ready(()) => {
|
|
self.tasks.remove(id.0);
|
|
self.wakers.remove(&id);
|
|
}
|
|
Poll::Pending => (),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
|
struct TaskId(usize);
|
|
|
|
/// Async task
|
|
struct Task {
|
|
future: Pin<Box<dyn Future<Output = ()>>>,
|
|
}
|
|
|
|
impl Task {
|
|
/// Create a new task from a future
|
|
fn new(future: impl Future<Output = ()> + 'static) -> Self {
|
|
Self {
|
|
future: Box::pin(future),
|
|
}
|
|
}
|
|
|
|
fn poll(&mut self, cx: &mut Context) -> Poll<()> {
|
|
self.future.as_mut().poll(cx)
|
|
}
|
|
}
|
|
|
|
struct TaskWaker {
|
|
id: TaskId,
|
|
queue: TaskQueue,
|
|
}
|
|
|
|
impl TaskWaker {
|
|
fn new(id: TaskId, queue: TaskQueue) -> Waker {
|
|
Waker::from(Arc::new(Self { id, queue }))
|
|
}
|
|
}
|
|
|
|
impl Wake for TaskWaker {
|
|
fn wake(self: Arc<Self>) {
|
|
self.wake_by_ref();
|
|
}
|
|
|
|
fn wake_by_ref(self: &Arc<Self>) {
|
|
self.queue.push(self.id);
|
|
}
|
|
}
|