esp-preempt: don't switch to sleeping tasks (#4081)

* Separate the alloc and run lists

* Replace circular task list with ready queue

* Remove separate SCHEDULER_STATE static

* Move scheduler to new file

* Reorganize, allow restarting scheduler

* Fix InternalMemory polyfill

* Use SingleShotTimer internally

* Implement a simple timer queue

* Extract run queue, wake tasks, store reason of scheduler event

* Add inherent function to get current task ptr

* Reimplement usleep with the timer queue

* Store current task in timer queue

* Sleep in timer queue task

* Remove ability to sleep arbitrary tasks

* More logging

* Clear timer interrupt in timer handler

* Even more logging

* Merge mutexes into semaphores

Author: bugadani

esp-preempt/Cargo.toml

@@ -49,6 +49,10 @@ esp-hal = { version = "1.0.0-rc.0", path = "../esp-hal", features = ["unstable"]
 default = ["esp-alloc"]
 
 ## Enable the use of the `esp-alloc` crate for dynamic memory allocation.
+##
+## If you opt-out, you need to provide implementations for the following functions:
+## - `pub extern "C" fn malloc_internal(size: usize) -> *mut u8`
+## - `pub extern "C" fn free_internal(ptr: *mut u8)`
 esp-alloc = ["dep:esp-alloc"]
 
 #! ### Chip selection

esp-preempt/src/lib.rs

@@ -23,27 +23,24 @@ extern crate alloc;
 // MUST be the first module
 mod fmt;
 
-mod mutex;
 mod queue;
+mod run_queue;
+mod scheduler;
 mod semaphore;
 mod task;
 mod timer;
 mod timer_queue;
 
-use core::ffi::c_void;
-
-use allocator_api2::boxed::Box;
+pub(crate) use esp_alloc::InternalMemory;
 use esp_hal::{
     Blocking,
-    time::{Duration, Instant, Rate},
-    timer::{AnyTimer, PeriodicTimer},
+    timer::{AnyTimer, OneShotTimer},
 };
-use esp_radio_preempt_driver::semaphore::{SemaphoreImplementation, SemaphorePtr};
-use esp_sync::NonReentrantMutex;
+pub(crate) use scheduler::SCHEDULER;
 
-use crate::{semaphore::Semaphore, task::Context, timer::TIMER};
+use crate::timer::TimeDriver;
 
-type TimeBase = PeriodicTimer<'static, Blocking>;
+type TimeBase = OneShotTimer<'static, Blocking>;
 
 // Polyfill the InternalMemory allocator
 #[cfg(not(feature = "esp-alloc"))]
@@ -52,32 +49,28 @@ mod esp_alloc {
 
     use allocator_api2::alloc::{AllocError, Allocator};
 
+    unsafe extern "C" {
+        fn malloc_internal(size: usize) -> *mut u8;
+
+        fn free_internal(ptr: *mut u8);
+    }
+
     /// An allocator that uses internal memory only.
     pub struct InternalMemory;
 
     unsafe impl Allocator for InternalMemory {
         fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
-            unsafe extern "C" {
-                fn esp_radio_allocate_from_internal_ram(size: usize) -> *mut u8;
-            }
-            let raw_ptr = unsafe { esp_radio_allocate_from_internal_ram(layout.size()) };
+            let raw_ptr = unsafe { malloc_internal(layout.size()) };
             let ptr = NonNull::new(raw_ptr).ok_or(AllocError)?;
             Ok(NonNull::slice_from_raw_parts(ptr, layout.size()))
         }
 
         unsafe fn deallocate(&self, ptr: NonNull<u8>, _layout: Layout) {
-            unsafe extern "C" {
-                fn esp_radio_deallocate_internal_ram(ptr: *mut u8);
-            }
-            unsafe {
-                esp_radio_deallocate_internal_ram(ptr.as_ptr());
-            }
+            unsafe { free_internal(ptr.as_ptr()) };
         }
     }
 }
 
-pub(crate) use esp_alloc::InternalMemory;
-
 /// A trait to allow better UX for initializing esp-preempt.
 ///
 /// This trait is meant to be used only for the `init` function.
@@ -107,174 +100,6 @@ where
     }
 }
 
