diff --git a/src/task/mod.rs b/src/task/mod.rs index cfd1c1e9..ff9e32f1 100644 --- a/src/task/mod.rs +++ b/src/task/mod.rs @@ -19,4 +19,8 @@ impl Task { fn poll(&mut self, context: &mut Context) -> Poll<()> { self.future.as_mut().poll(context) } + + fn id(&self) -> usize { + &*self.future as *const _ as *const () as usize + } } diff --git a/src/task/simple_executor.rs b/src/task/simple_executor.rs index 346df3c7..4329304a 100644 --- a/src/task/simple_executor.rs +++ b/src/task/simple_executor.rs @@ -1,18 +1,23 @@ use super::Task; -use alloc::collections::VecDeque; -use core::{ - ptr, - task::{Context, Poll, RawWaker, RawWakerVTable, Waker}, -}; +use crate::println; +use alloc::collections::{BTreeMap, VecDeque}; +use conquer_once::spin::OnceCell; +use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; +use crossbeam_queue::ArrayQueue; + +static WAKE_QUEUE: OnceCell> = OnceCell::uninit(); pub struct SimpleExecutor { task_queue: VecDeque, + waiting_tasks: BTreeMap, } impl SimpleExecutor { pub fn new() -> SimpleExecutor { + WAKE_QUEUE.init_once(|| ArrayQueue::new(100)); SimpleExecutor { task_queue: VecDeque::new(), + waiting_tasks: BTreeMap::new(), } } @@ -21,29 +26,55 @@ impl SimpleExecutor { } pub fn run(&mut self) { + loop { + self.handle_wakeups(); + self.run_ready_tasks(); + } + } + + fn handle_wakeups(&mut self) { + while let Ok(task_id) = WAKE_QUEUE.get().unwrap().pop() { + if let Some(task) = self.waiting_tasks.remove(&task_id) { + self.task_queue.push_back(task); + } + } + } + + fn run_ready_tasks(&mut self) { while let Some(mut task) = self.task_queue.pop_front() { - let waker = waker(); + let waker = waker(task.id()); let mut context = Context::from_waker(&waker); match task.poll(&mut context) { Poll::Ready(()) => {} // task done - Poll::Pending => self.task_queue.push_back(task), + Poll::Pending => { + if self.waiting_tasks.insert(task.id(), task).is_some() { + panic!("Same task inserted into waiting_tasks twice"); + } + } } } } } -fn raw_waker() -> RawWaker { - fn no_op(_: *const ()) {} - fn clone(_: *const ()) -> RawWaker { - raw_waker() +fn raw_waker(task_id: usize) -> RawWaker { + fn clone(id: *const ()) -> RawWaker { + raw_waker(id as usize) } + fn wake(id: *const ()) { + if let Err(_) = WAKE_QUEUE.try_get().unwrap().push(id as usize) { + println!("WARNING: WAKE_QUEUE full; dropping wakeup") + } + } + + fn drop(_id: *const ()) {} + RawWaker::new( - ptr::null(), - &RawWakerVTable::new(clone, no_op, no_op, no_op), + task_id as *const (), + &RawWakerVTable::new(clone, wake, wake, drop), ) } -fn waker() -> Waker { - unsafe { Waker::from_raw(raw_waker()) } +fn waker(task_id: usize) -> Waker { + unsafe { Waker::from_raw(raw_waker(task_id)) } }