Add scheduler interrupts and locking

kernel/src/cpu/isr.rs

@@ -1,9 +1,9 @@
-use core::arch::global_asm;
+use core::{arch::global_asm, sync::atomic::Ordering};
 
 use log::warn;
 use spin::Mutex;
 
-use crate::sys::{lapic::send_eoi, scheduler::scheduler};
+use crate::sys::{lapic::send_eoi, scheduler::scheduler, sync::IN_ISR_HANDLER};
 
 global_asm!(include_str!("isr.s"), options(att_syntax));
 
@@ -85,8 +85,10 @@ extern "C" fn isr_handler(state: &mut ISRState) {
     if 48 <= state.isr && state.isr <= 254 {
         send_eoi(); // APIC interrupt
     }
+    IN_ISR_HANDLER.store(true, Ordering::Relaxed);
     if state.isr == 254 {
         scheduler(state);
+        IN_ISR_HANDLER.store(false, Ordering::Relaxed);
         return;
     }
     let handler;
@@ -95,7 +97,8 @@ extern "C" fn isr_handler(state: &mut ISRState) {
     }
     if let Some(handler) = handler {
         handler();
-        return;
+    } else {
+        warn!("Unhandled interrupt: {}", state.isr);
     }
-    warn!("Unhandled interrupt: {}", state.isr);
+    IN_ISR_HANDLER.store(false, Ordering::Relaxed);
 }

kernel/src/sys/lapic.rs

@@ -46,6 +46,12 @@ fn calibrate_timer() {
         address.add(REGISTER_TIMER_INITIAL_COUNT).write_volatile(0);
     }
 }
+pub fn schedule_timer_interrupt() {
+    let address = ADDRESS.load(Ordering::Relaxed);
+    unsafe {
+        address.add(REGISTER_TIMER_INITIAL_COUNT).write_volatile(10 * TICKS_PER_MS.load(Ordering::Relaxed) as u32);
+    }
+}
 pub fn setup_lapic_timer() {
     let address = ADDRESS.load(Ordering::Relaxed);
     unsafe {

kernel/src/sys/mod.rs

@@ -5,4 +5,5 @@ pub mod lapic;
 pub mod madt;
 pub mod pic;
 pub mod scheduler;
+pub mod sync;
 pub mod task;

kernel/src/sys/scheduler.rs

@@ -5,12 +5,17 @@ use spin::Mutex;
 
 use crate::cpu::isr::ISRState;
 
-use super::task::{switch_task, Task, TaskState, CURRENT_TASK};
+use super::{
+    lapic::schedule_timer_interrupt,
+    sync::{lock_kernel, unlock_kernel},
+    task::{switch_task, Task, TaskState, CURRENT_TASK},
+};
 
 static SCHEDULER_LIST: Mutex<VecDeque<Task>> = Mutex::new(VecDeque::new());
 pub static IDLE_TASK: Mutex<Option<Task>> = Mutex::new(None);
 
 pub fn scheduler(state: &mut ISRState) {
+    lock_kernel();
     let mut switch_to_task = None;
     {
         let mut scheduler_list = SCHEDULER_LIST.lock();
@@ -21,6 +26,7 @@ pub fn scheduler(state: &mut ISRState) {
     }
     if let Some(task) = switch_to_task {
         switch_task(state, task);
+        unlock_kernel();
         return;
     }
     let mut switch_idle = false;
@@ -34,9 +40,16 @@ pub fn scheduler(state: &mut ISRState) {
     if switch_idle {
         switch_task(state, IDLE_TASK.lock().take().unwrap());
     }
+    unlock_kernel();
 }
 pub fn schedule_task(task: Task) {
-    SCHEDULER_LIST.lock().push_back(task);
+    lock_kernel();
+    let mut scheduler_list = SCHEDULER_LIST.lock();
+    if scheduler_list.is_empty() {
+        schedule_timer_interrupt();
+    }
+    scheduler_list.push_back(task);
+    unlock_kernel();
 }
 pub fn yield_task() {
     unsafe {

kernel/src/sys/sync.rs

@@ -0,0 +1,19 @@
+use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+
+use kernel_common::instructions::{cli, sti};
+
+static KERNEL_LOCKED: AtomicUsize = AtomicUsize::new(0);
+pub static IN_ISR_HANDLER: AtomicBool = AtomicBool::new(false);
+
+pub fn lock_kernel() {
+    cli();
+    KERNEL_LOCKED.fetch_add(1, Ordering::Relaxed);
+}
+pub fn unlock_kernel() {
+    KERNEL_LOCKED.fetch_sub(1, Ordering::Relaxed);
+    if !IN_ISR_HANDLER.load(Ordering::Relaxed) && KERNEL_LOCKED.load(Ordering::Relaxed) == 0 {
+        unsafe {
+            sti();
+        }
+    }
+}

kernel/src/sys/task.rs

@@ -10,7 +10,11 @@ use crate::{
     sys::scheduler::{schedule_task, IDLE_TASK},
 };
 
-use super::scheduler::yield_task;
+use super::{
+    lapic::schedule_timer_interrupt,
+    scheduler::yield_task,
+    sync::{lock_kernel, unlock_kernel},
+};
 
 #[derive(PartialEq)]
 pub enum TaskState {
@@ -44,6 +48,7 @@ fn destroy_task(task: Task) {
     }
 }
 pub fn switch_task(current_state: &mut ISRState, new_task: Task) {
+    lock_kernel();
     let mut _current_task = CURRENT_TASK.lock();
     if let Some(current_task) = _current_task.as_mut() {
         current_task.state = *current_state;
@@ -55,6 +60,8 @@ pub fn switch_task(current_state: &mut ISRState, new_task: Task) {
     }
     *current_state = new_task.state;
     *_current_task = Some(new_task);
+    schedule_timer_interrupt();
+    unlock_kernel();
 }
 pub fn create_task(func: fn()) -> Task {
     let stack_address = allocate_stack();
@@ -90,15 +97,19 @@ pub fn create_task(func: fn()) -> Task {
     task
 }
 extern "C" fn task_entry() -> ! {
+    lock_kernel();
     let func;
     {
         let task = CURRENT_TASK.lock();
         func = task.as_ref().unwrap().initial_func;
     }
+    unlock_kernel();
     func();
+    lock_kernel();
     {
         CURRENT_TASK.lock().as_mut().unwrap().task_state = TaskState::TERMINATED;
     }
+    unlock_kernel();
     yield_task();
     panic!("Failed to terminate task");
 }