ableos/kernel/src/task.rs

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