1
0
Fork 0
forked from AbleOS/ableos
ableos/kernel/src/task.rs

102 lines
2.3 KiB
Rust
Raw Normal View History

//! 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);
}
}