From 6dd0edd49207ea92d5dbafbaebf81572175f50bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?D=C3=A1niel=20Buga?= <bugadani@gmail.com>
Date: Fri, 5 Sep 2025 17:01:17 +0200
Subject: [PATCH] Teeny-tiny scheduler improvements (#4061)

* Avoid swapping to the same task

* Allow scheduling multiple tasks to be deleted

* Add task state
---
 esp-preempt/src/lib.rs      | 74 +++++++++++++++++++++++++------------
 esp-preempt/src/task/mod.rs | 17 +++++++++
 2 files changed, 68 insertions(+), 23 deletions(-)

diff --git a/esp-preempt/src/lib.rs b/esp-preempt/src/lib.rs
index 9e9d11567..310cc1ebb 100644
--- a/esp-preempt/src/lib.rs
+++ b/esp-preempt/src/lib.rs
@@ -160,48 +160,76 @@ impl SchedulerState {
         }
     }
 
-    #[cfg(xtensa)]
-    fn switch_task(&mut self, trap_frame: &mut esp_hal::trapframe::TrapFrame) {
-        task::save_task_context(unsafe { &mut *self.current_task }, trap_frame);
-
-        if !self.to_delete.is_null() {
+    fn delete_marked_tasks(&mut self) {
+        while !self.to_delete.is_null() {
             let task_to_delete = core::mem::take(&mut self.to_delete);
+            self.to_delete = unsafe { (*task_to_delete).next_to_delete };
             self.delete_task(task_to_delete);
         }
+    }
 
-        unsafe { self.current_task = (*self.current_task).next };
+    fn select_next_task(&mut self) -> Option<*mut Context> {
+        let mut current = self.current_task;
+        loop {
+            let next_task = unsafe { (*current).next };
+
+            if next_task == self.current_task {
+                // We didn't find a new task to switch to.
+                // TODO: mark the current task as Running
+                // Once we have actual task states, yield should marked the current task as Ready,
+                // other stuff as Waiting.
+                return None;
+            }
+
+            if unsafe { (*next_task).state }.is_ready() {
+                // TODO: mark the selected task as Running
+                return Some(next_task);
+            }
+            current = next_task;
+        }
+    }
+
+    #[cfg(xtensa)]
+    fn switch_task(&mut self, trap_frame: &mut esp_hal::trapframe::TrapFrame) {
+        self.delete_marked_tasks();
+
+        let Some(next_task) = self.select_next_task() else {
+            return;
+        };
+
+        task::save_task_context(unsafe { &mut *self.current_task }, trap_frame);
+
+        self.current_task = next_task;
 
         task::restore_task_context(unsafe { &mut *self.current_task }, trap_frame);
     }
 
     #[cfg(riscv)]
     fn switch_task(&mut self) {
-        if !self.to_delete.is_null() {
-            let task_to_delete = core::mem::take(&mut self.to_delete);
-            self.delete_task(task_to_delete);
-        }
+        self.delete_marked_tasks();
 
-        let task = self.current_task;
-        let context = unsafe { &mut (*task).trap_frame };
-        let old_ctx = core::ptr::addr_of_mut!(*context);
+        let Some(next_task) = self.select_next_task() else {
+            return;
+        };
 
-        let task = unsafe { (*self.current_task).next };
-        let context = unsafe { &mut (*task).trap_frame };
-        let new_ctx = core::ptr::addr_of_mut!(*context);
+        let old_ctx = unsafe { &raw mut (*self.current_task).trap_frame };
+        let new_ctx = unsafe { &raw mut (*next_task).trap_frame };
 
         if crate::task::arch_specific::task_switch(old_ctx, new_ctx) {
             unsafe { self.current_task = (*self.current_task).next };
         }
     }
 
-    fn schedule_task_deletion(&mut self, task: *mut Context) -> bool {
-        if task.is_null() {
-            self.to_delete = self.current_task;
-            true
-        } else {
-            self.to_delete = task;
-            core::ptr::eq(task, self.current_task)
+    fn schedule_task_deletion(&mut self, mut task_to_delete: *mut Context) -> bool {
+        if task_to_delete.is_null() {
+            task_to_delete = self.current_task;
         }
+        let is_current = core::ptr::eq(task_to_delete, self.current_task);
+
+        unsafe { (*task_to_delete).next_to_delete = self.to_delete };
+        self.to_delete = task_to_delete;
+
+        is_current
     }
 }
 
diff --git a/esp-preempt/src/task/mod.rs b/esp-preempt/src/task/mod.rs
index 72f9e3277..1a7fd52e8 100644
--- a/esp-preempt/src/task/mod.rs
+++ b/esp-preempt/src/task/mod.rs
@@ -14,6 +14,17 @@ use esp_radio_preempt_driver::semaphore::{SemaphoreHandle, SemaphorePtr};
 
 use crate::{InternalMemory, SCHEDULER_STATE, task, timer};
 
+#[derive(Clone, Copy)]
+pub(crate) enum TaskState {
+    Ready,
+}
+
+impl TaskState {
+    pub fn is_ready(self) -> bool {
+        matches!(self, Self::Ready)
+    }
+}
+
 #[repr(C)]
 pub(crate) struct Context {
     #[cfg(riscv)]
@@ -22,6 +33,8 @@ pub(crate) struct Context {
     pub trap_frame: TrapFrame,
     pub thread_semaphore: Option<SemaphorePtr>,
     pub next: *mut Context,
+    pub next_to_delete: *mut Context,
+    pub state: TaskState,
     pub _allocated_stack: Box<[MaybeUninit<u8>], InternalMemory>,
 }
 
@@ -41,6 +54,8 @@ impl Context {
             trap_frame: task::new_task_context(task_fn, param, stack_top),
             thread_semaphore: None,
             next: core::ptr::null_mut(),
+            next_to_delete: core::ptr::null_mut(),
+            state: TaskState::Ready,
             _allocated_stack: stack,
         }
     }
@@ -65,6 +80,8 @@ pub(super) fn allocate_main_task() {
             trap_frame: TrapFrame::default(),
             thread_semaphore: None,
             next: core::ptr::null_mut(),
+            next_to_delete: core::ptr::null_mut(),
+            state: TaskState::Ready,
             _allocated_stack: Box::<[u8], _>::new_uninit_slice_in(0, InternalMemory),
         },
         InternalMemory,