Track the timer item's owner

esp-hal-embassy/Cargo.toml

@@ -17,7 +17,7 @@ features       = ["esp32c6"]
 critical-section          = "1.2.0"
 defmt                     = { version = "0.3.8", optional = true }
 document-features         = "0.2.10"
-embassy-executor          = { version = "0.6.3", optional = true }
+embassy-executor          = { version = "0.6.3", features = ["timer-item-payload-size-4"], optional = true }
 embassy-sync              = { version = "0.6.1" }
 embassy-time              = { version = "0.3.0" }
 embassy-time-driver       = { version = "0.1.0", features = [ "tick-hz-1_000_000" ] }
@@ -51,9 +51,9 @@ defmt = ["dep:defmt", "embassy-executor?/defmt", "esp-hal/defmt"]
 log = ["dep:log"]
 ## Provide `Executor` and `InterruptExecutor`
 executors = ["dep:embassy-executor", "esp-hal/__esp_hal_embassy"]
-## Use the executor-integrated `embassy-time` timer queue. If not set, the crate provides a generic
-## timer queue that can be used with any executor.
-integrated-timers = ["embassy-time-queue-driver/integrated-timers", "executors"]
+## Use a single generic timer queue that can be used with any executor. If not set, the crate
+## provides an executor-integrated timer queue which does not need a set capacity.
+generic-queue = ["embassy-time-queue-driver/_generic-queue"]
 ## Use a single, global timer queue. This option only needs a single alarm, no matter how many
 ## executors are used. Ignored if `integrated-timers` is not set.
 single-queue = []

esp-hal-embassy/src/time_driver.rs

@@ -254,12 +254,48 @@ impl Driver for EmbassyTimer {
     fn schedule_wake(&self, at: u64, waker: &core::task::Waker) {
         #[cfg(not(single_queue))]
         unsafe {
+            // If we have multiple queues, we have integrated timers and our own timer queue
+            // implementation.
+            use embassy_executor::raw::{Executor as RawExecutor, TaskRef};
+            use portable_atomic::{AtomicPtr, Ordering};
+
             let task = embassy_executor::raw::task_from_waker(waker);
+
+            let mut executor = unsafe {
+                // SAFETY: it is impossible to schedule a task that has not yet been spawned,
+                // so the executor is guaranteed to be set to a non-null value.
+                task.executor().unwrap_unchecked() as *const RawExecutor
+            };
+
+            let owner = task
+                .timer_queue_item()
+                .payload
+                .as_ref::<AtomicPtr<RawExecutor>>();
+
+            // Try to take ownership over the timer item.
+            let owner = owner.compare_exchange(
+                core::ptr::null_mut(),
+                executor.cast_mut(),
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            );
+
+            // We can't take ownership, but we may still be able to enqueue the task. Point
+            // at the current owner.
+            if let Err(owner) = owner {
+                executor = owner;
+            };
+
+            // It is possible that the task's owner changes in the mean time. It doesn't
+            // matter, at this point the only interesting question is: can we enqueue in the
+            // currently loaded owner's timer queue?
+
+            // Try to enqueue in the current owner's timer queue. This will fail if the
+            // owner has a lower priority ceiling than the current context.
+
             // FIXME: this is UB, use Exposed Provenance API (or something better) when
             // available. Expose provenance in `InnerExecutor::init`, and use it here.
-            let executor = &*(task.executor().unwrap_unchecked()
-                as *const embassy_executor::raw::Executor)
-                .cast::<crate::executor::InnerExecutor>();
+            let executor = &*(executor.cast::<crate::executor::InnerExecutor>());
             executor.timer_queue.schedule_wake(at, waker);
         }
 

qa-test/src/bin/embassy_executor_benchmark.rs

@@ -1,9 +1,9 @@
 //! Embassy executor benchmark, used to try out optimization ideas.
 
 //% CHIPS: esp32c2 esp32c3 esp32c6 esp32h2 esp32s2 esp32s3
-// FEATURES:
+//% FEATURES:
 // FEATURES: esp-hal-embassy/single-queue
-//% FEATURES: esp-hal-embassy/generic-queue
+// FEATURES: esp-hal-embassy/generic-queue
 
 #![no_std]
 #![no_main]