4 files changed, 98 insertions, 83 deletions

diff --git a/embassy-stm32/src/adc/v3.rs b/embassy-stm32/src/adc/v3.rs
index 62c2da557..170b08a25 100644
--- a/embassy-stm32/src/adc/v3.rs
+++ b/embassy-stm32/src/adc/v3.rs
@@ -463,7 +463,7 @@ impl<'d, T: Instance> Adc<'d, T> {
         );
 
         #[cfg(all(feature = "low-power", stm32wlex))]
-        let _device_busy = crate::low_power::DeviceBusy::new();
+        let _device_busy = crate::low_power::DeviceBusy::new_stop1();
 
         // Ensure no conversions are ongoing and ADC is enabled.
         Self::cancel_conversions();
diff --git a/embassy-stm32/src/i2c/v2.rs b/embassy-stm32/src/i2c/v2.rs
index 6f2d03bd1..978a401d5 100644
--- a/embassy-stm32/src/i2c/v2.rs
+++ b/embassy-stm32/src/i2c/v2.rs
@@ -820,7 +820,7 @@ impl<'d, IM: MasterMode> I2c<'d, Async, IM> {
     /// Write.
     pub async fn write(&mut self, address: u8, write: &[u8]) -> Result<(), Error> {
         #[cfg(all(feature = "low-power", stm32wlex))]
-        let _device_busy = crate::low_power::DeviceBusy::new();
+        let _device_busy = crate::low_power::DeviceBusy::new_stop1();
         let timeout = self.timeout();
         if write.is_empty() {
             self.write_internal(address.into(), write, true, timeout)
@@ -836,7 +836,7 @@ impl<'d, IM: MasterMode> I2c<'d, Async, IM> {
     /// The buffers are concatenated in a single write transaction.
     pub async fn write_vectored(&mut self, address: Address, write: &[&[u8]]) -> Result<(), Error> {
         #[cfg(all(feature = "low-power", stm32wlex))]
-        let _device_busy = crate::low_power::DeviceBusy::new();
+        let _device_busy = crate::low_power::DeviceBusy::new_stop1();
         let timeout = self.timeout();
 
         if write.is_empty() {
@@ -861,7 +861,7 @@ impl<'d, IM: MasterMode> I2c<'d, Async, IM> {
     /// Read.
     pub async fn read(&mut self, address: u8, buffer: &mut [u8]) -> Result<(), Error> {
         #[cfg(all(feature = "low-power", stm32wlex))]
-        let _device_busy = crate::low_power::DeviceBusy::new();
+        let _device_busy = crate::low_power::DeviceBusy::new_stop1();
         let timeout = self.timeout();
 
         if buffer.is_empty() {
@@ -875,7 +875,7 @@ impl<'d, IM: MasterMode> I2c<'d, Async, IM> {
     /// Write, restart, read.
     pub async fn write_read(&mut self, address: u8, write: &[u8], read: &mut [u8]) -> Result<(), Error> {
         #[cfg(all(feature = "low-power", stm32wlex))]
-        let _device_busy = crate::low_power::DeviceBusy::new();
+        let _device_busy = crate::low_power::DeviceBusy::new_stop1();
         let timeout = self.timeout();
 
         if write.is_empty() {
@@ -902,7 +902,7 @@ impl<'d, IM: MasterMode> I2c<'d, Async, IM> {
     /// [transaction contract]: embedded_hal_1::i2c::I2c::transaction
     pub async fn transaction(&mut self, addr: u8, operations: &mut [Operation<'_>]) -> Result<(), Error> {
         #[cfg(all(feature = "low-power", stm32wlex))]
-        let _device_busy = crate::low_power::DeviceBusy::new();
+        let _device_busy = crate::low_power::DeviceBusy::new_stop1();
         let _ = addr;
         let _ = operations;
         todo!()
diff --git a/embassy-stm32/src/low_power.rs b/embassy-stm32/src/low_power.rs
index cde3153f6..938db2686 100644
--- a/embassy-stm32/src/low_power.rs
+++ b/embassy-stm32/src/low_power.rs
@@ -57,10 +57,11 @@ use core::marker::PhantomData;
 use core::sync::atomic::{Ordering, compiler_fence};
 
 use cortex_m::peripheral::SCB;
+use critical_section::CriticalSection;
 use embassy_executor::*;
 
 use crate::interrupt;
-use crate::time_driver::{RtcDriver, get_driver};
+use crate::time_driver::get_driver;
 
 const THREAD_PENDER: usize = usize::MAX;
 
@@ -68,46 +69,59 @@ use crate::rtc::Rtc;
 
 static mut EXECUTOR: Option<Executor> = None;
 
-#[cfg(stm32wlex)]
-pub(crate) use self::busy::DeviceBusy;
-#[cfg(stm32wlex)]
-mod busy {
-    use core::sync::atomic::{AtomicU32, Ordering};
+/// Prevent the device from going into the stop mode if held
+pub struct DeviceBusy(StopMode);
 
-    // Count of devices blocking STOP
-    static STOP_BLOCKED: AtomicU32 = AtomicU32::new(0);
+impl DeviceBusy {
+    /// Create a new DeviceBusy with stop1.
+    pub fn new_stop1() -> Self {
+        Self::new(StopMode::Stop1)
+    }
 
-    /// Check if STOP1 is blocked.
-    pub(crate) fn stop_blocked() -> bool {
-        STOP_BLOCKED.load(Ordering::SeqCst) > 0
+    /// Create a new DeviceBusy with stop2.
+    pub fn new_stop2() -> Self {
+        Self::new(StopMode::Stop2)
     }
 
-    /// When ca device goes busy it will construct one of these where it will be dropped when the device goes idle.
-    pub(crate) struct DeviceBusy {}
+    /// Create a new DeviceBusy.
+    pub fn new(stop_mode: StopMode) -> Self {
+        critical_section::with(|_| unsafe {
+            match stop_mode {
+                StopMode::Stop1 => {
+                    crate::rcc::REFCOUNT_STOP1 += 1;
+                }
+                StopMode::Stop2 => {
+                    crate::rcc::REFCOUNT_STOP2 += 1;
+                }
+            }
+        });
 
-    impl DeviceBusy {
-        /// Create a new DeviceBusy.
-        pub(crate) fn new() -> Self {
-            STOP_BLOCKED.fetch_add(1, Ordering::SeqCst);
-            trace!("low power: device busy: Stop:{}", STOP_BLOCKED.load(Ordering::SeqCst));
-            Self {}
-        }
+        Self(stop_mode)
     }
+}
 
-    impl Drop for DeviceBusy {
-        fn drop(&mut self) {
-            STOP_BLOCKED.fetch_sub(1, Ordering::SeqCst);
-            trace!("low power: device idle: Stop:{}", STOP_BLOCKED.load(Ordering::SeqCst));
-        }
+impl Drop for DeviceBusy {
+    fn drop(&mut self) {
+        critical_section::with(|_| unsafe {
+            match self.0 {
+                StopMode::Stop1 => {
+                    crate::rcc::REFCOUNT_STOP1 -= 1;
+                }
+                StopMode::Stop2 => {
+                    crate::rcc::REFCOUNT_STOP2 -= 1;
+                }
+            }
+        });
     }
 }
+
 #[cfg(not(stm32u0))]
 foreach_interrupt! {
     (RTC, rtc, $block:ident, WKUP, $irq:ident) => {
         #[interrupt]
         #[allow(non_snake_case)]
         unsafe fn $irq() {
-            EXECUTOR.as_mut().unwrap().on_wakeup_irq();
+            Executor::on_wakeup_irq();
         }
     };
 }
@@ -118,27 +132,28 @@ foreach_interrupt! {
         #[interrupt]
         #[allow(non_snake_case)]
         unsafe fn $irq() {
-            EXECUTOR.as_mut().unwrap().on_wakeup_irq();
+            Executor::on_wakeup_irq();
         }
     };
 }
 
 #[allow(dead_code)]
 pub(crate) unsafe fn on_wakeup_irq() {
-    if EXECUTOR.is_some() {
-        trace!("low power: wakeup irq");
-        EXECUTOR.as_mut().unwrap().on_wakeup_irq();
-    }
+    Executor::on_wakeup_irq();
 }
 
 /// Configure STOP mode with RTC.
 pub fn stop_with_rtc(rtc: Rtc) {
-    unsafe { EXECUTOR.as_mut().unwrap() }.stop_with_rtc(rtc)
+    assert!(unsafe { EXECUTOR.is_some() });
+
+    Executor::stop_with_rtc(rtc)
 }
 
 /// Reconfigure the RTC, if set.
 pub fn reconfigure_rtc(f: impl FnOnce(&mut Rtc)) {
-    unsafe { EXECUTOR.as_mut().unwrap() }.reconfigure_rtc(f);
+    assert!(unsafe { EXECUTOR.is_some() });
+
+    Executor::reconfigure_rtc(f);
 }
 
 /// Get whether the core is ready to enter the given stop mode.
@@ -146,11 +161,11 @@ pub fn reconfigure_rtc(f: impl FnOnce(&mut Rtc)) {
 /// This will return false if some peripheral driver is in use that
 /// prevents entering the given stop mode.
 pub fn stop_ready(stop_mode: StopMode) -> bool {
-    match unsafe { EXECUTOR.as_mut().unwrap() }.stop_mode() {
+    critical_section::with(|cs| match Executor::stop_mode(cs) {
         Some(StopMode::Stop2) => true,
         Some(StopMode::Stop1) => stop_mode == StopMode::Stop1,
         None => false,
-    }
+    })
 }
 
 /// Available Stop modes.
@@ -193,7 +208,6 @@ pub struct Executor {
     inner: raw::Executor,
     not_send: PhantomData<*mut ()>,
     scb: SCB,
-    time_driver: &'static RtcDriver,
 }
 
 impl Executor {
@@ -206,7 +220,6 @@ impl Executor {
                 inner: raw::Executor::new(THREAD_PENDER as *mut ()),
                 not_send: PhantomData,
                 scb: cortex_m::Peripherals::steal().SCB,
-                time_driver: get_driver(),
             });
 
             let executor = EXECUTOR.as_mut().unwrap();
@@ -215,54 +228,50 @@ impl Executor {
         })
     }
 
-    unsafe fn on_wakeup_irq(&mut self) {
-        #[cfg(stm32wlex)]
-        {
-            let extscr = crate::pac::PWR.extscr().read();
-            if extscr.c1stop2f() || extscr.c1stopf() {
-                // when we wake from any stop mode we need to re-initialize the rcc
-                crate::rcc::apply_resume_config();
-                if extscr.c1stop2f() {
-                    // when we wake from STOP2, we need to re-initialize the time driver
-                    critical_section::with(|cs| crate::time_driver::init_timer(cs));
-                    // reset the refcounts for STOP2 and STOP1 (initializing the time driver will increment one of them for the timer)
-                    // and given that we just woke from STOP2, we can reset them
-                    crate::rcc::REFCOUNT_STOP2 = 0;
-                    crate::rcc::REFCOUNT_STOP1 = 0;
+    pub(self) unsafe fn on_wakeup_irq() {
+        critical_section::with(|cs| {
+            if !get_driver().is_rtc_set(cs) {
+                trace!("low power: wakeup irq; rtc not set");
+
+                return;
+            }
+
+            #[cfg(stm32wlex)]
+            {
+                let extscr = crate::pac::PWR.extscr().read();
+                if extscr.c1stop2f() || extscr.c1stopf() {
+                    // when we wake from any stop mode we need to re-initialize the rcc
+                    crate::rcc::apply_resume_config();
+                    if extscr.c1stop2f() {
+                        // when we wake from STOP2, we need to re-initialize the time driver
+                        crate::time_driver::init_timer(cs);
+                        // reset the refcounts for STOP2 and STOP1 (initializing the time driver will increment one of them for the timer)
+                        // and given that we just woke from STOP2, we can reset them
+                        crate::rcc::REFCOUNT_STOP2 = 0;
+                        crate::rcc::REFCOUNT_STOP1 = 0;
+                    }
                 }
             }
-        }
-        self.time_driver.resume_time();
-        trace!("low power: resume");
+            get_driver().resume_time();
+            trace!("low power: resume");
+        });
     }
 
-    pub(self) fn stop_with_rtc(&mut self, rtc: Rtc) {
-        self.time_driver.set_rtc(rtc);
-        self.time_driver.reconfigure_rtc(|rtc| rtc.enable_wakeup_line());
+    pub(self) fn stop_with_rtc(rtc: Rtc) {
+        get_driver().set_rtc(rtc);
+        get_driver().reconfigure_rtc(|rtc| rtc.enable_wakeup_line());
 
         trace!("low power: stop with rtc configured");
     }
 
-    pub(self) fn reconfigure_rtc(&mut self, f: impl FnOnce(&mut Rtc)) {
-        self.time_driver.reconfigure_rtc(f);
-    }
-
-    fn stop1_ok_to_enter(&self) -> bool {
-        #[cfg(stm32wlex)]
-        if self::busy::stop_blocked() {
-            return false;
-        }
-        unsafe { crate::rcc::REFCOUNT_STOP1 == 0 }
-    }
-
-    fn stop2_ok_to_enter(&self) -> bool {
-        self.stop1_ok_to_enter() && unsafe { crate::rcc::REFCOUNT_STOP2 == 0 }
+    pub(self) fn reconfigure_rtc(f: impl FnOnce(&mut Rtc)) {
+        get_driver().reconfigure_rtc(f);
     }
 
-    fn stop_mode(&self) -> Option<StopMode> {
-        if self.stop2_ok_to_enter() {
+    fn stop_mode(_cs: CriticalSection) -> Option<StopMode> {
+        if unsafe { crate::rcc::REFCOUNT_STOP2 == 0 && crate::rcc::REFCOUNT_STOP1 == 0 } {
             Some(StopMode::Stop2)
-        } else if self.stop1_ok_to_enter() {
+        } else if unsafe { crate::rcc::REFCOUNT_STOP1 == 0 } {
             Some(StopMode::Stop1)
         } else {
             None
@@ -291,14 +300,14 @@ impl Executor {
 
         compiler_fence(Ordering::SeqCst);
 
-        let stop_mode = self.stop_mode();
+        let stop_mode = critical_section::with(|cs| Self::stop_mode(cs));
 
         if stop_mode.is_none() {
             trace!("low power: not ready to stop");
             return;
         }
 
-        if self.time_driver.pause_time().is_err() {
+        if get_driver().pause_time().is_err() {
             trace!("low power: failed to pause time");
             return;
         }
diff --git a/embassy-stm32/src/time_driver.rs b/embassy-stm32/src/time_driver.rs
index 4956d1f68..152d3d22e 100644
--- a/embassy-stm32/src/time_driver.rs
+++ b/embassy-stm32/src/time_driver.rs
@@ -413,6 +413,12 @@ impl RtcDriver {
     }
 
     #[cfg(feature = "low-power")]
+    /// Has the rtc been set?
+    pub(crate) fn is_rtc_set(&self, cs: CriticalSection) -> bool {
+        self.rtc.borrow(cs).borrow().is_some()
+    }
+
+    #[cfg(feature = "low-power")]
     /// Set the rtc but panic if it's already been set
     pub(crate) fn reconfigure_rtc(&self, f: impl FnOnce(&mut Rtc)) {
         critical_section::with(|cs| f(self.rtc.borrow(cs).borrow_mut().as_mut().unwrap()));
@@ -543,7 +549,7 @@ impl Driver for RtcDriver {
 }
 
 #[cfg(feature = "low-power")]
-pub(crate) fn get_driver() -> &'static RtcDriver {
+pub(crate) const fn get_driver() -> &'static RtcDriver {
     &DRIVER
 }