ember_os 0.1.3

A simple OS kernel implemented in rust, which has referenced https://os.phil-opp.com/
Documentation
use super::{Task, TaskId};
use alloc::task::Wake;
use alloc::{collections::BTreeMap, sync::Arc};
use core::task::{Context, Poll, Waker};
use crossbeam_queue::ArrayQueue;

pub struct Executor {
  tasks: BTreeMap<TaskId, Task>,
  task_queue: Arc<ArrayQueue<TaskId>>,
  waker_cache: BTreeMap<TaskId, Waker>,
}

impl Executor {
  pub fn new() -> Self {
    Executor {
      tasks: BTreeMap::new(),
      task_queue: Arc::new(ArrayQueue::new(100)),
      waker_cache: BTreeMap::new(),
    }
  }

  pub fn spawn(&mut self, task: Task) {
    let task_id = task.id;
    if self.tasks.insert(task.id, task).is_some() {
      panic!("task with same ID already in tasks!\n");
    }
    self.task_queue.push(task_id).expect("queue full!\n");
  }

  fn run_ready_tasks(&mut self) {
    // destructure `self` to avoid borrow checker errors
    let Self {
      tasks,
      task_queue,
      waker_cache,
    } = self;

    while let Some(task_id) = task_queue.pop() {
      let task = match tasks.get_mut(&task_id) {
        Some(task) => task,
        None => continue, // task no longer exists
      };
      let waker = waker_cache
        .entry(task_id)
        .or_insert_with(|| TaskWaker::new_waker(task_id, task_queue.clone()));
      let mut context = Context::from_waker(waker);
      match task.poll(&mut context) {
        Poll::Ready(()) => {
          // task done -> remove it and its cached waker
          tasks.remove(&task_id);
          waker_cache.remove(&task_id);
        }
        Poll::Pending => {}
      }
    }
  }
}

impl Default for Executor {
  fn default() -> Self {
    Self::new()
  }
}

impl Executor {
  pub fn run(&mut self) -> ! {
    loop {
      self.run_ready_tasks();
      self.sleep_if_idle();
    }
  }

  pub fn run_until_all_task_finished(&mut self) {
    while !self.task_queue.is_empty() {
      self.run_ready_tasks();
    }
  }

  fn sleep_if_idle(&self) {
    use x86_64::instructions::interrupts::{self, enable_and_hlt};

    // to avoid race conditions, disable all interruptions temporarily
    interrupts::disable();

    if self.task_queue.is_empty() {
      // enable interruptions again, hlt cpu
      enable_and_hlt();
    } else {
      // only enable interruptions
      interrupts::enable();
    }
  }
}

struct TaskWaker {
  task_id: TaskId,
  task_queue: Arc<ArrayQueue<TaskId>>,
}

impl TaskWaker {
  fn new_waker(task_id: TaskId, task_queue: Arc<ArrayQueue<TaskId>>) -> Waker {
    Waker::from(Arc::new(TaskWaker {
      task_id,
      task_queue,
    }))
  }

  fn wake_task(&self) {
    self
      .task_queue
      .push(self.task_id)
      .expect("task_queue full!\n");
  }
}

impl Wake for TaskWaker {
  fn wake(self: Arc<Self>) {
    self.wake_task();
  }

  fn wake_by_ref(self: &Arc<Self>) {
    self.wake_task();
  }
}