post-12 async/await

src/allocator.rs

@@ -53,7 +53,7 @@ pub fn init_heap(
             .ok_or(MapToError::FrameAllocationFailed)?;
 
         let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;
-        unsafe { mapper.map_to(page, frame, flags, frame_allocator)?.flush() };
+        mapper.map_to(page, frame, flags, frame_allocator)?.flush();
     }
 
     unsafe {

src/lib.rs

@@ -6,6 +6,7 @@
 #![feature(const_fn)]
 #![feature(const_in_array_repeat_expressions)]
 #![feature(custom_test_frameworks)]
+#![feature(wake_trait)]
 #![test_runner(crate::test_runner)]
 #![reexport_test_harness_main = "test_main"]
 

src/main.rs

@@ -10,7 +10,7 @@ use alloc::{boxed::Box, rc::Rc, vec, vec::Vec};
 use core::panic::PanicInfo;
 
 use blog_os::println;
-use blog_os::task::{keyboard, Task, simple_executor::SimpleExecutor};
+use blog_os::task::{keyboard, Task, executor::Executor};
 use bootloader::{entry_point, BootInfo};
 
 /// This function is called on panic.
@@ -66,7 +66,7 @@ fn kernel_main(boot_info: &'static BootInfo) -> ! {
         Rc::strong_count(&cloned_reference)
     );
 
-    let mut executor = SimpleExecutor::new();
+    let mut executor = Executor::new();
     executor.spawn(Task::new(example_task()));
     executor.spawn(Task::new(keyboard::print_keypresses()));
     executor.run();

src/task.rs

@@ -3,6 +3,7 @@ use core::sync::atomic::{AtomicU64, Ordering};
 use core::task::{Context, Poll};
 use alloc::boxed::Box;
 
+pub mod executor;
 pub mod keyboard;
 pub mod simple_executor;
 

src/task/executor.rs

@@ -0,0 +1,108 @@
+use super::{Task, TaskId};
+use alloc::{collections::{BTreeMap, VecDeque}, sync::Arc};
+use alloc::task::Wake;
+use core::task::{Waker, Poll, Context};
+use crossbeam_queue::ArrayQueue;
+
+pub struct Executor {
+    task_queue: VecDeque<Task>,
+    waiting_tasks: BTreeMap<TaskId, Task>,
+    wake_queue: Arc<ArrayQueue<TaskId>>,
+    waker_cache: BTreeMap<TaskId, Waker>,
+}
+
+impl Executor {
+    pub fn new() -> Self {
+        Executor {
+            task_queue: VecDeque::new(),
+            waiting_tasks: BTreeMap::new(),
+            wake_queue: Arc::new(ArrayQueue::new(100)),
+            waker_cache: BTreeMap::new(),
+        }
+    }
+
+    pub fn spawn(&mut self, task: Task) {
+        self.task_queue.push_back(task)
+    }
+
+    fn run_ready_tasks(&mut self) {
+        while let Some(mut task) = self.task_queue.pop_front() {
+            let task_id = task.id;
+            if !self.waker_cache.contains_key(&task_id) {
+                self.waker_cache.insert(task_id, self.create_waker(task_id));
+            }
+            let waker = self.waker_cache.get(&task_id).expect("should exist");
+            let mut context = Context::from_waker(waker);
+            match task.poll(&mut context) {
+                Poll::Ready(()) => {
+                    // task done -> remove cached waker
+                    self.waker_cache.remove(&task_id);
+                }
+                Poll::Pending => {
+                    if self.waiting_tasks.insert(task_id, task).is_some() {
+                        panic!("task with same ID already in waiting_tasks");
+                    }
+                }
+            }
+        }
+    }
+
+    fn create_waker(&self, task_id: TaskId) -> Waker {
+        Waker::from(Arc::new(TaskWaker {
+            task_id,
+            wake_queue: self.wake_queue.clone(),
+        }))
+    }
+
+    fn wake_tasks(&mut self) {
+        while let Ok(task_id) = self.wake_queue.pop() {
+            if let Some(task) = self.waiting_tasks.remove(&task_id) {
+                self.task_queue.push_back(task);
+            }
+        }
+    }
+
+    pub fn run(&mut self) -> ! {
+        loop {
+            self.wake_tasks();
+            self.run_ready_tasks();
+            self.sleep_if_idle();
+        }
+    }
+
+    fn sleep_if_idle(&self) {
+        use x86_64::instructions::interrupts::{self, enable_interrupts_and_hlt};
+
+        if !self.wake_queue.is_empty() {
+            return;
+        }
+
+        interrupts::disable();
+        if self.wake_queue.is_empty() {
+            enable_interrupts_and_hlt();
+        } else {
+            interrupts::enable();
+        }
+    }
+}
+
+struct TaskWaker {
+    task_id: TaskId,
+    wake_queue: Arc<ArrayQueue<TaskId>>,
+}
+
+impl TaskWaker {
+    fn wake_task(&self) {
+        self.wake_queue.push(self.task_id).expect("wake_queue full");
+    }
+}
+
+impl Wake for TaskWaker {
+    fn wake(self: Arc<Self>) {
+        self.wake_task();
+    }
+
+    fn wake_by_ref(self: &Arc<Self>) {
+        self.wake_task();
+    }
+}

src/vga_buffer.rs

@@ -6,8 +6,6 @@ use volatile::Volatile;
 #[cfg(test)]
 use crate::{serial_print, serial_println};
 
-use crate::println;
-
 #[allow(dead_code)]
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 #[repr(u8)]