Improve dispatch to a real scheduler using preemption

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "lrnrtos"
-version = "0.4.4"
+version = "0.4.5"
 edition = "2024"
 
 [[bin]]

Documentation/kernel/data_structure.md

@@ -7,6 +7,8 @@
       - [Invariants](#invariants)
     - [AlignedStack16](#alignedstack16)
       - [Invariants](#invariants-1)
+    - [IndexedLinkedList](#indexedlinkedlist)
+      - [Invariants](#invariants-2)
 <!--toc:end-->
 
 ## Description
@@ -37,3 +39,19 @@ This type is used when you need a stack on a buffer, and the `sp` must be aligne
 - The backing storage is always 16-byte aligned.
 - Any stack pointer derived from this type must remain 16-byte aligned at all call boundaries.
 - This type provides a memory-layout guarantee only; it does not validate stack usage correctness.
+
+### IndexedLinkedList
+
+A linked list but store in an array, so each node is accessed from it's index in the array.
+This is better to use this data structure as a side storage, like we use it to store blocked task.
+Task are already store in a TaskList, so in the IndexedLinkedList used for blocked task we only store task id, and task awake tick for now.
+
+#### Invariants
+
+- All node should be accessible from the head node.
+- The list is sorted naturally from the `value` field of a node.
+- The `count` field should reflect the number of accessible node in the list.
+- The list is empty when `count`, `head` and `tail` are equal to 0.
+- If the `next_node` of a node is some, this `next_node` is valid.
+- If the `next_node` of a node is none, then this node is the `tail`.
+- The node `id` is unique, you can't add the same `id` in the list.

Documentation/kernel/scheduler.md

@@ -0,0 +1,57 @@
+# Kernel scheduler
+
+<!--toc:start-->
+- [Kernel scheduler](#kernel-scheduler)
+  - [Description](#description)
+  - [Queues](#queues)
+    - [Run queue](#run-queue)
+    - [Blocked queue](#blocked-queue)
+  - [Preemption](#preemption)
+  - [Scheduling model](#scheduling-model)
+  - [Invariants](#invariants)
+<!--toc:end-->
+
+## Description
+
+The kernel scheduler is in charge of managing which task to run, updating the run queue and the blocked queue and updating the task state.
+
+## Queues
+
+There's 2 queues used in the scheduler. The `run queue` and the `blocked queue`. Each queue is specific for a `CPU core`. The `CPU core 1` has a different `run queue` than the `CPU core 2`.
+
+### Run queue
+
+The `run queue` contains all the task ready to be run. In the `ready` task state.
+
+The `run queue` work using a `FIFO` queue, there's a queue for each `priority`, up to `32 queues`.
+To find which queue to use to execute a task, instead of iterating over all queues, we use a `bitmap`.
+There's only `32 priorities`, so we use a `u32 bitmap`, each bit representing a `run queue`, if the bit is set, there's at least 1 task to run.
+Else, the queue is empty.
+
+### Blocked queue
+
+The `blocked queue` contains all the task currently blocked. With different block reasons.
+But currently the `blocked queue` contains uniquely the task blocked using the `sleep` task primitive.
+
+The `blocked queue` work using an `indexed linked list`, the list is sorted from the shortest awake tick to the largest awake tick.
+The list is manage using `head` and `tail`, a bit like a `ring buffer` data structure.
+So when we need to check the next task to awake, we just check the `head` of the `blocked queue`.
+If it can be awake, it will update the `need_resched` flag, then it'll trigger a reschedule, the scheduler will be able to awake the task, and move it from the `blocked queue` to the `run queue`.
+
+## Preemption
+
+The scheduler is preemptive, meaning that if a task as a higher priority than the current task, it will run this higher priority task.
+The current task having a lowest priority, will be saved, and re-execute when there's no higher priority task to run.
+
+## Scheduling model
+
+The current scheduling model is a `cooperative priority based`. Meaning that if there's `2 task` with the same `priority`, if they don't call `yield`, one task will `always run`.
+Making the other task `starving`.
+So if you want to switch from a `task` to another one, you need to use `cooperative` functions, like `yield` or `sleep`.
+
+## Invariants
+
+- The scheduler need at least one task in the `run queue`, if the `run queue` is empty, it will try to run the idle task. Make sure that there's always at least one task in the `run queue`, or enable the `idle task` feature.
+- The scheduler assume that the `CpusState` is initialized to access the `CPU core scheduler state`.
+- The scheduler can be called from the `trap epilogue`, if so, a `trap frame`, should be available and accessible for the scheduler to run on.
+- For now the scheduler assume that there's always a `current running task`, if not, and the scheduler is triggered, this could lead to UB.

linkers/linker_test_mode.ld

@@ -1,9 +1,9 @@
 ENTRY(kstart)
 
 MEMORY {
-  RAM (rwx) : ORIGIN = 0x80200000, LENGTH = 128K
+  RAM (rwx) : ORIGIN = 0x80200000, LENGTH = 256K
   /* Not real ROM or flash from qemu virt machine, just use another RAM reg for now */
-  ROM (rx) : ORIGIN = 0x80000000, LENGTH = 210K
+  ROM (rx) : ORIGIN = 0x80000000, LENGTH = 256K
 }
 
 SECTIONS {

src/arch/riscv32/asm/context_offset.S

@@ -11,7 +11,6 @@
 .set OFFSET_SP, 140
 # Return address
 .set OFFSET_RA, 144
-.set OFFSET_MSTATUS, 148
 # Optionnal flags, will be use later maybe
-.set OFFSET_FLAGS, 152
-.set OFFSET_INSTRUCTION_REG, 155
+.set OFFSET_FLAGS, 148
+.set OFFSET_INSTRUCTION_REG, 151

src/arch/riscv32/asm/restore_context.S

@@ -12,14 +12,38 @@ restore_context:
   # Restore sp from current task context structure
   lw t0, OFFSET_SP(t6)
   mv sp, t0
-  ## Update mstatus
-  lw t0, OFFSET_MSTATUS(t6)
-  csrw mstatus, t0
   # Restore current task context using GNU macro
   # GNU macros from `src/arch/riscv32/asm/gnu_macro.S`
   .set	i, 1
   .rept	31
     load_gp_context %i
     .set	i, i+1
   .endr
+  # Update mstatus
+  li   t0, ((3 << 11) | (1 << 3))
+  csrrs x0, mstatus, t0
   ret
+
+.global trap_restore_context
+.type trap_restore_context, @function
+trap_restore_context:
+  # Move current task context struct save in caller in a0 reg to t6
+  mv t6, a0
+  # Update sp
+  # Restore sp from current task context structure
+  lw t0, OFFSET_SP(t6)
+  mv sp, t0
+  # Update mepc
+  lw t0, OFFSET_PC(t6)
+  csrw mepc, t0
+  # Restore current task context using GNU macro
+  # GNU macros from `src/arch/riscv32/asm/gnu_macro.S`
+  .set	i, 1
+  .rept	31
+    load_gp_context %i
+    .set	i, i+1
+  .endr
+  # Update mstatus
+  li   t0, ((3 << 11) | (1 << 7))
+  csrrs x0, mstatus, t0
+  mret

src/arch/riscv32/asm/save_context.S

@@ -3,15 +3,37 @@
 .global save_context
 .type save_context, @function
 save_context:
+  # Clear mstatus.MPP and mstatus.mie to avoid interruption while saving and switching context
+  li   t0, ((3 << 11) | (1 << 3))
+  csrrc x0, mstatus, t0
   # Move current task context struct save in caller in a0 reg to t6
   mv t6, a0
   # Store ra in task context structure
   sw a1, OFFSET_RA(t6)
   # Store word from t0(sp) in context structure
   sw a2, OFFSET_SP(t6)
-  # Save mstatus
-  csrr t0, mstatus
-  sw t0, OFFSET_MSTATUS(t6)
+  sw a3, OFFSET_PC(t6)
+  # Save current task context using GNU macro
+  # GNU macros from `src/arch/riscv32/asm/gnu_macro.S`
+  .set	i, 1
+  .rept	31
+    save_gp_context %i
+    .set	i, i+1
+  .endr
+  ret
+
+.global trap_save_context
+.type trap_save_context, @function
+trap_save_context:
+  # Clear mstatus.MPP and mstatus.mie to avoid interruption while saving and switching context
+  li   t0, ((3 << 11) | (1 << 3))
+  csrrc x0, mstatus, t0
+   # Move current task context struct save in caller in a0 reg to t6
+  mv t6, a0
+  # Store ra in task context structure
+  sw a1, OFFSET_PC(t6)
+  # Store word from t0(sp) in context structure
+  sw a2, OFFSET_SP(t6)
   # Save current task context using GNU macro
   # GNU macros from `src/arch/riscv32/asm/gnu_macro.S`
   .set	i, 1

src/arch/riscv32/asm/trap_entry.S

@@ -11,22 +11,19 @@
 .global trap_entry
 .type trap_entry, @function
 trap_entry:
-  # Read mscratch into t6
+  # Save task context
+  # Get current task ptr
+  la t0, TASK_HANDLER # Address in RAM
+  lw t1, 0(t0) # Get the value behind the ref
+  mv a0, t1
+  # Save current pc
+  csrr a1, mepc
+  # Save current sp
+  mv a2, sp
+  # Call the save context function
+  call trap_save_context
+  # Read mscratch into t7
   csrr	t6, mscratch
-  # Save all GP registers using GNU macro
-  # Just a loop 
-  .set 	i, 1
-  .rept	31
-    save_gp %i, t6
-    .set	i, i+1
-  .endr
-  # Read CSR and store word in correct offset of trap frame structure
-  csrr    t0, satp
-  sw      t0, OFFSET_SATP(t6)       # store satp
-  csrr    t0, mhartid
-  sw      t0, OFFSET_HARTID(t6)     # store hartid (optional, redundant)
-  sw      sp, (OFFSET_TRAP_STACK - 4)(t6)   # store original sp just before trap_stack slot
-  # Load content of trap stack offset into t1
   lw      t1, OFFSET_TRAP_STACK(t6)   # t1 = trap_stack pointer
   # Branch instruction to check if t1 = 0
   beqz    t1, trap_no_trapstack
@@ -42,22 +39,18 @@ trap_entry:
   mv a5, t6
   # Call trap_handler rust function
   call trap_handler 
-  # Load all register back to previous state
-  csrr	t6, mscratch
-  # Restore all GP registers
-  .set	i, 1
-  .rept	31
-    load_gp %i
-    .set	i, i+1
-  .endr
+
   # Check if a re-schedule is needed or not.
   call read_need_reschedule
   # If a0 != 0, goto 1f, else mret
   bnez a0, 1f
-  mret
+  # If there's no need to a reschedule, restore the task context
+  la t0, TASK_HANDLER # Address in RAM
+  lw t1, 0(t0) # Get the value behind the ref
+  mv a0, t1
+  call trap_restore_context
 1:
-  call dispatch
-  mret
+  call scheduler
 
 # Function used when the trap strack = 0, just infinite loop for debuging purpose 
 trap_no_trapstack:

src/arch/riscv32/asm/yield.S

@@ -12,4 +12,4 @@ yield:
   # Call the save context function
   call save_context
   # Once save_context return, call the scheduler
-  call dispatch
+  call scheduler

src/arch/riscv32/task/task_context.rs

@@ -25,9 +25,8 @@ pub struct TaskContext {
     pub pc: u32,                  // Offset 136
     pub sp: u32,                  // Offset 140
     pub ra: u32,                  // Offset 144
-    pub mstatus: u32,             // Offset 148
-    pub flags: [u8; 3],           // Offset 152 (first index 152; second index 153, third index 154)
-    pub instruction_register: u8, // Offset 155
+    pub flags: [u8; 3],           // Offset 148 (first index 148; second index 149, third index 150)
+    pub instruction_register: u8, // Offset 151
 }
 
 impl TaskContext {
@@ -39,7 +38,6 @@ impl TaskContext {
             sp: size[0] as u32,
             ra: func as usize as u32,
             // Set mstatus to 8 by default to enable mie
-            mstatus: 8,
             flags: [0u8; 3],
             instruction_register: 0,
         }

src/config.rs

@@ -17,7 +17,10 @@ pub static LOG_LEVEL: LogLevel = LogLevel::Debug;
 
 // Define the uart address to use in kprint
 pub static KPRINT_ADDRESS: usize = 0x1000_0000;
-
+// ————————————————————————————————————————————————————————————
+// ——————— Define the max priority available for a task ———————
+// ————————————————————————————————————————————————————————————
+pub static TASK_MAX_PRIORITY: usize = 32;
 // ————————————————————————————————————————————————————————————
 // ———————— Define the max size of devices sub-systems ————————
 // ————————————————————————————————————————————————————————————
@@ -36,10 +39,11 @@ pub static FDT_MAX_PROPS: usize = 128;
 // ————————————————————————————————————————————————————————————
 pub static TASK_LIST_MAX_SIZE: usize = 4;
 // ————————————————————————————————————————————————————————————
-// ————————————— Define the max size of the Run queue —————————
+// ———— Define the max size of the task run/blocked queue —————
 // ————————————————————————————————————————————————————————————
 // The run queue is len - 1, if the size is 4, it will only use 3 slot in the queue.
 pub static RUN_QUEUE_MAX_SIZE: usize = 3;
+pub static BLOCK_QUEUE_MAX_SIZE: usize = 3;
 // ————————————————————————————————————————————————————————————
 // ————————————— Define the number of CPU core ————————————————
 // ————————————————————————————————————————————————————————————

src/info.rs

@@ -1,4 +1,4 @@
 // The kernel exposes build-time version information at runtime.
 // This information is immutable and reflects the exact binary currently running.
 // It is intended for debugging, logging, and diagnostic purposes.
-pub static KERNEL_VERSION: &str = "0.4.4";
+pub static KERNEL_VERSION: &str = "0.4.5";

src/main.rs

@@ -6,6 +6,7 @@
 #![deny(clippy::expect_used)]
 #![deny(clippy::todo)]
 #![deny(clippy::unimplemented)]
+#![feature(stmt_expr_attributes)]
 
 // Config module
 pub mod config;

src/primitives/bitmap.rs

@@ -0,0 +1,44 @@
+#[derive(Copy, Clone)]
+pub struct Bitmap {
+    pub map: u32,
+}
+
+impl Bitmap {
+    #[allow(clippy::new_without_default)]
+    pub const fn new() -> Self {
+        Bitmap { map: 0 }
+    }
+
+    /// Set the given bit to 1.
+    pub fn set_bit(&mut self, bit: usize) {
+        let mask = 1 << bit;
+        self.map |= mask;
+    }
+
+    /// Clear the given bit. Set it to 0.
+    pub fn clear_bit(&mut self, bit: usize) {
+        self.map &= !(1 << bit);
+    }
+
+    /// Iterate over the bitmap and return the heavier bit.
+    /// Example: bitmap set to u8 -> map: 01001010. The function will return 6, because the first
+    /// bit set to 1, from the highest bit, is the bit 6.
+    /// Arguments:
+    /// &mut self: call the map initialized. Must be mutable.
+    pub fn find_leading_bit(&mut self) -> usize {
+        let bits = core::mem::size_of::<usize>() * 8;
+        let mut value: usize = 0;
+        for i in (0..bits).rev() {
+            let bit = (self.map >> i) & 1;
+            if bit == 1 {
+                value = i;
+                break;
+            }
+        }
+        value
+    }
+
+    pub fn is_bitmap_zero(&self) -> bool {
+        self.map == 0
+    }
+}