Rewrite time driver (#2559)

* Rewrite time driver

* Don't store priority

* Changelog

* Fix doc example

* Use separate locks for Alarms

* Mention generic queues

* Cache the next wakeup timestamp

* Immediately repoll if timestamp is in the past

* Add benchmark

* Remove equality check for now

* Enable interrupts when allocating the alarm

* Clean up

* Use relaxed ordering

* wut

* Typo

* Move benchmar

* fmt

.cargo/config.toml

@@ -1,3 +1,4 @@
 [alias]
 xtask = "run --package xtask --"
 xfmt = "xtask fmt-packages"
+qa = "xtask run-example qa-test"

esp-hal-embassy/CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 - Reduce memory footprint by 4 bytes on multi-core MCUs.
+- The time driver no longer uses cross-core critical sections. (#2559)
 
 ### Fixed
 

esp-hal-embassy/src/executor/thread.rs

@@ -30,7 +30,7 @@ pub(crate) fn pend_thread_mode(_core: usize) {
     #[cfg(low_power_wait)]
     {
         // Signal that there is work to be done.
-        SIGNAL_WORK_THREAD_MODE[_core].store(true, Ordering::SeqCst);
+        SIGNAL_WORK_THREAD_MODE[_core].store(true, Ordering::Relaxed);
 
         // If we are pending a task on the current core, we're done. Otherwise, we
         // need to make sure the other core wakes up.
@@ -123,7 +123,8 @@ This will use software-interrupt 3 which isn't available for anything else to wa
 
         // we do not care about race conditions between the load and store operations,
         // interrupts will only set this value to true.
-        if SIGNAL_WORK_THREAD_MODE[cpu].load(Ordering::SeqCst) {
+        // Acquire makes no sense but at this time it's slightly faster than Relaxed.
+        if SIGNAL_WORK_THREAD_MODE[cpu].load(Ordering::Acquire) {
             // if there is work to do, exit critical section and loop back to polling
             unsafe {
                 core::arch::asm!(
@@ -140,7 +141,7 @@ This will use software-interrupt 3 which isn't available for anything else to wa
             unsafe { core::arch::asm!("waiti 0") };
         }
         // If this races and some waker sets the signal, we'll reset it, but still poll.
-        SIGNAL_WORK_THREAD_MODE[cpu].store(false, Ordering::SeqCst);
+        SIGNAL_WORK_THREAD_MODE[cpu].store(false, Ordering::Relaxed);
     }
 
     #[cfg(all(riscv, low_power_wait))]
@@ -149,14 +150,14 @@ This will use software-interrupt 3 which isn't available for anything else to wa
         // interrupts will only set this value to true.
         critical_section::with(|_| {
             // if there is work to do, loop back to polling
-            if !SIGNAL_WORK_THREAD_MODE[cpu].load(Ordering::SeqCst) {
+            if !SIGNAL_WORK_THREAD_MODE[cpu].load(Ordering::Relaxed) {
                 // if not, wait for interrupt
                 unsafe { core::arch::asm!("wfi") };
             }
         });
         // if an interrupt occurred while waiting, it will be serviced here
         // If this races and some waker sets the signal, we'll reset it, but still poll.
-        SIGNAL_WORK_THREAD_MODE[cpu].store(false, Ordering::SeqCst);
+        SIGNAL_WORK_THREAD_MODE[cpu].store(false, Ordering::Relaxed);
     }
 }
 

esp-hal-embassy/src/lib.rs

@@ -146,7 +146,7 @@ impl_array!(4);
 ///
 /// ```rust, no_run
 #[doc = esp_hal::before_snippet!()]
-/// use esp_hal::timg::TimerGroup;
+/// use esp_hal::timer::timg::TimerGroup;
 ///
 /// let timg0 = TimerGroup::new(peripherals.TIMG0);
 /// esp_hal_embassy::init(timg0.timer0);

esp-hal-embassy/src/time_driver.rs

@@ -1,94 +1,163 @@
-use core::cell::{Cell, RefCell};
+use core::cell::Cell;
 
-use critical_section::Mutex;
 use embassy_time_driver::{AlarmHandle, Driver};
 use esp_hal::{
     interrupt::{InterruptHandler, Priority},
     prelude::*,
+    sync::Locked,
     time::now,
     timer::{AnyTimer, OneShotTimer},
 };
 
-pub const MAX_SUPPORTED_ALARM_COUNT: usize = 7;
-
 pub type Timer = OneShotTimer<'static, AnyTimer>;
 
-static TIMERS: Mutex<RefCell<Option<&'static mut [Timer]>>> = Mutex::new(RefCell::new(None));
-
-#[allow(clippy::type_complexity)]
-struct AlarmState {
-    pub callback: Cell<Option<(fn(*mut ()), *mut ())>>,
-    pub allocated: Cell<bool>,
+enum AlarmState {
+    Created(extern "C" fn()),
+    Allocated(extern "C" fn()),
+    Initialized(&'static mut Timer),
+}
+impl AlarmState {
+    fn initialize(
+        timer: &'static mut OneShotTimer<'_, AnyTimer>,
+        interrupt_handler: extern "C" fn(),
+    ) -> AlarmState {
+        // If the driver is initialized, bind the interrupt handler to the
+        // timer. This ensures that alarms allocated after init are correctly
+        // bound to the core that created the executor.
+        timer.set_interrupt_handler(InterruptHandler::new(interrupt_handler, Priority::max()));
+        timer.enable_interrupt(true);
+        AlarmState::Initialized(timer)
+    }
 }
 
-unsafe impl Send for AlarmState {}
+struct AlarmInner {
+    pub callback: Cell<(*const (), *mut ())>,
+    pub state: AlarmState,
+}
 
-impl AlarmState {
-    pub const fn new() -> Self {
+struct Alarm {
+    pub inner: Locked<AlarmInner>,
+}
+
+unsafe impl Send for Alarm {}
+
+impl Alarm {
+    pub const fn new(handler: extern "C" fn()) -> Self {
         Self {
-            callback: Cell::new(None),
-            allocated: Cell::new(false),
+            inner: Locked::new(AlarmInner {
+                callback: Cell::new((core::ptr::null(), core::ptr::null_mut())),
+                state: AlarmState::Created(handler),
+            }),
         }
     }
 }
 
 pub(super) struct EmbassyTimer {
-    alarms: Mutex<[AlarmState; MAX_SUPPORTED_ALARM_COUNT]>,
+    alarms: [Alarm; MAX_SUPPORTED_ALARM_COUNT],
+    available_timers: Locked<Option<&'static mut [Timer]>>,
 }
 
+/// Repeats the `Alarm::new` constructor for each alarm, creating an interrupt
+/// handler for each of them.
+macro_rules! alarms {
+    ($($idx:literal),*) => {
+        [$(
+            Alarm::new({
+                // Not #[handler] so we don't have to store the priority - which is constant.
+                extern "C" fn handler() {
+                    DRIVER.on_interrupt($idx);
+                }
+                handler
+            })
+        ),*]
+    };
+}
+
+const MAX_SUPPORTED_ALARM_COUNT: usize = 7;
 embassy_time_driver::time_driver_impl!(static DRIVER: EmbassyTimer = EmbassyTimer {
-    alarms: Mutex::new([const { AlarmState::new() }; MAX_SUPPORTED_ALARM_COUNT]),
+    alarms: alarms!(0, 1, 2, 3, 4, 5, 6),
+    available_timers: Locked::new(None),
 });
 
 impl EmbassyTimer {
-    pub(super) fn init(timers: &'static mut [Timer]) {
-        if timers.len() > MAX_SUPPORTED_ALARM_COUNT {
-            panic!(
-                "Maximum of {} timers can be used.",
-                MAX_SUPPORTED_ALARM_COUNT
-            );
-        }
+    pub(super) fn init(mut timers: &'static mut [Timer]) {
+        assert!(
+            timers.len() <= MAX_SUPPORTED_ALARM_COUNT,
+            "Maximum {} timers can be used.",
+            MAX_SUPPORTED_ALARM_COUNT
+        );
 
-        critical_section::with(|cs| {
-            timers.iter_mut().enumerate().for_each(|(n, timer)| {
-                timer.enable_interrupt(false);
-                timer.stop();
+        // Reset timers
+        timers.iter_mut().for_each(|timer| {
+            timer.enable_interrupt(false);
+            timer.stop();
+        });
 
-                if DRIVER.alarms.borrow(cs)[n].allocated.get() {
-                    // FIXME: we should track which core allocated an alarm and bind the interrupt
-                    // to that core.
-                    timer.set_interrupt_handler(HANDLERS[n]);
+        // Initialize already allocated timers
+        for alarm in DRIVER.alarms.iter() {
+            timers = alarm.inner.with(move |alarm| {
+                if let AlarmState::Allocated(interrupt_handler) = alarm.state {
+                    // Pluck off a timer
+
+                    let Some((timer, remaining_timers)) = timers.split_first_mut() else {
+                        not_enough_timers();
+                    };
+
+                    alarm.state = AlarmState::initialize(timer, interrupt_handler);
+
+                    remaining_timers
+                } else {
+                    timers
                 }
             });
+        }
 
-            TIMERS.replace(cs, Some(timers));
-        });
+        // Store the available timers
+        DRIVER
+            .available_timers
+            .with(|available_timers| *available_timers = Some(timers));
     }
 
     fn on_interrupt(&self, id: usize) {
-        let cb = critical_section::with(|cs| {
-            let mut timers = TIMERS.borrow_ref_mut(cs);
-            let timers = unwrap!(timers.as_mut(), "Time driver not initialized");
-            let timer = &mut timers[id];
-
-            timer.clear_interrupt();
-
-            let alarm = &self.alarms.borrow(cs)[id];
-            alarm.callback.get()
+        let (cb, ctx) = self.alarms[id].inner.with(|alarm| {
+            if let AlarmState::Initialized(timer) = &mut alarm.state {
+                timer.clear_interrupt();
+                alarm.callback.get()
+            } else {
+                unsafe {
+                    // SAFETY: `on_interrupt` is registered right when the alarm is initialized.
+                    core::hint::unreachable_unchecked()
+                }
+            }
         });
 
-        if let Some((f, ctx)) = cb {
-            f(ctx);
-        }
+        let cb: fn(*mut ()) = unsafe {
+            // Safety:
+            // - we can ignore the possibility of `f` being unset (null) because of the
+            //   safety contract of `allocate_alarm`.
+            // - other than that we only store valid function pointers into alarm.callback
+            core::mem::transmute(cb)
+        };
+
+        cb(ctx);
     }
 
-    fn arm(timer: &mut Timer, timestamp: u64) {
+    /// Returns `true` if the timer was armed, `false` if the timestamp is in
+    /// the past.
+    fn arm(timer: &mut Timer, timestamp: u64) -> bool {
         let now = now().duration_since_epoch();
         let ts = timestamp.micros();
-        // if the TS is already in the past make the timer fire immediately
-        let timeout = if ts > now { ts - now } else { 0.micros() };
-        unwrap!(timer.schedule(timeout));
-        timer.enable_interrupt(true);
+
+        if ts > now {
+            let timeout = ts - now;
+            unwrap!(timer.schedule(timeout));
+            true
+        } else {
+            // If the timestamp is past, we return `false` to ask embassy to poll again
+            // immediately.
+            timer.stop();
+            false
+        }
     }
 }
 
@@ -98,42 +167,60 @@ impl Driver for EmbassyTimer {
     }
 
     unsafe fn allocate_alarm(&self) -> Option<AlarmHandle> {
-        critical_section::with(|cs| {
-            for (i, alarm) in self.alarms.borrow(cs).iter().enumerate() {
-                if alarm.allocated.get() {
-                    continue;
-                }
-                let mut timer = TIMERS.borrow_ref_mut(cs);
-                // `allocate_alarm` may be called before `esp_hal_embassy::init()`, so
-                // we need to check if we have timers.
-                if let Some(timer) = &mut *timer {
-                    // If we do, bind the interrupt handler to the timer.
-                    // This ensures that alarms allocated after init are correctly bound to the
-                    // core that created the executor.
-                    let timer = unwrap!(
-                        timer.get_mut(i),
-                        "There are not enough timers to allocate a new alarm. Call `esp_hal_embassy::init()` with the correct number of timers."
-                    );
-                    timer.set_interrupt_handler(HANDLERS[i]);
-                }
+        for (i, alarm) in self.alarms.iter().enumerate() {
+            let handle = alarm.inner.with(|alarm| {
+                let AlarmState::Created(interrupt_handler) = alarm.state else {
+                    return None;
+                };
 
-                // set alarm so it is not overwritten
-                alarm.allocated.set(true);
-                return Some(AlarmHandle::new(i as u8));
+                let timer = self.available_timers.with(|available_timers| {
+                    // `allocate_alarm` may be called before `esp_hal_embassy::init()`. If
+                    // `timers` is `None`, we return `None` to signal that the alarm cannot be
+                    // initialized yet. These alarms will be initialized when `init` is called.
+                    if let Some(timers) = available_timers.take() {
+                        // If the driver is initialized, we can allocate a timer.
+                        // If this fails, we can't do anything about it.
+                        let Some((timer, rest)) = timers.split_first_mut() else {
+                            not_enough_timers();
+                        };
+                        *available_timers = Some(rest);
+                        Some(timer)
+                    } else {
+                        None
+                    }
+                });
+
+                alarm.state = match timer {
+                    Some(timer) => AlarmState::initialize(timer, interrupt_handler),
+
+                    None => {
+                        // No timers are available yet, mark the alarm as allocated.
+                        AlarmState::Allocated(interrupt_handler)
+                    }
+                };
+
+                Some(AlarmHandle::new(i as u8))
+            });
+
+            if handle.is_some() {
+                return handle;
             }
-            None
-        })
+        }
+
+        None
     }
 
     fn set_alarm_callback(&self, alarm: AlarmHandle, callback: fn(*mut ()), ctx: *mut ()) {
         let n = alarm.id() as usize;
-        critical_section::with(|cs| {
-            let alarm = &self.alarms.borrow(cs)[n];
-            alarm.callback.set(Some((callback, ctx)));
+
+        self.alarms[n].inner.with(|alarm| {
+            alarm.callback.set((callback as *const (), ctx));
         })
     }
 
     fn set_alarm(&self, alarm: AlarmHandle, timestamp: u64) -> bool {
+        let alarm = &self.alarms[alarm.id() as usize];
+
         // If `embassy-executor/integrated-timers` is enabled and there are no pending
         // timers, embassy still calls `set_alarm` with `u64::MAX`. By returning
         // `true` we signal that no re-polling is necessary.
@@ -148,47 +235,22 @@ impl Driver for EmbassyTimer {
         // This is correct behavior. See https://docs.rs/embassy-time-driver/0.1.0/embassy_time_driver/trait.Driver.html#tymethod.set_alarm
         // (... the driver should return true and arrange to call the alarm callback as
         // soon as possible, but not synchronously.)
-        critical_section::with(|cs| {
-            let mut timers = TIMERS.borrow_ref_mut(cs);
-            let timers = unwrap!(timers.as_mut(), "Time driver not initialized");
-            let timer = &mut timers[alarm.id() as usize];
 
-            Self::arm(timer, timestamp);
-        });
-
-        true
+        alarm.inner.with(|alarm| {
+            if let AlarmState::Initialized(timer) = &mut alarm.state {
+                Self::arm(timer, timestamp)
+            } else {
+                panic!("set_alarm called before esp_hal_embassy::init()")
+            }
+        })
     }
 }
 
-static HANDLERS: [InterruptHandler; MAX_SUPPORTED_ALARM_COUNT] = [
-    handler0, handler1, handler2, handler3, handler4, handler5, handler6,
-];
-
-#[handler(priority = Priority::max())]
-fn handler0() {
-    DRIVER.on_interrupt(0);
-}
-#[handler(priority = Priority::max())]
-fn handler1() {
-    DRIVER.on_interrupt(1);
-}
-#[handler(priority = Priority::max())]
-fn handler2() {
-    DRIVER.on_interrupt(2);
-}
-#[handler(priority = Priority::max())]
-fn handler3() {
-    DRIVER.on_interrupt(3);
-}
-#[handler(priority = Priority::max())]
-fn handler4() {
-    DRIVER.on_interrupt(4);
-}
-#[handler(priority = Priority::max())]
-fn handler5() {
-    DRIVER.on_interrupt(5);
-}
-#[handler(priority = Priority::max())]
-fn handler6() {
-    DRIVER.on_interrupt(6);
+#[cold]
+#[track_caller]
+fn not_enough_timers() -> ! {
+    // This is wrapped in a separate function because rustfmt does not like
+    // extremely long strings. Also, if log is used, this avoids storing the string
+    // twice.
+    panic!("There are not enough timers to allocate a new alarm. Call esp_hal_embassy::init() with the correct number of timers, or consider using one of the embassy-timer/generic-queue-X features.");
 }

esp-hal-procmacros/CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Functions marked with `#[handler]` can now be referenced in `const` context. (#2559)
+
 ### Removed
 
 - Removed the `enum-dispatch`, `interrupt`, and `ram` features (#2594)

esp-hal-procmacros/src/lib.rs

@@ -321,7 +321,7 @@ pub fn handler(args: TokenStream, input: TokenStream) -> TokenStream {
         #f
 
         #[allow(non_upper_case_globals)]
-        #vis static #orig: #root::interrupt::InterruptHandler = #root::interrupt::InterruptHandler::new(#new, #priority);
+        #vis const #orig: #root::interrupt::InterruptHandler = #root::interrupt::InterruptHandler::new(#new, #priority);
     )
     .into()
 }

qa-test/src/bin/embassy_executor_benchmark.rs

@@ -0,0 +1,75 @@
+//! Embassy executor benchmark, used to try out optimization ideas.
+
+//% CHIPS: esp32c2 esp32c3 esp32c6 esp32h2 esp32s2 esp32s3
+//% FEATURES: esp-hal-embassy/integrated-timers
+
+#![no_std]
+#![no_main]
+
+use core::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use embassy_executor::{raw::TaskStorage, Spawner};
+use esp_backtrace as _;
+use esp_hal::{
+    prelude::*,
+    time::Duration,
+    timer::{systimer::SystemTimer, OneShotTimer},
+};
+use esp_println::println;
+
+static mut COUNTER: u32 = 0;
+
+const CLOCK: CpuClock = CpuClock::max();
+const TEST_MILLIS: u64 = 50;
+
+#[handler]
+fn timer_handler() {
+    let c = unsafe { COUNTER } as u64;
+    let cpu_clock = CLOCK.hz() as u64;
+    let timer_ticks_per_second = SystemTimer::ticks_per_second();
+    let cpu_cycles_per_timer_ticks = cpu_clock / timer_ticks_per_second;
+    println!(
+        "Test OK, count={}, cycles={}/100",
+        c,
+        (100 * timer_ticks_per_second * cpu_cycles_per_timer_ticks * TEST_MILLIS / 1000) / c
+    );
+    loop {}
+}
+
+struct Task1 {}
+impl Future for Task1 {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        unsafe { COUNTER += 1 };
+        cx.waker().wake_by_ref();
+        Poll::Pending
+    }
+}
+
+static TASK1: TaskStorage<Task1> = TaskStorage::new();
+
+#[esp_hal_embassy::main]
+async fn main(spawner: Spawner) {
+    let peripherals = esp_hal::init({
+        let mut config = esp_hal::Config::default();
+        config.cpu_clock = CLOCK;
+        config
+    });
+    let systimer = SystemTimer::new(peripherals.SYSTIMER);
+    esp_hal_embassy::init(systimer.alarm0);
+    println!("Embassy initialized!");
+
+    spawner.spawn(TASK1.spawn(|| Task1 {})).unwrap();
+
+    println!("Starting test");
+
+    let mut timer = OneShotTimer::new(systimer.alarm1);
+    timer.set_interrupt_handler(timer_handler);
+    timer.enable_interrupt(true);
+    timer.schedule(Duration::millis(TEST_MILLIS)).unwrap();
+}