1  use super::{Task, TaskId};
  2  use alloc::{collections::BTreeMap, sync::Arc};
  3  use core::task::Waker;
  4  use crossbeam_queue::ArrayQueue;
  5  use core::task::{Context, Poll};
  6  use alloc::task::Wake;
  7  
  8  pub struct Executor {
  9      tasks: BTreeMap<TaskId, Task>,
 10      task_queue: Arc<ArrayQueue<TaskId>>,
 11      waker_cache: BTreeMap<TaskId, Waker>,
 12  }
 13  
 14  impl Executor {
 15      pub fn new() -> Self {
 16          Executor {
 17              tasks: BTreeMap::new(),
 18              task_queue: Arc::new(ArrayQueue::new(100)),
 19              waker_cache: BTreeMap::new(),
 20          }
 21      }
 22  
 23      pub fn spawn(&mut self, task: Task) {
 24          let task_id = task.id;
 25          if self.tasks.insert(task.id, task).is_some() {
 26              panic!("task with some ID already in tasks");
 27          }
 28          self.task_queue.push(task_id).expect("queue full");
 29      }
 30  
 31      fn run_ready_tasks(&mut self) {
 32          let Self {
 33              tasks,
 34              task_queue,
 35              waker_cache,
 36          } = self;
 37  
 38          while let Some(task_id) = task_queue.pop() {
 39              let task = match tasks.get_mut(&task_id) {
 40                  Some(task) => task,
 41                  None => continue,
 42              };
 43              let waker = waker_cache
 44                  .entry(task_id)
 45                  .or_insert_with(|| TaskWaker::new(task_id, task_queue.clone()));
 46              let mut context = Context::from_waker(waker);
 47              match task.poll(&mut context) {
 48                  Poll::Ready(()) => {
 49                      tasks.remove(&task_id);
 50                      waker_cache.remove(&task_id);
 51                  }
 52                  Poll::Pending => {}
 53              }
 54          }
 55      }
 56  
 57      pub fn run(&mut self) -> ! {
 58          loop {
 59              self.run_ready_tasks();
 60              self.sleep_if_idle();
 61          }
 62      }
 63  
 64      fn sleep_if_idle(&self) { 
 65          use x86_64::instructions::interrupts::{self, enable_and_hlt};
 66  
 67          interrupts::disable();
 68          if self.task_queue.is_empty() {
 69              enable_and_hlt();
 70          } else {
 71              interrupts::enable();
 72          }
 73      }
 74  }
 75  
 76  struct TaskWaker {
 77      task_id: TaskId,
 78      task_queue: Arc<ArrayQueue<TaskId>>,
 79  }
 80  
 81  impl TaskWaker {
 82      fn wake_task(&self) {
 83          self.task_queue.push(self.task_id).expect("task_queue full");
 84      }
 85  
 86      fn new(task_id: TaskId, task_queue: Arc<ArrayQueue<TaskId>>) -> Waker {
 87          Waker::from(Arc::new(TaskWaker {
 88              task_id,
 89              task_queue,
 90          }))
 91      }
 92  }
 93  
 94  impl Wake for TaskWaker {
 95      fn wake(self: Arc<Self>) {
 96          self.wake_task();
 97      }
 98  
 99      fn wake_by_ref(self: &Arc<Self>) {
100          self.wake_task();
101      }
102  }