diff --git a/src/allocator.rs b/src/allocator.rs index 858c79c..1cb5a8b 100644 --- a/src/allocator.rs +++ b/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 { diff --git a/src/lib.rs b/src/lib.rs index f93e5bf..a612f18 100644 --- a/src/lib.rs +++ b/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"] diff --git a/src/main.rs b/src/main.rs index a602d6f..e83fbf4 100644 --- a/src/main.rs +++ b/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(); diff --git a/src/task.rs b/src/task.rs index 6adcd50..fba58be 100644 --- a/src/task.rs +++ b/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; diff --git a/src/task/executor.rs b/src/task/executor.rs new file mode 100644 index 0000000..dac026b --- /dev/null +++ b/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, + waiting_tasks: BTreeMap, + wake_queue: Arc>, + waker_cache: BTreeMap, +} + +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>, +} + +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.wake_task(); + } + + fn wake_by_ref(self: &Arc) { + self.wake_task(); + } +} \ No newline at end of file diff --git a/src/vga_buffer.rs b/src/vga_buffer.rs index 84cae0f..7512794 100644 --- a/src/vga_buffer.rs +++ b/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)]