refactor esp-wifi to ease using external task scheduler (#3106)

* esp-wifi: untangle `setup_timer_isr`

* esp-wifi: bundle scheduler functions into `preempt`

* esp-wifi: split `timer` into `preempt::timer`, `radio` and `time`

* move more deinit stuff into `preempt::disable()`

* move getter of `thread_semaphore` into `preempt`

esp-wifi/src/ble/btdm.rs

@@ -233,7 +233,7 @@ unsafe extern "C" fn task_create(
         extern "C" fn(*mut esp_wifi_sys::c_types::c_void),
     >(func);
 
-    let task = crate::preempt::arch_specific::task_create(task_func, param, stack_depth as usize);
+    let task = crate::preempt::task_create(task_func, param, stack_depth as usize);
     *(handle as *mut usize) = task as usize;
 
     1

esp-wifi/src/ble/npl.rs

@@ -15,7 +15,7 @@ use crate::{
         self,
         common::{str_from_c, ConcurrentQueue},
     },
-    timer::yield_task,
+    preempt::yield_task,
 };
 
 #[cfg_attr(esp32c2, path = "os_adapter_esp32c2.rs")]
@@ -403,7 +403,7 @@ unsafe extern "C" fn task_create(
         extern "C" fn(*mut esp_wifi_sys::c_types::c_void),
     >(task_func);
 
-    let task = crate::preempt::arch_specific::task_create(task_func, param, stack_depth as usize);
+    let task = crate::preempt::task_create(task_func, param, stack_depth as usize);
     *(task_handle as *mut usize) = task as usize;
 
     1

esp-wifi/src/compat/common.rs

@@ -14,8 +14,7 @@ use crate::{
     binary::c_types::{c_int, c_void},
     hal::sync::Locked,
     memory_fence::memory_fence,
-    preempt::current_task,
-    timer::yield_task,
+    preempt::{current_task, yield_task},
 };
 
 pub(crate) const OSI_FUNCS_TIME_BLOCKING: u32 = u32::MAX;
@@ -206,7 +205,7 @@ pub(crate) fn sem_take(semphr: *mut c_void, tick: u32) -> i32 {
 
     let forever = tick == OSI_FUNCS_TIME_BLOCKING;
     let timeout = tick as u64;
-    let start = crate::timer::systimer_count();
+    let start = crate::time::systimer_count();
 
     let sem = semphr as *mut u32;
 
@@ -227,7 +226,7 @@ pub(crate) fn sem_take(semphr: *mut c_void, tick: u32) -> i32 {
             return 1;
         }
 
-        if !forever && crate::timer::elapsed_time_since(start) > timeout {
+        if !forever && crate::time::elapsed_time_since(start) > timeout {
             break 'outer;
         }
 
@@ -251,7 +250,7 @@ pub(crate) fn sem_give(semphr: *mut c_void) -> i32 {
 
 pub(crate) fn thread_sem_get() -> *mut c_void {
     trace!("wifi_thread_semphr_get");
-    unsafe { &mut ((*current_task()).thread_semaphore) as *mut _ as *mut c_void }
+    crate::preempt::current_task_thread_semaphore()
 }
 
 pub(crate) fn create_recursive_mutex() -> *mut c_void {
@@ -376,7 +375,7 @@ pub(crate) fn receive_queued(
 
     let forever = block_time_tick == OSI_FUNCS_TIME_BLOCKING;
     let timeout = block_time_tick as u64;
-    let start = crate::timer::systimer_count();
+    let start = crate::time::systimer_count();
 
     loop {
         if unsafe { (*queue).try_dequeue(item) } {
@@ -384,7 +383,7 @@ pub(crate) fn receive_queued(
             return 1;
         }
 
-        if !forever && crate::timer::elapsed_time_since(start) > timeout {
+        if !forever && crate::time::elapsed_time_since(start) > timeout {
             trace!("queue_recv returns with timeout");
             return -1;
         }

esp-wifi/src/compat/timer_compat.rs

@@ -76,7 +76,7 @@ impl TimerQueue {
         let mut current = self.head.as_mut();
         while let Some(timer) = current {
             if timer.active
-                && crate::timer::time_diff(timer.started, current_timestamp) >= timer.timeout
+                && crate::time::time_diff(timer.started, current_timestamp) >= timer.timeout
             {
                 return Some(timer);
             }
@@ -149,8 +149,8 @@ pub(crate) fn compat_timer_arm(ets_timer: *mut ets_timer, tmout: u32, repeat: bo
 }
 
 pub(crate) fn compat_timer_arm_us(ets_timer: *mut ets_timer, us: u32, repeat: bool) {
-    let systick = crate::timer::systimer_count();
-    let ticks = crate::timer::micros_to_ticks(us as u64);
+    let systick = crate::time::systimer_count();
+    let ticks = crate::time::micros_to_ticks(us as u64);
 
     trace!(
         "timer_arm_us {:x} current: {} ticks: {} repeat: {}",

esp-wifi/src/lib.rs

@@ -124,17 +124,20 @@ use portable_atomic::Ordering;
 #[cfg(feature = "wifi")]
 use crate::wifi::WifiError;
 use crate::{
+    preempt::yield_task,
+    radio::{setup_radio_isr, shutdown_radio_isr},
     tasks::init_tasks,
-    timer::{setup_timer_isr, shutdown_timer_isr},
 };
 
 mod binary {
     pub use esp_wifi_sys::*;
 }
 mod compat;
+
 mod preempt;
 
-mod timer;
+mod radio;
+mod time;
 
 #[cfg(feature = "wifi")]
 pub mod wifi;
@@ -236,7 +239,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!");
 };
 
-type TimeBase = PeriodicTimer<'static, Blocking>;
+pub(crate) type TimeBase = PeriodicTimer<'static, Blocking>;
 
 pub(crate) mod flags {
     use portable_atomic::{AtomicBool, AtomicUsize};
@@ -382,8 +385,14 @@ pub fn init<'d, T: EspWifiTimerSource, R: EspWifiRngSource>(
     info!("esp-wifi configuration {:?}", crate::CONFIG);
     crate::common_adapter::chip_specific::enable_wifi_power_domain();
     phy_mem_init();
+
+    setup_radio_isr();
+
+    // This initializes the task switcher and timer tick interrupt.
+    preempt::setup(unsafe { timer.clone_unchecked() }.timer());
+
     init_tasks();
-    setup_timer_isr(unsafe { timer.clone_unchecked() }.timer());
+    yield_task();
 
     wifi_set_log_verbose();
     init_clocks();
@@ -435,10 +444,10 @@ pub unsafe fn deinit_unchecked() -> Result<(), InitializationError> {
         crate::flags::BLE.store(false, Ordering::Release);
     }
 
-    shutdown_timer_isr();
-    crate::preempt::delete_all_tasks();
+    shutdown_radio_isr();
 
-    crate::timer::TIMER.with(|timer| timer.take());
+    // This shuts down the task switcher and timer tick interrupt.
+    preempt::disable();
 
     crate::flags::ESP_WIFI_INITIALIZED.store(false, Ordering::Release);
 

esp-wifi/src/preempt/mod.rs

@@ -1,12 +1,16 @@
 #[cfg_attr(target_arch = "riscv32", path = "preempt_riscv.rs")]
 #[cfg_attr(target_arch = "xtensa", path = "preempt_xtensa.rs")]
-pub mod arch_specific;
+mod arch_specific;
+pub mod timer;
 
 use core::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 crate::{compat::malloc::free, hal::trapframe::TrapFrame, memory_fence::memory_fence};
 
@@ -19,7 +23,22 @@ 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 allocate_main_task() -> *mut Context {
+pub(crate) fn setup(timer: crate::TimeBase) {
+    // allocate the main task
+    allocate_main_task();
+    setup_timer(timer);
+    setup_multitasking();
+}
+
+pub(crate) fn disable() {
+    disable_timer();
+    disable_multitasking();
+    delete_all_tasks();
+
+    timer::TIMER.with(|timer| timer.take());
+}
+
+fn allocate_main_task() -> *mut Context {
     CTX_NOW.with(|ctx_now| unsafe {
         if !ctx_now.0.is_null() {
             panic!("Tried to allocate main task multiple times");
@@ -114,8 +133,6 @@ pub(crate) fn current_task() -> *mut Context {
 }
 
 pub(crate) fn schedule_task_deletion(task: *mut Context) {
-    use crate::timer::yield_task;
-
     unsafe {
         SCHEDULED_TASK_TO_DELETE = task;
     }
@@ -140,3 +157,9 @@ 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/timer/mod.rs

@@ -0,0 +1,11 @@
+use esp_hal::sync::Locked;
+
+#[cfg_attr(xtensa, path = "xtensa.rs")]
+#[cfg_attr(riscv, path = "riscv.rs")]
+mod arch_specific;
+
+pub(crate) use arch_specific::*;
+
+use crate::TimeBase;
+
+pub(crate) static TIMER: Locked<Option<TimeBase>> = Locked::new(None);

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

@@ -18,9 +18,6 @@ use crate::{
 /// The timer responsible for time slicing.
 const TIMESLICE_FREQUENCY: Rate = Rate::from_hz(crate::CONFIG.tick_rate_hz);
 
-// Time keeping
-pub const TICKS_PER_SECOND: u64 = 1_000_000;
-
 use super::TIMER;
 
 pub(crate) fn setup_timer(mut alarm0: TimeBase) {
@@ -88,17 +85,3 @@ pub(crate) fn yield_task() {
         .cpu_intr_from_cpu_3()
         .modify(|_, w| w.cpu_intr_from_cpu_3().set_bit());
 }
-
-/// Current systimer count value
-/// A tick is 1 / 1_000_000 seconds
-pub(crate) fn systimer_count() -> u64 {
-    esp_hal::time::Instant::now()
-        .duration_since_epoch()
-        .as_micros()
-}
-
-// TODO: use an Instance type instead...
-pub(crate) fn time_diff(start: u64, end: u64) -> u64 {
-    // 52-bit wrapping sub
-    end.wrapping_sub(start) & 0x000f_ffff_ffff_ffff
-}

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

@@ -11,17 +11,6 @@ const TIMESLICE_FREQUENCY: Rate = Rate::from_hz(crate::CONFIG.tick_rate_hz);
 
 use super::TIMER;
 
-// Time keeping
-pub const TICKS_PER_SECOND: u64 = 1_000_000;
-
-/// This function must not be called in a critical section. Doing so may return
-/// an incorrect value.
-pub(crate) fn systimer_count() -> u64 {
-    esp_hal::time::Instant::now()
-        .duration_since_epoch()
-        .as_micros()
-}
-
 pub(crate) fn setup_timer(mut timer1: TimeBase) {
     timer1.set_interrupt_handler(InterruptHandler::new(
         unsafe { core::mem::transmute::<*const (), extern "C" fn()>(handler as *const ()) },
@@ -89,8 +78,3 @@ pub(crate) fn yield_task() {
         core::arch::asm!("wsr.intset  {0}", in(reg) intr, options(nostack));
     }
 }
-
-// TODO: use an Instance type instead...
-pub(crate) fn time_diff(start: u64, end: u64) -> u64 {
-    end.wrapping_sub(start)
-}

esp-wifi/src/radio/mod.rs

@@ -0,0 +1,10 @@
+#[cfg_attr(esp32, path = "radio_esp32.rs")]
+#[cfg_attr(esp32c2, path = "radio_esp32c2.rs")]
+#[cfg_attr(esp32c3, path = "radio_esp32c3.rs")]
+#[cfg_attr(esp32c6, path = "radio_esp32c6.rs")]
+#[cfg_attr(esp32h2, path = "radio_esp32h2.rs")]
+#[cfg_attr(esp32s3, path = "radio_esp32s3.rs")]
+#[cfg_attr(esp32s2, path = "radio_esp32s2.rs")]
+mod chip_specific;
+
+pub(crate) use chip_specific::*;

esp-wifi/src/tasks.rs

@@ -1,14 +1,11 @@
 use crate::{
     compat::timer_compat::TIMERS,
-    preempt::arch_specific::task_create,
-    timer::{systimer_count, yield_task},
+    preempt::{task_create, yield_task},
+    time::systimer_count,
 };
 
-/// Initializes the `main` and `timer` tasks for the Wi-Fi driver.
+/// Initializes the `timer` task for the Wi-Fi driver.
 pub(crate) fn init_tasks() {
-    // allocate the main task
-    crate::preempt::allocate_main_task();
-
     // schedule the timer task
     task_create(timer_task, core::ptr::null_mut(), 8192);
 }

esp-wifi/src/time.rs

@@ -0,0 +1,51 @@
+// Time keeping
+pub const TICKS_PER_SECOND: u64 = 1_000_000;
+
+/// Current systimer count value
+/// A tick is 1 / 1_000_000 seconds
+/// This function must not be called in a critical section. Doing so may return
+/// an incorrect value.
+pub(crate) fn systimer_count() -> u64 {
+    esp_hal::time::Instant::now()
+        .duration_since_epoch()
+        .as_micros()
+}
+
+// TODO: use an Instance type instead...
+#[cfg(target_arch = "riscv32")]
+pub(crate) fn time_diff(start: u64, end: u64) -> u64 {
+    // 52-bit wrapping sub
+    end.wrapping_sub(start) & 0x000f_ffff_ffff_ffff
+}
+
+// TODO: use an Instance type instead...
+#[cfg(target_arch = "xtensa")]
+pub(crate) fn time_diff(start: u64, end: u64) -> u64 {
+    end.wrapping_sub(start)
+}
+
+#[allow(unused)]
+pub(crate) fn micros_to_ticks(us: u64) -> u64 {
+    us * (TICKS_PER_SECOND / 1_000_000)
+}
+
+#[allow(unused)]
+pub(crate) fn millis_to_ticks(ms: u64) -> u64 {
+    ms * (TICKS_PER_SECOND / 1_000)
+}
+
+#[allow(unused)]
+pub(crate) fn ticks_to_micros(ticks: u64) -> u64 {
+    ticks / (TICKS_PER_SECOND / 1_000_000)
+}
+
+#[allow(unused)]
+pub(crate) fn ticks_to_millis(ticks: u64) -> u64 {
+    ticks / (TICKS_PER_SECOND / 1_000)
+}
+
+/// Do not call this in a critical section!
+pub(crate) fn elapsed_time_since(start: u64) -> u64 {
+    let now = systimer_count();
+    time_diff(start, now)
+}

esp-wifi/src/timer/mod.rs

@@ -1,65 +0,0 @@
-use esp_hal::sync::Locked;
-
-#[cfg_attr(esp32, path = "timer_esp32.rs")]
-#[cfg_attr(esp32c2, path = "timer_esp32c2.rs")]
-#[cfg_attr(esp32c3, path = "timer_esp32c3.rs")]
-#[cfg_attr(esp32c6, path = "timer_esp32c6.rs")]
-#[cfg_attr(esp32h2, path = "timer_esp32h2.rs")]
-#[cfg_attr(esp32s3, path = "timer_esp32s3.rs")]
-#[cfg_attr(esp32s2, path = "timer_esp32s2.rs")]
-mod chip_specific;
-
-#[cfg_attr(any(esp32, esp32s2, esp32s3), path = "xtensa.rs")]
-#[cfg_attr(any(esp32c2, esp32c3, esp32c6, esp32h2), path = "riscv.rs")]
-mod arch_specific;
-
-pub(crate) use arch_specific::*;
-pub(crate) use chip_specific::*;
-
-use crate::TimeBase;
-
-pub(crate) static TIMER: Locked<Option<TimeBase>> = Locked::new(None);
-
-pub(crate) fn setup_timer_isr(timebase: TimeBase) {
-    setup_radio_isr();
-
-    setup_timer(timebase);
-
-    setup_multitasking();
-
-    yield_task();
-}
-
-pub(crate) fn shutdown_timer_isr() {
-    shutdown_radio_isr();
-
-    disable_timer();
-
-    disable_multitasking();
-}
-
-#[allow(unused)]
-pub(crate) fn micros_to_ticks(us: u64) -> u64 {
-    us * (TICKS_PER_SECOND / 1_000_000)
-}
-
-#[allow(unused)]
-pub(crate) fn millis_to_ticks(ms: u64) -> u64 {
-    ms * (TICKS_PER_SECOND / 1_000)
-}
-
-#[allow(unused)]
-pub(crate) fn ticks_to_micros(ticks: u64) -> u64 {
-    ticks / (TICKS_PER_SECOND / 1_000_000)
-}
-
-#[allow(unused)]
-pub(crate) fn ticks_to_millis(ticks: u64) -> u64 {
-    ticks / (TICKS_PER_SECOND / 1_000)
-}
-
-/// Do not call this in a critical section!
-pub(crate) fn elapsed_time_since(start: u64) -> u64 {
-    let now = systimer_count();
-    time_diff(start, now)
-}

esp-wifi/src/wifi/mod.rs

@@ -1915,7 +1915,7 @@ mod sealed {
 
         fn tx_token(self) -> Option<WifiTxToken<Self>> {
             if !self.can_send() {
-                crate::timer::yield_task();
+                crate::preempt::yield_task();
             }
 
             if self.can_send() {
@@ -1928,7 +1928,7 @@ mod sealed {
         fn rx_token(self) -> Option<(WifiRxToken<Self>, WifiTxToken<Self>)> {
             let is_empty = self.data_queue_rx().with(|q| q.is_empty());
             if is_empty || !self.can_send() {
-                crate::timer::yield_task();
+                crate::preempt::yield_task();
             }
 
             let is_empty = is_empty && self.data_queue_rx().with(|q| q.is_empty());

esp-wifi/src/wifi/os_adapter.rs

@@ -35,7 +35,7 @@ use crate::{
         sync::{Locked, RawMutex},
     },
     memory_fence::memory_fence,
-    timer::yield_task,
+    preempt::yield_task,
 };
 
 static WIFI_LOCK: RawMutex = RawMutex::new();
@@ -665,7 +665,7 @@ pub unsafe extern "C" fn task_create_pinned_to_core(
         extern "C" fn(*mut esp_wifi_sys::c_types::c_void),
     >(task_func);
 
-    let task = crate::preempt::arch_specific::task_create(task_func, param, stack_depth as usize);
+    let task = crate::preempt::task_create(task_func, param, stack_depth as usize);
     *(task_handle as *mut usize) = task as usize;
 
     1
@@ -741,8 +741,8 @@ pub unsafe extern "C" fn task_delete(task_handle: *mut crate::binary::c_types::c
 /// *************************************************************************
 pub unsafe extern "C" fn task_delay(tick: u32) {
     trace!("task_delay tick {}", tick);
-    let start_time = crate::timer::systimer_count();
-    while crate::timer::elapsed_time_since(start_time) < tick as u64 {
+    let start_time = crate::time::systimer_count();
+    while crate::time::elapsed_time_since(start_time) < tick as u64 {
         yield_task();
     }
 }
@@ -762,7 +762,7 @@ pub unsafe extern "C" fn task_delay(tick: u32) {
 /// *************************************************************************
 pub unsafe extern "C" fn task_ms_to_tick(ms: u32) -> i32 {
     trace!("task_ms_to_tick ms {}", ms);
-    crate::timer::millis_to_ticks(ms as u64) as i32
+    crate::time::millis_to_ticks(ms as u64) as i32
 }
 
 /// **************************************************************************
@@ -1130,7 +1130,7 @@ pub unsafe extern "C" fn wifi_rtc_disable_iso() {
 #[no_mangle]
 pub unsafe extern "C" fn esp_timer_get_time() -> i64 {
     trace!("esp_timer_get_time");
-    crate::timer::ticks_to_micros(crate::timer::systimer_count()) as i64
+    crate::time::ticks_to_micros(crate::time::systimer_count()) as i64
 }
 
 /// **************************************************************************