esp-wifi: support using an external scheduler (#3115)

* support using an external scheduler

* update changelog for `preempt-extern` feature

* change to "extern trait impl" scheme

* make built-in scheduler use the `Scheduler` trait

* pull timer tick out of  trait

* lint `esp-wifi` with `builtin-scheduler`

* esp-wifi: silence warning when not using `builtin-scheduler`

* enable timer tick before scheduler

* esp-wifi: fix preempt.rs type `now` -> `know`

esp-wifi/CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added support for using an external scheduler (#3115)
+
 ### Changed
 
 - `esp_wifi::init` now takes an `impl Peripheral` for RNG source (#2992)

esp-wifi/Cargo.toml

@@ -50,7 +50,7 @@ esp-config   = { version = "0.3.0", path = "../esp-config", features = ["build"]
 esp-metadata = { version = "0.5.0", path = "../esp-metadata" }
 
 [features]
-default = ["esp-alloc"]
+default = ["builtin-scheduler", "esp-alloc"]
 
 ## Use `esp-alloc` for dynamic allocations.
 ##
@@ -131,6 +131,8 @@ smoltcp = ["dep:smoltcp"]
 ## Provide utilities for smoltcp initialization. Adds smoltcp dependency
 utils = ["smoltcp"]
 
+## Use builtin scheduler
+builtin-scheduler = []
 
 # Implement serde Serialize / Deserialize
 serde = ["dep:serde", "enumset?/serde", "heapless/serde"]

esp-wifi/src/lib.rs

@@ -134,7 +134,10 @@ mod binary {
 }
 mod compat;
 
-mod preempt;
+#[cfg(feature = "builtin-scheduler")]
+mod preempt_builtin;
+
+pub mod preempt;
 
 mod radio;
 mod time;
@@ -239,7 +242,7 @@ const _: () = {
     core::assert!(CONFIG.rx_ba_win < (CONFIG.static_rx_buf_num * 2), "WiFi configuration check: rx_ba_win should not be larger than double of the static_rx_buf_num!");
 };
 
-pub(crate) type TimeBase = PeriodicTimer<'static, Blocking>;
+type TimeBase = PeriodicTimer<'static, Blocking>;
 
 pub(crate) mod flags {
     use portable_atomic::{AtomicBool, AtomicUsize};
@@ -388,8 +391,12 @@ pub fn init<'d, T: EspWifiTimerSource, R: EspWifiRngSource>(
 
     setup_radio_isr();
 
-    // This initializes the task switcher and timer tick interrupt.
-    preempt::setup(unsafe { timer.clone_unchecked() }.timer());
+    // Enable timer tick interrupt
+    #[cfg(feature = "builtin-scheduler")]
+    preempt_builtin::setup_timer(unsafe { timer.clone_unchecked() }.timer());
+
+    // This initializes the task switcher
+    preempt::enable();
 
     init_tasks();
     yield_task();
@@ -405,6 +412,9 @@ pub fn init<'d, T: EspWifiTimerSource, R: EspWifiRngSource>(
 
     crate::flags::ESP_WIFI_INITIALIZED.store(true, Ordering::Release);
 
+    #[cfg(not(feature = "builtin-scheduler"))]
+    let _ = timer; // mark used to suppress warning
+
     Ok(EspWifiController {
         _inner: PhantomData,
     })
@@ -446,6 +456,9 @@ pub unsafe fn deinit_unchecked() -> Result<(), InitializationError> {
 
     shutdown_radio_isr();
 
+    #[cfg(feature = "builtin-scheduler")]
+    preempt_builtin::disable_timer();
+
     // This shuts down the task switcher and timer tick interrupt.
     preempt::disable();
 

esp-wifi/src/preempt.rs

@@ -0,0 +1,146 @@
+//! This module allows hooking `esp-wifi` into an external scheduler, instead
+//! of using the integrated one as provided by the `preempt` module.
+//!
+//! In order to hook up an external scheduler, enable the `preempt-extern`
+//! feature, implement the `Scheduler` trait for a struct, and create the
+//! necessary `extern` functions using the `scheduler_impl!()` macro.
+//!
+//! Example:
+//!
+//! ```ignore
+//! use esp_wifi::preempt::Scheduler;
+//!
+//! struct MyScheduler {}
+//!
+//! impl Scheduler for MyScheduler {
+//!     // impl goes here
+//! }
+//!
+//! esp_wifi::scheduler_impl!(static SCHEDULER: MyScheduler = MyScheduler {});
+//! ```
+use core::ffi::c_void;
+
+pub trait Scheduler: Send + Sync + 'static {
+    /// This function is called by `esp-wifi` when starting up the WiFi stack.
+    fn enable(&self);
+
+    /// This function is called by `esp-wifi` when shutting down the WiFi stack.
+    fn disable(&self);
+
+    /// This function is called by threads and should switch to the next thread.
+    fn yield_task(&self);
+
+    /// This function is called by threads and should return an opaque handle
+    /// for the calling thread. The same handle will be passed to
+    /// `esp_wifi_preempt_schedule_task_deletion`.
+    fn current_task(&self) -> *mut c_void;
+
+    /// This function is used to create threads.
+    /// It should allocate the stack.
+    fn task_create(
+        &self,
+        task: extern "C" fn(*mut c_void),
+        param: *mut c_void,
+        task_stack_size: usize,
+    ) -> *mut c_void;
+
+    /// This function is called to let the scheduler know this thread is not
+    /// needed anymore and should be deleted. After this function is called,
+    /// the thread should not be scheduled anymore. The thread stack can be
+    /// free'ed.
+    fn schedule_task_deletion(&self, task_handle: *mut c_void);
+
+    /// This function should return an opaque per-thread pointer to an
+    /// usize-sized memory location, which will be used to store a pointer
+    /// to a semaphore for this thread.
+    fn current_task_thread_semaphore(&self) -> *mut c_void;
+}
+
+extern "Rust" {
+    fn esp_wifi_preempt_enable();
+    fn esp_wifi_preempt_disable();
+    fn esp_wifi_preempt_yield_task();
+    fn esp_wifi_preempt_current_task() -> *mut c_void;
+    fn esp_wifi_preempt_task_create(
+        task: extern "C" fn(*mut c_void),
+        param: *mut c_void,
+        task_stack_size: usize,
+    ) -> *mut c_void;
+    fn esp_wifi_preempt_schedule_task_deletion(task_handle: *mut c_void);
+    fn esp_wifi_preempt_current_task_thread_semaphore() -> *mut c_void;
+}
+
+pub(crate) fn enable() {
+    unsafe { esp_wifi_preempt_enable() }
+}
+
+pub(crate) fn disable() {
+    unsafe { esp_wifi_preempt_disable() }
+}
+
+pub(crate) fn yield_task() {
+    unsafe { esp_wifi_preempt_yield_task() }
+}
+
+pub(crate) fn current_task() -> *mut c_void {
+    unsafe { esp_wifi_preempt_current_task() }
+}
+
+pub(crate) fn task_create(
+    task: extern "C" fn(*mut c_void),
+    param: *mut c_void,
+    task_stack_size: usize,
+) -> *mut c_void {
+    unsafe { esp_wifi_preempt_task_create(task, param, task_stack_size) }
+}
+
+pub(crate) fn schedule_task_deletion(task_handle: *mut c_void) {
+    unsafe { esp_wifi_preempt_schedule_task_deletion(task_handle) }
+}
+
+pub(crate) fn current_task_thread_semaphore() -> *mut c_void {
+    unsafe { esp_wifi_preempt_current_task_thread_semaphore() }
+}
+
+/// Set the Scheduler implementation.
+///
+/// See the module documentation for an example.
+#[macro_export]
+macro_rules! scheduler_impl {
+    (static $name:ident: $t: ty = $val:expr) => {
+        static $name: $t = $val;
+
+        #[no_mangle]
+        fn esp_wifi_preempt_enable() {
+            <$t as $crate::preempt::Scheduler>::enable(&$name)
+        }
+        #[no_mangle]
+        fn esp_wifi_preempt_disable() {
+            <$t as $crate::preempt::Scheduler>::disable(&$name)
+        }
+        #[no_mangle]
+        fn esp_wifi_preempt_yield_task() {
+            <$t as $crate::preempt::Scheduler>::yield_task(&$name)
+        }
+        #[no_mangle]
+        fn esp_wifi_preempt_current_task() -> *mut c_void {
+            <$t as $crate::preempt::Scheduler>::current_task(&$name)
+        }
+        #[no_mangle]
+        fn esp_wifi_preempt_task_create(
+            task: extern "C" fn(*mut c_void),
+            param: *mut c_void,
+            task_stack_size: usize,
+        ) -> *mut c_void {
+            <$t as $crate::preempt::Scheduler>::task_create(&$name, task, param, task_stack_size)
+        }
+        #[no_mangle]
+        fn esp_wifi_preempt_schedule_task_deletion(task_handle: *mut c_void) {
+            <$t as $crate::preempt::Scheduler>::schedule_task_deletion(&$name, task_handle)
+        }
+        #[no_mangle]
+        fn esp_wifi_preempt_current_task_thread_semaphore() -> *mut c_void {
+            <$t as $crate::preempt::Scheduler>::current_task_thread_semaphore(&$name)
+        }
+    };
+}

esp-wifi/src/preempt_builtin/mod.rs

@@ -3,14 +3,13 @@
 mod arch_specific;
 pub mod timer;
 
-use core::mem::size_of;
+use core::{ffi::c_void, mem::size_of};
 
 use arch_specific::*;
 use esp_hal::sync::Locked;
 use esp_wifi_sys::include::malloc;
-use timer::{disable_multitasking, disable_timer, setup_multitasking, setup_timer};
-//
-pub(crate) use {arch_specific::task_create, timer::yield_task};
+use timer::{disable_multitasking, setup_multitasking};
+pub(crate) use timer::{disable_timer, setup_timer};
 
 use crate::{compat::malloc::free, hal::trapframe::TrapFrame, memory_fence::memory_fence};
 
@@ -23,19 +22,53 @@ static CTX_NOW: Locked<ContextWrapper> = Locked::new(ContextWrapper(core::ptr::n
 
 static mut SCHEDULED_TASK_TO_DELETE: *mut Context = core::ptr::null_mut();
 
-pub(crate) fn setup(timer: crate::TimeBase) {
-    // allocate the main task
-    allocate_main_task();
-    setup_timer(timer);
-    setup_multitasking();
-}
+use crate::preempt::Scheduler;
 
-pub(crate) fn disable() {
-    disable_timer();
-    disable_multitasking();
-    delete_all_tasks();
+struct BuiltinScheduler {}
 
-    timer::TIMER.with(|timer| timer.take());
+crate::scheduler_impl!(static SCHEDULER: BuiltinScheduler = BuiltinScheduler {});
+
+impl Scheduler for BuiltinScheduler {
+    fn enable(&self) {
+        // allocate the main task
+        allocate_main_task();
+        setup_multitasking();
+    }
+
+    fn disable(&self) {
+        disable_multitasking();
+        delete_all_tasks();
+
+        timer::TIMER.with(|timer| timer.take());
+    }
+
+    fn yield_task(&self) {
+        timer::yield_task()
+    }
+
+    fn task_create(
+        &self,
+        task: extern "C" fn(*mut c_void),
+        param: *mut c_void,
+        task_stack_size: usize,
+    ) -> *mut c_void {
+        arch_specific::task_create(task, param, task_stack_size) as *mut c_void
+    }
+
+    fn current_task(&self) -> *mut c_void {
+        current_task() as *mut c_void
+    }
+
+    fn schedule_task_deletion(&self, task_handle: *mut c_void) {
+        schedule_task_deletion(task_handle as *mut Context)
+    }
+
+    fn current_task_thread_semaphore(&self) -> *mut crate::binary::c_types::c_void {
+        unsafe {
+            &mut ((*current_task()).thread_semaphore) as *mut _
+                as *mut crate::binary::c_types::c_void
+        }
+    }
 }
 
 fn allocate_main_task() -> *mut Context {
@@ -75,7 +108,7 @@ fn next_task() {
 /// Delete the given task.
 ///
 /// This will also free the memory (stack and context) allocated for it.
-pub(crate) fn delete_task(task: *mut Context) {
+fn delete_task(task: *mut Context) {
     CTX_NOW.with(|ctx_now| unsafe {
         let mut ptr = ctx_now.0;
         let initial = ptr;
@@ -99,7 +132,7 @@ pub(crate) fn delete_task(task: *mut Context) {
     });
 }
 
-pub(crate) fn delete_all_tasks() {
+fn delete_all_tasks() {
     CTX_NOW.with(|ctx_now| unsafe {
         let current_task = ctx_now.0;
 
@@ -128,18 +161,18 @@ pub(crate) fn delete_all_tasks() {
     });
 }
 
-pub(crate) fn current_task() -> *mut Context {
+fn current_task() -> *mut Context {
     CTX_NOW.with(|ctx_now| ctx_now.0)
 }
 
-pub(crate) fn schedule_task_deletion(task: *mut Context) {
+fn schedule_task_deletion(task: *mut Context) {
     unsafe {
         SCHEDULED_TASK_TO_DELETE = task;
     }
 
     if task == current_task() {
         loop {
-            yield_task();
+            timer::yield_task();
         }
     }
 }
@@ -157,9 +190,3 @@ pub(crate) fn task_switch(trap_frame: &mut TrapFrame) {
     next_task();
     restore_task_context(current_task(), trap_frame);
 }
-
-pub(crate) fn current_task_thread_semaphore() -> *mut crate::binary::c_types::c_void {
-    unsafe {
-        &mut ((*current_task()).thread_semaphore) as *mut _ as *mut crate::binary::c_types::c_void
-    }
-}

esp-wifi/src/preempt_builtin/timer/riscv.rs

@@ -11,7 +11,7 @@ use crate::{
         riscv,
         time::Rate,
     },
-    preempt::task_switch,
+    preempt_builtin::task_switch,
     TimeBase,
 };
 

esp-wifi/src/preempt_builtin/timer/xtensa.rs

@@ -2,7 +2,7 @@ use esp_hal::interrupt::InterruptHandler;
 
 use crate::{
     hal::{interrupt, time::Rate, trapframe::TrapFrame, xtensa_lx, xtensa_lx_rt},
-    preempt::task_switch,
+    preempt_builtin::task_switch,
     TimeBase,
 };
 

xtask/src/main.rs

@@ -971,7 +971,9 @@ fn lint_packages(workspace: &Path, args: LintPackagesArgs) -> Result<()> {
                 }
 
                 Package::EspWifi => {
-                    let mut features = format!("--features={chip},defmt,sys-logs,esp-hal/unstable");
+                    let mut features = format!(
+                        "--features={chip},defmt,sys-logs,esp-hal/unstable,builtin-scheduler"
+                    );
 
                     if device.contains("wifi") {
                         features.push_str(",esp-now,sniffer")