Protect SYSTIMER/TIMG shared registers (#2051)

* Protect SYSTIMER/TIMG shared registers

* some review comments

* more review comments

* more review comments

* bring back portable_atomic

---------

Co-authored-by: Dominic Fischer <git@dominicfischer.me>

esp-hal/src/lib.rs

@@ -572,6 +572,94 @@ mod critical_section_impl {
     }
 }
 
+// The state of a re-entrant lock
+pub(crate) struct LockState {
+    #[cfg(multi_core)]
+    core: portable_atomic::AtomicUsize,
+}
+
+impl LockState {
+    #[cfg(multi_core)]
+    const UNLOCKED: usize = usize::MAX;
+
+    pub const fn new() -> Self {
+        Self {
+            #[cfg(multi_core)]
+            core: portable_atomic::AtomicUsize::new(Self::UNLOCKED),
+        }
+    }
+}
+
+// This is preferred over critical-section as this allows you to have multiple
+// locks active at the same time rather than using the global mutex that is
+// critical-section.
+#[allow(unused_variables)]
+pub(crate) fn lock<T>(state: &LockState, f: impl FnOnce() -> T) -> T {
+    // In regards to disabling interrupts, we only need to disable
+    // the interrupts that may be calling this function.
+
+    #[cfg(not(multi_core))]
+    {
+        // Disabling interrupts is enough on single core chips to ensure mutual
+        // exclusion.
+
+        #[cfg(riscv)]
+        return riscv::interrupt::free(f);
+        #[cfg(xtensa)]
+        return xtensa_lx::interrupt::free(|_| f());
+    }
+
+    #[cfg(multi_core)]
+    {
+        use portable_atomic::Ordering;
+
+        let current_core = get_core() as usize;
+
+        let mut f = f;
+
+        loop {
+            let func = || {
+                // Use Acquire ordering in success to ensure `f()` "happens after" the lock is
+                // taken. Use Relaxed ordering in failure as there's no
+                // synchronisation happening.
+                if let Err(locked_core) = state.core.compare_exchange(
+                    LockState::UNLOCKED,
+                    current_core,
+                    Ordering::Acquire,
+                    Ordering::Relaxed,
+                ) {
+                    assert_ne!(
+                        locked_core, current_core,
+                        "esp_hal::lock is not re-entrant!"
+                    );
+
+                    Err(f)
+                } else {
+                    let result = f();
+
+                    // Use Release ordering here to ensure `f()` "happens before" this lock is
+                    // released.
+                    state.core.store(LockState::UNLOCKED, Ordering::Release);
+
+                    Ok(result)
+                }
+            };
+
+            #[cfg(riscv)]
+            let result = riscv::interrupt::free(func);
+            #[cfg(xtensa)]
+            let result = xtensa_lx::interrupt::free(|_| func());
+
+            match result {
+                Ok(result) => break result,
+                Err(the_function) => f = the_function,
+            }
+
+            // Consider using core::hint::spin_loop(); Might need SW_INT.
+        }
+    }
+}
+
 /// Default (unhandled) interrupt handler
 pub const DEFAULT_INTERRUPT_HANDLER: interrupt::InterruptHandler = interrupt::InterruptHandler::new(
     unsafe { core::mem::transmute::<*const (), extern "C" fn()>(EspDefaultHandler as *const ()) },

esp-hal/src/timer/systimer.rs

@@ -80,6 +80,7 @@ use fugit::{Instant, MicrosDurationU32, MicrosDurationU64};
 use super::{Error, Timer as _};
 use crate::{
     interrupt::{self, InterruptHandler},
+    lock,
     peripheral::Peripheral,
     peripherals::{Interrupt, SYSTIMER},
     system::{Peripheral as PeripheralEnable, PeripheralClockControl},
@@ -87,6 +88,7 @@ use crate::{
     Blocking,
     Cpu,
     InterruptConfigurable,
+    LockState,
     Mode,
 };
 
@@ -240,7 +242,7 @@ pub trait Unit {
         let systimer = unsafe { &*SYSTIMER::ptr() };
         let conf = systimer.conf();
 
-        critical_section::with(|_| {
+        lock(&CONF_LOCK, || {
             conf.modify(|_, w| match config {
                 UnitConfig::Disabled => match self.channel() {
                     0 => w.timer_unit0_work_en().clear_bit(),
@@ -418,7 +420,7 @@ pub trait Comparator {
     fn set_enable(&self, enable: bool) {
         let systimer = unsafe { &*SYSTIMER::ptr() };
 
-        critical_section::with(|_| {
+        lock(&CONF_LOCK, || {
             #[cfg(not(esp32s2))]
             systimer.conf().modify(|_, w| match self.channel() {
                 0 => w.target0_work_en().bit(enable),
@@ -426,12 +428,14 @@ pub trait Comparator {
                 2 => w.target2_work_en().bit(enable),
                 _ => unreachable!(),
             });
-
-            #[cfg(esp32s2)]
-            systimer
-                .target_conf(self.channel() as usize)
-                .modify(|_r, w| w.work_en().bit(enable));
         });
+
+        // Note: The ESP32-S2 doesn't require a lock because each
+        // comparator's enable bit in a different register.
+        #[cfg(esp32s2)]
+        systimer
+            .target_conf(self.channel() as usize)
+            .modify(|_r, w| w.work_en().bit(enable));
     }
 
     /// Returns true if the comparator has been enabled. This means
@@ -964,9 +968,11 @@ where
     }
 
     fn enable_interrupt(&self, state: bool) {
-        unsafe { &*SYSTIMER::PTR }
+        lock(&INT_ENA_LOCK, || {
-            .int_ena()
+            unsafe { &*SYSTIMER::PTR }
-            .modify(|_, w| w.target(self.comparator.channel()).bit(state));
+                .int_ena()
+                .modify(|_, w| w.target(self.comparator.channel()).bit(state));
+        });
     }
 
     fn clear_interrupt(&self) {
@@ -1004,6 +1010,9 @@ where
     }
 }
 
+static CONF_LOCK: LockState = LockState::new();
+static INT_ENA_LOCK: LockState = LockState::new();
+
 // Async functionality of the system timer.
 #[cfg(feature = "async")]
 mod asynch {
@@ -1090,27 +1099,33 @@ mod asynch {
 
     #[handler]
     fn target0_handler() {
-        unsafe { &*crate::peripherals::SYSTIMER::PTR }
+        lock(&INT_ENA_LOCK, || {
-            .int_ena()
+            unsafe { &*crate::peripherals::SYSTIMER::PTR }
-            .modify(|_, w| w.target0().clear_bit());
+                .int_ena()
+                .modify(|_, w| w.target0().clear_bit());
+        });
 
         WAKERS[0].wake();
     }
 
     #[handler]
     fn target1_handler() {
-        unsafe { &*crate::peripherals::SYSTIMER::PTR }
+        lock(&INT_ENA_LOCK, || {
-            .int_ena()
+            unsafe { &*crate::peripherals::SYSTIMER::PTR }
-            .modify(|_, w| w.target1().clear_bit());
+                .int_ena()
+                .modify(|_, w| w.target1().clear_bit());
+        });
 
         WAKERS[1].wake();
     }
 
     #[handler]
     fn target2_handler() {
-        unsafe { &*crate::peripherals::SYSTIMER::PTR }
+        lock(&INT_ENA_LOCK, || {
-            .int_ena()
+            unsafe { &*crate::peripherals::SYSTIMER::PTR }
-            .modify(|_, w| w.target2().clear_bit());
+                .int_ena()
+                .modify(|_, w| w.target2().clear_bit());
+        });
 
         WAKERS[2].wake();
     }

esp-hal/src/timer/timg.rs

@@ -77,6 +77,7 @@ use crate::soc::constants::TIMG_DEFAULT_CLK_SRC;
 use crate::{
     clock::Clocks,
     interrupt::{self, InterruptHandler},
+    lock,
     peripheral::{Peripheral, PeripheralRef},
     peripherals::{timg0::RegisterBlock, Interrupt, TIMG0},
     private::Sealed,
@@ -84,9 +85,12 @@ use crate::{
     Async,
     Blocking,
     InterruptConfigurable,
+    LockState,
     Mode,
 };
 
+static INT_ENA_LOCK: LockState = LockState::new();
+
 /// A timer group consisting of up to 2 timers (chip dependent) and a watchdog
 /// timer.
 pub struct TimerGroup<'d, T, DM>
@@ -481,9 +485,11 @@ where
             .config()
             .modify(|_, w| w.level_int_en().set_bit());
 
-        self.register_block()
+        lock(&INT_ENA_LOCK, || {
-            .int_ena()
+            self.register_block()
-            .modify(|_, w| w.t(self.timer_number()).bit(state));
+                .int_ena()
+                .modify(|_, w| w.t(self.timer_number()).bit(state));
+        });
     }
 
     fn clear_interrupt(&self) {
@@ -706,15 +712,19 @@ where
             .config()
             .modify(|_, w| w.level_int_en().set_bit());
 
-        self.register_block()
+        lock(&INT_ENA_LOCK, || {
-            .int_ena()
+            self.register_block()
-            .modify(|_, w| w.t(T).set_bit());
+                .int_ena()
+                .modify(|_, w| w.t(T).set_bit());
+        });
     }
 
     fn unlisten(&self) {
-        self.register_block()
+        lock(&INT_ENA_LOCK, || {
-            .int_ena()
+            self.register_block()
-            .modify(|_, w| w.t(T).clear_bit());
+                .int_ena()
+                .modify(|_, w| w.t(T).clear_bit());
+        });
     }
 
     fn clear_interrupt(&self) {