Fix us waits in embassy hanging (#906)

* Remove unused field from AlarmState

* Clean up non-idiomatic code

* Simplify timg0 time driver

* Better explain instant firing, hide fn, re-add check

* Add changelog entry

CHANGELOG.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - ESP32-C2/C3 examples: fix build error (#899)
 - ESP32-S3: Fix GPIO interrupt handler crashing when using GPIO48. (#898)
+- Fixed short wait times in embassy causing hangs (#906)
 
 ### Removed
 

esp-hal-common/src/embassy/mod.rs

@@ -101,7 +101,6 @@ pub fn init(clocks: &Clocks, td: time_driver::TimerType) {
 }
 
 pub struct AlarmState {
-    pub timestamp: Cell<u64>,
     pub callback: Cell<Option<(fn(*mut ()), *mut ())>>,
     pub allocated: Cell<bool>,
 }
@@ -111,7 +110,6 @@ unsafe impl Send for AlarmState {}
 impl AlarmState {
     pub const fn new() -> Self {
         Self {
-            timestamp: Cell::new(0),
             callback: Cell::new(None),
             allocated: Cell::new(false),
         }
@@ -124,18 +122,17 @@ impl Driver for EmbassyTimer {
     }
 
     unsafe fn allocate_alarm(&self) -> Option<AlarmHandle> {
-        return critical_section::with(|cs| {
+        critical_section::with(|cs| {
-            let alarms = self.alarms.borrow(cs);
+            for (i, alarm) in self.alarms.borrow(cs).iter().enumerate() {
-            for i in 0..time_driver::ALARM_COUNT {
+                if !alarm.allocated.get() {
-                let c = alarms.get_unchecked(i);
-                if !c.allocated.get() {
                     // set alarm so it is not overwritten
-                    c.allocated.set(true);
+                    alarm.allocated.set(true);
-                    return Option::Some(AlarmHandle::new(i as u8));
+                    self.on_alarm_allocated(i);
+                    return Some(AlarmHandle::new(i as u8));
                 }
             }
-            return Option::None;
+            None
-        });
+        })
     }
 
     fn set_alarm_callback(

esp-hal-common/src/embassy/time_driver_systimer.rs

@@ -32,15 +32,23 @@ impl EmbassyTimer {
         SystemTimer::now()
     }
 
-    pub(crate) fn trigger_alarm(&self, n: usize, cs: CriticalSection) {
+    fn trigger_alarm(&self, n: usize, cs: CriticalSection) {
         let alarm = &self.alarms.borrow(cs)[n];
-        alarm.timestamp.set(u64::MAX);
 
         if let Some((f, ctx)) = alarm.callback.get() {
             f(ctx);
         }
     }
 
+    pub(super) fn on_alarm_allocated(&self, n: usize) {
+        match n {
+            0 => self.alarm0.enable_interrupt(true),
+            1 => self.alarm1.enable_interrupt(true),
+            2 => self.alarm2.enable_interrupt(true),
+            _ => {}
+        }
+    }
+
     fn on_interrupt(&self, id: usize) {
         critical_section::with(|cs| {
             self.clear_interrupt(id);
@@ -83,17 +91,20 @@ impl EmbassyTimer {
         alarm: embassy_time::driver::AlarmHandle,
         timestamp: u64,
     ) -> bool {
-        critical_section::with(|cs| {
+        critical_section::with(|_cs| {
             let n = alarm.id() as usize;
-            let alarm_state = &self.alarms.borrow(cs)[n];
 
             // The hardware fires the alarm even if timestamp is lower than the current
-            // time.
+            // time. In this case the interrupt handler will pend a wakeup when we exit the
-            alarm_state.timestamp.set(timestamp);
+            // critical section.
             self.arm(n, timestamp);
+        });
 
-            true
+        // In theory, the above comment is true. However, in practice, we seem to be
-        })
+        // missing interrupt for very short timeouts, so let's make sure and catch
+        // timestamps that already passed. Returning `false` means embassy will
+        // run one more poll loop.
+        Self::now() < timestamp
     }
 
     fn clear_interrupt(&self, id: usize) {
@@ -107,18 +118,9 @@ impl EmbassyTimer {
 
     fn arm(&self, id: usize, timestamp: u64) {
         match id {
-            0 => {
+            0 => self.alarm0.set_target(timestamp),
-                self.alarm0.set_target(timestamp);
+            1 => self.alarm1.set_target(timestamp),
-                self.alarm0.enable_interrupt(true);
+            2 => self.alarm2.set_target(timestamp),
-            }
-            1 => {
-                self.alarm1.set_target(timestamp);
-                self.alarm1.enable_interrupt(true);
-            }
-            2 => {
-                self.alarm2.set_target(timestamp);
-                self.alarm2.enable_interrupt(true);
-            }
             _ => {}
         }
     }

esp-hal-common/src/embassy/time_driver_timg.rs

@@ -33,15 +33,16 @@ impl EmbassyTimer {
         unsafe { Timer0::<TIMG0>::steal() }.now()
     }
 
-    pub(crate) fn trigger_alarm(&self, n: usize, cs: CriticalSection) {
+    fn trigger_alarm(&self, n: usize, cs: CriticalSection) {
         let alarm = &self.alarms.borrow(cs)[n];
-        alarm.timestamp.set(u64::MAX);
 
         if let Some((f, ctx)) = alarm.callback.get() {
             f(ctx);
         }
     }
 
+    pub(super) fn on_alarm_allocated(&self, _n: usize) {}
+
     fn on_interrupt<Timer: Instance>(&self, id: u8, mut timer: Timer) {
         critical_section::with(|cs| {
             timer.clear_interrupt();
@@ -89,29 +90,19 @@ impl EmbassyTimer {
         timestamp: u64,
     ) -> bool {
         critical_section::with(|cs| {
-            let n = alarm.id() as usize;
-            let alarm_state = &self.alarms.borrow(cs)[n];
-
-            #[cfg(any(esp32, esp32s2, esp32s3))]
-            if n == 1 {
-                let tg = unsafe { Timer1::<TIMG0>::steal() };
-                return Self::set_alarm_impl(tg, alarm_state, timestamp);
-            }
-
-            let tg = unsafe { Timer0::<TIMG0>::steal() };
-            Self::set_alarm_impl(tg, alarm_state, timestamp)
-        })
-    }
-
-    fn set_alarm_impl<Timer: Instance>(
-        mut tg: Timer,
-        alarm_state: &AlarmState,
-        timestamp: u64,
-    ) -> bool {
             // The hardware fires the alarm even if timestamp is lower than the current
-        // time.
+            // time. In this case the interrupt handler will pend a wakeup when we exit the
-        alarm_state.timestamp.set(timestamp);
+            // critical section.
+            #[cfg(any(esp32, esp32s2, esp32s3))]
+            if alarm.id() == 1 {
+                let mut tg = unsafe { Timer1::<TIMG0>::steal() };
                 Self::arm(&mut tg, timestamp);
+                return;
+            }
+
+            let mut tg = unsafe { Timer0::<TIMG0>::steal() };
+            Self::arm(&mut tg, timestamp);
+        });
 
         true
     }