Replace some esp-wifi critical sections with locks

esp-hal/src/sync.rs

@@ -280,6 +280,15 @@ impl<T> Locked<T> {
     pub fn with<R>(&self, f: impl FnOnce(&mut T) -> R) -> R {
         lock(&self.lock_state, || f(unsafe { &mut *self.data.get() }))
     }
+
+    /// Provide exclusive access to the protected data to the given closure.
+    pub fn with_cs<R>(
+        &self,
+        _cs: critical_section::CriticalSection<'_>,
+        f: impl FnOnce(&T) -> R,
+    ) -> R {
+        f(unsafe { &*self.data.get() })
+    }
 }
 
 unsafe impl<T> Sync for Locked<T> {}

esp-wifi/src/compat/timer_compat.rs

@@ -1,7 +1,6 @@
 use alloc::boxed::Box;
-use core::cell::RefCell;
 
-use critical_section::Mutex;
+use esp_hal::sync::Locked;
 
 use crate::binary::{
     c_types,
@@ -138,7 +137,7 @@ impl TimerQueue {
 
 unsafe impl Send for TimerQueue {}
 
-pub(crate) static TIMERS: Mutex<RefCell<TimerQueue>> = Mutex::new(RefCell::new(TimerQueue::new()));
+pub(crate) static TIMERS: Locked<TimerQueue> = Locked::new(TimerQueue::new());
 
 pub(crate) fn compat_timer_arm(ets_timer: *mut ets_timer, tmout: u32, repeat: bool) {
     compat_timer_arm_us(ets_timer, tmout * 1000, repeat);
@@ -156,8 +155,8 @@ pub(crate) fn compat_timer_arm_us(ets_timer: *mut ets_timer, us: u32, repeat: bo
         repeat
     );
 
-    critical_section::with(|cs| {
-        if let Some(timer) = TIMERS.borrow_ref_mut(cs).find(ets_timer) {
+    TIMERS.with(|timers| {
+        if let Some(timer) = timers.find(ets_timer) {
             timer.started = systick;
             timer.timeout = ticks;
             timer.active = true;
@@ -170,8 +169,8 @@ pub(crate) fn compat_timer_arm_us(ets_timer: *mut ets_timer, us: u32, repeat: bo
 
 pub fn compat_timer_disarm(ets_timer: *mut ets_timer) {
     trace!("timer disarm");
-    critical_section::with(|cs| {
-        if let Some(timer) = TIMERS.borrow_ref_mut(cs).find(ets_timer) {
+    TIMERS.with(|timers| {
+        if let Some(timer) = timers.find(ets_timer) {
             trace!("timer_disarm {:x}", timer.id());
             timer.active = false;
         } else {
@@ -182,8 +181,7 @@ pub fn compat_timer_disarm(ets_timer: *mut ets_timer) {
 
 pub fn compat_timer_done(ets_timer: *mut ets_timer) {
     trace!("timer done");
-    critical_section::with(|cs| {
-        let mut timers = TIMERS.borrow_ref_mut(cs);
+    TIMERS.with(|timers| {
         if let Some(timer) = timers.find(ets_timer) {
             trace!("timer_done {:x}", timer.id());
             timer.active = false;
@@ -211,8 +209,7 @@ pub(crate) fn compat_timer_setfn(
         pfunction,
         parg
     );
-    let set = critical_section::with(|cs| unsafe {
-        let mut timers = TIMERS.borrow_ref_mut(cs);
+    let set = TIMERS.with(|timers| unsafe {
         if let Some(timer) = timers.find(ets_timer) {
             timer.callback = TimerCallback::new(pfunction, parg);
             timer.active = false;

esp-wifi/src/lib.rs

@@ -438,7 +438,7 @@ pub unsafe fn deinit_unchecked() -> Result<(), InitializationError> {
     shutdown_timer_isr();
     crate::preempt::delete_all_tasks();
 
-    critical_section::with(|cs| crate::timer::TIMER.borrow_ref_mut(cs).take());
+    crate::timer::TIMER.with(|timer| timer.take());
 
     crate::flags::ESP_WIFI_INITIALIZED.store(false, Ordering::Release);
 

esp-wifi/src/preempt/mod.rs

@@ -2,10 +2,10 @@
 #[cfg_attr(target_arch = "xtensa", path = "preempt_xtensa.rs")]
 pub mod arch_specific;
 
-use core::{cell::RefCell, mem::size_of};
+use core::mem::size_of;
 
 use arch_specific::*;
-use critical_section::Mutex;
+use esp_hal::sync::Locked;
 use esp_wifi_sys::include::malloc;
 
 use crate::{compat::malloc::free, hal::trapframe::TrapFrame, memory_fence::memory_fence};
@@ -15,14 +15,12 @@ struct ContextWrapper(*mut Context);
 
 unsafe impl Send for ContextWrapper {}
 
-static CTX_NOW: Mutex<RefCell<ContextWrapper>> =
-    Mutex::new(RefCell::new(ContextWrapper(core::ptr::null_mut())));
+static CTX_NOW: Locked<ContextWrapper> = Locked::new(ContextWrapper(core::ptr::null_mut()));
 
 static mut SCHEDULED_TASK_TO_DELETE: *mut Context = core::ptr::null_mut();
 
 pub(crate) fn allocate_main_task() -> *mut Context {
-    critical_section::with(|cs| unsafe {
-        let mut ctx_now = CTX_NOW.borrow_ref_mut(cs);
+    CTX_NOW.with(|ctx_now| unsafe {
         if !ctx_now.0.is_null() {
             panic!("Tried to allocate main task multiple times");
         }
@@ -36,8 +34,7 @@ pub(crate) fn allocate_main_task() -> *mut Context {
 }
 
 fn allocate_task() -> *mut Context {
-    critical_section::with(|cs| unsafe {
-        let mut ctx_now = CTX_NOW.borrow_ref_mut(cs);
+    CTX_NOW.with(|ctx_now| unsafe {
         if ctx_now.0.is_null() {
             panic!("Called `allocate_task` before allocating main task");
         }
@@ -51,8 +48,7 @@ fn allocate_task() -> *mut Context {
 }
 
 fn next_task() {
-    critical_section::with(|cs| unsafe {
-        let mut ctx_now = CTX_NOW.borrow_ref_mut(cs);
+    CTX_NOW.with(|ctx_now| unsafe {
         ctx_now.0 = (*ctx_now.0).next;
     });
 }
@@ -61,8 +57,8 @@ fn next_task() {
 ///
 /// This will also free the memory (stack and context) allocated for it.
 pub(crate) fn delete_task(task: *mut Context) {
-    critical_section::with(|cs| unsafe {
-        let mut ptr = CTX_NOW.borrow_ref_mut(cs).0;
+    CTX_NOW.with(|ctx_now| unsafe {
+        let mut ptr = ctx_now.0;
         let initial = ptr;
         loop {
             if (*ptr).next == task {
@@ -85,9 +81,8 @@ pub(crate) fn delete_task(task: *mut Context) {
 }
 
 pub(crate) fn delete_all_tasks() {
-    critical_section::with(|cs| unsafe {
-        let mut ctx_now_ref = CTX_NOW.borrow_ref_mut(cs);
-        let current_task = ctx_now_ref.0;
+    CTX_NOW.with(|ctx_now| unsafe {
+        let current_task = ctx_now.0;
 
         if current_task.is_null() {
             return;
@@ -108,14 +103,14 @@ pub(crate) fn delete_all_tasks() {
             task_to_delete = next_task;
         }
 
-        ctx_now_ref.0 = core::ptr::null_mut();
+        ctx_now.0 = core::ptr::null_mut();
 
         memory_fence();
     });
 }
 
 pub(crate) fn current_task() -> *mut Context {
-    critical_section::with(|cs| CTX_NOW.borrow_ref(cs).0)
+    CTX_NOW.with(|ctx_now| ctx_now.0)
 }
 
 pub(crate) fn schedule_task_deletion(task: *mut Context) {

esp-wifi/src/tasks.rs

@@ -18,21 +18,16 @@ pub(crate) fn init_tasks() {
 pub(crate) extern "C" fn timer_task(_param: *mut esp_wifi_sys::c_types::c_void) {
     loop {
         let current_timestamp = systimer_count();
-        let to_run = critical_section::with(|cs| unsafe {
-            let mut timers = TIMERS.borrow_ref_mut(cs);
-            let to_run = timers.find_next_due(current_timestamp);
+        let to_run = TIMERS.with(|timers| {
+            let to_run = unsafe { timers.find_next_due(current_timestamp) }?;
 
-            if let Some(to_run) = to_run {
-                to_run.active = to_run.periodic;
+            to_run.active = to_run.periodic;
 
-                if to_run.periodic {
-                    to_run.started = current_timestamp;
-                }
-
-                Some(to_run.callback)
-            } else {
-                None
+            if to_run.periodic {
+                to_run.started = current_timestamp;
             }
+
+            Some(to_run.callback)
         });
 
         // run the due timer callback NOT in an interrupt free context

esp-wifi/src/timer/mod.rs

@@ -1,6 +1,4 @@
-use core::cell::RefCell;
-
-use critical_section::Mutex;
+use esp_hal::sync::Locked;
 
 #[cfg_attr(esp32, path = "timer_esp32.rs")]
 #[cfg_attr(esp32c2, path = "timer_esp32c2.rs")]
@@ -20,7 +18,7 @@ pub(crate) use chip_specific::*;
 
 use crate::TimeBase;
 
-pub(crate) static TIMER: Mutex<RefCell<Option<TimeBase>>> = Mutex::new(RefCell::new(None));
+pub(crate) static TIMER: Locked<Option<TimeBase>> = Locked::new(None);
 
 pub(crate) fn setup_timer_isr(timebase: TimeBase) {
     setup_radio_isr();

esp-wifi/src/timer/riscv.rs

@@ -30,16 +30,17 @@ pub(crate) fn setup_timer(mut alarm0: TimeBase) {
     let cb: extern "C" fn() = unsafe { core::mem::transmute(handler as *const ()) };
     alarm0.set_interrupt_handler(InterruptHandler::new(cb, interrupt::Priority::Priority1));
     unwrap!(alarm0.start(TIMESLICE_FREQUENCY.into_duration()));
-    critical_section::with(|cs| {
+    TIMER.with(|timer| {
         alarm0.enable_interrupt(true);
-        TIMER.borrow_ref_mut(cs).replace(alarm0);
+        timer.replace(alarm0);
     });
 }
 
 pub(crate) fn disable_timer() {
-    critical_section::with(|cs| {
-        unwrap!(TIMER.borrow_ref_mut(cs).as_mut()).enable_interrupt(false);
-        unwrap!(unwrap!(TIMER.borrow_ref_mut(cs).as_mut()).cancel());
+    TIMER.with(|timer| {
+        let timer = unwrap!(timer.as_mut());
+        timer.enable_interrupt(false);
+        unwrap!(timer.cancel());
     });
 }
 
@@ -60,8 +61,8 @@ pub(crate) fn disable_multitasking() {
 
 extern "C" fn handler(trap_frame: &mut TrapFrame) {
     // clear the systimer intr
-    critical_section::with(|cs| {
-        unwrap!(TIMER.borrow_ref_mut(cs).as_mut()).clear_interrupt();
+    TIMER.with(|timer| {
+        unwrap!(timer.as_mut()).clear_interrupt();
     });
 
     task_switch(trap_frame);
@@ -76,8 +77,7 @@ extern "C" fn FROM_CPU_INTR3(trap_frame: &mut TrapFrame) {
             .modify(|_, w| w.cpu_intr_from_cpu_3().clear_bit());
     }
 
-    critical_section::with(|cs| {
-        let mut alarm0 = TIMER.borrow_ref_mut(cs);
+    TIMER.with(|alarm0| {
         let alarm0 = unwrap!(alarm0.as_mut());
         alarm0.clear_interrupt();
     });

esp-wifi/src/timer/xtensa.rs

@@ -27,16 +27,17 @@ pub(crate) fn setup_timer(mut timer1: TimeBase) {
         interrupt::Priority::Priority2,
     ));
     unwrap!(timer1.start(TIMESLICE_FREQUENCY.into_duration()));
-    critical_section::with(|cs| {
+    TIMER.with(|timer| {
         timer1.enable_interrupt(true);
-        TIMER.borrow_ref_mut(cs).replace(timer1);
+        timer.replace(timer1);
     });
 }
 
 pub(crate) fn disable_timer() {
-    critical_section::with(|cs| {
-        unwrap!(TIMER.borrow_ref_mut(cs).as_mut()).enable_interrupt(false);
-        unwrap!(unwrap!(TIMER.borrow_ref_mut(cs).as_mut()).cancel());
+    TIMER.with(|timer| {
+        let timer = unwrap!(timer.as_mut());
+        timer.enable_interrupt(false);
+        unwrap!(timer.cancel());
     });
 }
 
@@ -58,8 +59,7 @@ pub(crate) fn disable_multitasking() {
 }
 
 fn do_task_switch(context: &mut TrapFrame) {
-    critical_section::with(|cs| {
-        let mut timer = TIMER.borrow_ref_mut(cs);
+    TIMER.with(|timer| {
         let timer = unwrap!(timer.as_mut());
         timer.clear_interrupt();
     });