Use cooked-waker crate for creating wakers

Cargo.lock

@@ -24,6 +24,7 @@ version = "0.1.0"
 dependencies = [
  "bootloader",
  "conquer-once",
+ "cooked-waker",
  "crossbeam-queue",
  "futures-util",
  "lazy_static",
@@ -61,6 +62,27 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "654fb2472cc369d311c547103a1fa81d467bef370ae7a0680f65939895b1182a"
 
+[[package]]
+name = "cooked-waker"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b0032e8be0680b63daf14b0fd7fd57f85fd6897951e110b041d32e839a831914"
+dependencies = [
+ "cooked-waker-derive",
+ "stowaway",
+]
+
+[[package]]
+name = "cooked-waker-derive"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "74ded02b04a8ff53ea7328cd71297cde598bee70d165ad92512bdb7e20be9d74"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "cpuio"
 version = "0.2.0"
@@ -158,6 +180,24 @@ version = "0.1.0-alpha.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5894c618ce612a3fa23881b152b608bafb8c56cfc22f434a3ba3120b40f7b587"
 
+[[package]]
+name = "proc-macro2"
+version = "1.0.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c09721c6781493a2a492a96b5a5bf19b65917fe6728884e7c44dd0c60ca3435"
+dependencies = [
+ "unicode-xid",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2bdc6c187c65bca4260c9011c9e3132efe4909da44726bad24cf7572ae338d7f"
+dependencies = [
+ "proc-macro2",
+]
+
 [[package]]
 name = "scopeguard"
 version = "1.1.0"
@@ -179,6 +219,23 @@ dependencies = [
  "lock_api",
 ]
 
+[[package]]
+name = "stowaway"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32b13d61e782863c0eaeeaf7e8e580dc956a625851fd2fb73e5e92db9601c891"
+
+[[package]]
+name = "syn"
+version = "1.0.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0df0eb663f387145cab623dea85b09c2c5b4b0aef44e945d928e682fce71bb03"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-xid",
+]
+
 [[package]]
 name = "uart_16550"
 version = "0.2.4"
@@ -189,6 +246,12 @@ dependencies = [
  "x86_64",
 ]
 
+[[package]]
+name = "unicode-xid"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "826e7639553986605ec5979c7dd957c7895e93eabed50ab2ffa7f6128a75097c"
+
 [[package]]
 name = "volatile"
 version = "0.2.6"

Cargo.toml

@@ -21,6 +21,7 @@ uart_16550 = "0.2.0"
 pic8259_simple = "0.1.1"
 pc-keyboard = "0.5.0"
 linked_list_allocator = "0.8.0"
+cooked-waker = "1.0.2"
 
 [dependencies.lazy_static]
 version = "1.0"

src/task/simple_executor.rs

@@ -1,23 +1,24 @@
 use super::Task;
-use crate::println;
-use alloc::collections::{BTreeMap, VecDeque};
-use conquer_once::spin::OnceCell;
-use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
+use alloc::{
+    collections::{BTreeMap, VecDeque},
+    sync::Arc,
+};
+use cooked_waker::IntoWaker;
+use core::task::{Context, Poll};
 use crossbeam_queue::ArrayQueue;
 
-static WAKE_QUEUE: OnceCell<ArrayQueue<usize>> = OnceCell::uninit();
-
 pub struct SimpleExecutor {
     task_queue: VecDeque<Task>,
     waiting_tasks: BTreeMap<usize, Task>,
+    wake_queue: Arc<ArrayQueue<usize>>,
 }
 
 impl SimpleExecutor {
     pub fn new() -> SimpleExecutor {
-        WAKE_QUEUE.init_once(|| ArrayQueue::new(100));
         SimpleExecutor {
             task_queue: VecDeque::new(),
             waiting_tasks: BTreeMap::new(),
+            wake_queue: Arc::new(ArrayQueue::new(100)),
         }
     }
 
@@ -33,7 +34,7 @@ impl SimpleExecutor {
     }
 
     fn handle_wakeups(&mut self) {
-        while let Ok(task_id) = WAKE_QUEUE.get().unwrap().pop() {
+        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);
             }
@@ -42,7 +43,11 @@ impl SimpleExecutor {
 
     fn run_ready_tasks(&mut self) {
         while let Some(mut task) = self.task_queue.pop_front() {
-            let waker = waker(task.id());
+            let waker = TaskWaker {
+                task_id: task.id(),
+                wake_queue: self.wake_queue.clone(),
+            }
+            .into_waker();
             let mut context = Context::from_waker(&waker);
             match task.poll(&mut context) {
                 Poll::Ready(()) => {} // task done
@@ -56,25 +61,26 @@ impl SimpleExecutor {
     }
 }
 
-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(
-        task_id as *const (),
-        &RawWakerVTable::new(clone, wake, wake, drop),
-    )
+#[derive(Debug, Clone, IntoWaker)]
+struct TaskWaker {
+    task_id: usize,
+    wake_queue: Arc<ArrayQueue<usize>>,
 }
 
-fn waker(task_id: usize) -> Waker {
-    unsafe { Waker::from_raw(raw_waker(task_id)) }
+impl TaskWaker {
+    fn wake_task(&self) {
+        self.wake_queue.push(self.task_id).expect("wake queue full");
+    }
+}
+
+impl cooked_waker::WakeRef for TaskWaker {
+    fn wake_by_ref(&self) {
+        self.wake_task();
+    }
+}
+
+impl cooked_waker::Wake for TaskWaker {
+    fn wake(self) {
+        self.wake_task();
+    }
 }