-struct SchedulerState {
-    /// Pointer to the current task.
-    ///
-    /// Tasks are stored in a circular linked list. CTX_NOW points to the
-    /// current task.
-    current_task: *mut Context,
-
-    /// Pointer to the task that is scheduled for deletion.
-    to_delete: *mut Context,
-}
-
-unsafe impl Send for SchedulerState {}
-
-impl SchedulerState {
-    const fn new() -> Self {
-        Self {
-            current_task: core::ptr::null_mut(),
-            to_delete: core::ptr::null_mut(),
-        }
-    }
-
-    fn delete_task(&mut self, task: *mut Context) {
-        let mut current_task = self.current_task;
-        // Save the first pointer so we can prevent an accidental infinite loop.
-        let initial = current_task;
-        loop {
-            // We don't have the previous pointer, so we need to walk forward in the circle
-            // even if we need to delete the first task.
-
-            // If the next task is the one we want to delete, we need to remove it from the
-            // list, then drop it.
-            let next_task = unsafe { (*current_task).next };
-            if core::ptr::eq(next_task, task) {
-                unsafe {
-                    (*current_task).next = (*next_task).next;
-
-                    core::ptr::drop_in_place(task);
-                    break;
-                }
-            }
-
-            // If the next task is the first task, we can stop. If we needed to delete the
-            // first task, we have already handled it in the above case. If we needed to
-            // delete another task, it has already been deleted in a previous iteration.
-            if core::ptr::eq(next_task, initial) {
-                break;
-            }
-
-            // Move to the next task.
-            current_task = next_task;
-        }
-    }
-
-    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);
-        }
-    }
-
-    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) {
-        self.delete_marked_tasks();
-
-        let Some(next_task) = self.select_next_task() else {
-            return;
-        };
-
-        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, 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
-    }
-}
-
-fn usleep(us: u32) {
-    trace!("usleep");
-    unsafe extern "C" {
-        fn esp_rom_delay_us(us: u32);
-    }
-
-    const MIN_YIELD_TIME: u32 = 1_000_000 / TICK_RATE;
-    if us < MIN_YIELD_TIME {
-        // Short wait, just sleep
-        unsafe { esp_rom_delay_us(us) };
-    } else {
-        const MIN_YIELD_DURATION: Duration = Duration::from_micros(MIN_YIELD_TIME as u64);
-        let sleep_for = Duration::from_micros(us as u64);
-        let start = Instant::now();
-        loop {
-            // Yield to other tasks
-            timer::yield_task();
-
-            let elapsed = start.elapsed();
-            if elapsed.as_micros() > us as u64 {
-                break;
-            }
-
-            let remaining = sleep_for - elapsed;
-
-            if remaining < MIN_YIELD_DURATION {
-                // If the remaining time is less than the minimum yield time, we can just sleep
-                // for the remaining time.
-                unsafe { esp_rom_delay_us(remaining.as_micros() as u32) };
-                break;
-            }
-        }
-    }
-}
-
-static SCHEDULER_STATE: NonReentrantMutex<SchedulerState> =
-    NonReentrantMutex::new(SchedulerState::new());
-
-struct Scheduler {}
-
-esp_radio_preempt_driver::scheduler_impl!(static SCHEDULER: Scheduler = Scheduler {});
-
 /// Initializes the scheduler.
 ///
 /// # The `timer` argument
@@ -288,99 +113,7 @@ esp_radio_preempt_driver::scheduler_impl!(static SCHEDULER: Scheduler = Schedule
 ///
 /// For an example, see the [crate-level documentation][self].
 pub fn init(timer: impl TimerSource) {
-    timer::setup_timebase(timer.timer());
+    SCHEDULER.with(move |scheduler| scheduler.set_time_driver(TimeDriver::new(timer.timer())))
 }
 
 const TICK_RATE: u32 = esp_config::esp_config_int!(u32, "ESP_PREEMPT_CONFIG_TICK_RATE_HZ");
-const TIMESLICE_FREQUENCY: Rate = Rate::from_hz(TICK_RATE);
-
-impl esp_radio_preempt_driver::Scheduler for Scheduler {
-    fn initialized(&self) -> bool {
-        timer::initialized()
-    }
-
-    fn enable(&self) {
-        // allocate the main task
-        task::allocate_main_task();
-        timer::setup_multitasking();
-        timer_queue::create_timer_task();
-
-        TIMER.with(|t| {
-            let t = unwrap!(t.as_mut());
-            unwrap!(t.start(TIMESLICE_FREQUENCY.as_duration()));
-        });
-    }
-
-    fn disable(&self) {
-        timer::disable_timebase();
-        timer::disable_multitasking();
-        task::delete_all_tasks();
-    }
-
-    fn yield_task(&self) {
-        timer::yield_task()
-    }
-
-    fn yield_task_from_isr(&self) {
-        self.yield_task();
-    }
-
-    fn max_task_priority(&self) -> u32 {
-        255
-    }
-
-    fn task_create(
-        &self,
-        task: extern "C" fn(*mut c_void),
-        param: *mut c_void,
-        _priority: u32,
-        _pin_to_core: Option<u32>,
-        task_stack_size: usize,
-    ) -> *mut c_void {
-        let task = Box::new_in(Context::new(task, param, task_stack_size), InternalMemory);
-        let task_ptr = Box::into_raw(task);
-
-        SCHEDULER_STATE.with(|state| unsafe {
-            let current_task = state.current_task;
-            debug_assert!(
-                !current_task.is_null(),
-                "Tried to allocate a task before allocating the main task"
-            );
-            // Insert the new task at the next position.
-            let next = (*current_task).next;
-            (*task_ptr).next = next;
-            (*current_task).next = task_ptr;
-        });
-
-        task_ptr as *mut c_void
-    }
-
-    fn current_task(&self) -> *mut c_void {
-        task::current_task() as *mut c_void
-    }
-
-    fn schedule_task_deletion(&self, task_handle: *mut c_void) {
-        task::schedule_task_deletion(task_handle as *mut Context)
-    }
-
-    fn current_task_thread_semaphore(&self) -> SemaphorePtr {
-        task::with_current_task(|task| {
-            if task.thread_semaphore.is_none() {
-                task.thread_semaphore = Some(Semaphore::create(1, 0));
-            }
-
-            unwrap!(task.thread_semaphore)
-        })
-    }
-
-    fn usleep(&self, us: u32) {
-        usleep(us)
-    }
-
-    fn now(&self) -> u64 {
-        // FIXME: this function needs to return the timestamp of the scheduler's timer
-        esp_hal::time::Instant::now()
-            .duration_since_epoch()
-            .as_micros()
-    }
-}