13 files changed, 80 insertions, 187 deletions

diff --git a/embassy-executor/CHANGELOG.md b/embassy-executor/CHANGELOG.md
index 068156210..59ea88d0c 100644
--- a/embassy-executor/CHANGELOG.md
+++ b/embassy-executor/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - embassy-executor no longer provides an `embassy-time-queue-driver` implementation
 - Added `TaskRef::executor` to obtain a reference to a task's executor
 - integrated-timers are no longer processed when polling the executor.
+- Added the option to store data in timer queue items
 
 ## 0.6.3 - 2024-11-12
 
diff --git a/embassy-executor/Cargo.toml b/embassy-executor/Cargo.toml
index 60fe7087a..2a64b9c83 100644
--- a/embassy-executor/Cargo.toml
+++ b/embassy-executor/Cargo.toml
@@ -92,6 +92,22 @@ trace = []
 ## Enable support for rtos-trace framework
 rtos-trace = ["dep:rtos-trace", "trace", "dep:embassy-time-driver"]
 
+#! ### Timer Item Payload Size
+#! Sets the size of the payload for timer items, allowing integrated timer implementors to store
+#! additional data in the timer item. The payload field will be aligned to this value as well.
+#! If these features are not defined, the timer item will contain no payload field.
+
+_timer-item-payload = [] # A size was picked
+
+## 1 bytes
+timer-item-payload-size-1 = ["_timer-item-payload"]
+## 2 bytes
+timer-item-payload-size-2 = ["_timer-item-payload"]
+## 4 bytes
+timer-item-payload-size-4 = ["_timer-item-payload"]
+## 8 bytes
+timer-item-payload-size-8 = ["_timer-item-payload"]
+
 #! ### Task Arena Size
 #! Sets the [task arena](#task-arena) size. Necessary if you’re not using `nightly`.
 #!
diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 5a476213b..bdd5ff5ae 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -96,29 +96,6 @@ impl TaskRef {
         &self.header().timer_queue_item
     }
 
-    /// Mark the task as timer-queued. Return whether it should be actually enqueued
-    /// using `_embassy_time_schedule_wake`.
-    ///
-    /// Entering this state prevents the task from being respawned while in a timer queue.
-    ///
-    /// Safety:
-    ///
-    /// This functions should only be called by the timer queue driver, before
-    /// enqueueing the timer item.
-    pub unsafe fn timer_enqueue(&self) -> timer_queue::TimerEnqueueOperation {
-        self.header().state.timer_enqueue()
-    }
-
-    /// Unmark the task as timer-queued.
-    ///
-    /// Safety:
-    ///
-    /// This functions should only be called by the timer queue implementation, after the task has
-    /// been removed from the timer queue.
-    pub unsafe fn timer_dequeue(&self) {
-        self.header().state.timer_dequeue()
-    }
-
     /// The returned pointer is valid for the entire TaskStorage.
     pub(crate) fn as_ptr(self) -> *const TaskHeader {
         self.ptr.as_ptr()
@@ -195,25 +172,7 @@ impl<F: Future + 'static> TaskStorage<F> {
         match future.poll(&mut cx) {
             Poll::Ready(_) => {
                 this.future.drop_in_place();
-
-                // Mark this task to be timer queued.
-                // We're splitting the enqueue in two parts, so that we can change task state
-                // to something that prevent re-queueing.
-                let op = this.raw.state.timer_enqueue();
-
-                // Now mark the task as not spawned, so that
-                // - it can be spawned again once it has been removed from the timer queue
-                // - it can not be timer-queued again
-                // We must do this before scheduling the wake, to prevent the task from being
-                // dequeued by the time driver while it's still SPAWNED.
                 this.raw.state.despawn();
-
-                // Now let's finish enqueueing. While we shouldn't get an `Ignore` here, it's
-                // better to be safe.
-                if op == timer_queue::TimerEnqueueOperation::Enqueue {
-                    // Schedule the task in the past, so it gets dequeued ASAP.
-                    unsafe { _embassy_time_schedule_wake(0, &waker) }
-                }
             }
             Poll::Pending => {}
         }
@@ -232,10 +191,6 @@ impl<F: Future + 'static> TaskStorage<F> {
     }
 }
 
-extern "Rust" {
-    fn _embassy_time_schedule_wake(at: u64, waker: &core::task::Waker);
-}
-
 /// An uninitialized [`TaskStorage`].
 pub struct AvailableTask<F: Future + 'static> {
     task: &'static TaskStorage<F>,
diff --git a/embassy-executor/src/raw/state_atomics.rs b/embassy-executor/src/raw/state_atomics.rs
index d7350464f..6f5266bda 100644
--- a/embassy-executor/src/raw/state_atomics.rs
+++ b/embassy-executor/src/raw/state_atomics.rs
@@ -1,7 +1,5 @@
 use core::sync::atomic::{AtomicU32, Ordering};
 
-use super::timer_queue::TimerEnqueueOperation;
-
 #[derive(Clone, Copy)]
 pub(crate) struct Token(());
 
@@ -16,8 +14,6 @@ pub(crate) fn locked<R>(f: impl FnOnce(Token) -> R) -> R {
 pub(crate) const STATE_SPAWNED: u32 = 1 << 0;
 /// Task is in the executor run queue
 pub(crate) const STATE_RUN_QUEUED: u32 = 1 << 1;
-/// Task is in the executor timer queue
-pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 2;
 
 pub(crate) struct State {
     state: AtomicU32,
@@ -71,32 +67,4 @@ impl State {
         let state = self.state.fetch_and(!STATE_RUN_QUEUED, Ordering::AcqRel);
         state & STATE_SPAWNED != 0
     }
-
-    /// Mark the task as timer-queued. Return whether it can be enqueued.
-    #[inline(always)]
-    pub fn timer_enqueue(&self) -> TimerEnqueueOperation {
-        if self
-            .state
-            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |state| {
-                // If not started, ignore it
-                if state & STATE_SPAWNED == 0 {
-                    None
-                } else {
-                    // Mark it as enqueued
-                    Some(state | STATE_TIMER_QUEUED)
-                }
-            })
-            .is_ok()
-        {
-            TimerEnqueueOperation::Enqueue
-        } else {
-            TimerEnqueueOperation::Ignore
-        }
-    }
-
-    /// Unmark the task as timer-queued.
-    #[inline(always)]
-    pub fn timer_dequeue(&self) {
-        self.state.fetch_and(!STATE_TIMER_QUEUED, Ordering::Relaxed);
-    }
 }
diff --git a/embassy-executor/src/raw/state_atomics_arm.rs b/embassy-executor/src/raw/state_atomics_arm.rs
index c1e8f69ab..4896b33c5 100644
--- a/embassy-executor/src/raw/state_atomics_arm.rs
+++ b/embassy-executor/src/raw/state_atomics_arm.rs
@@ -1,8 +1,6 @@
 use core::arch::asm;
 use core::sync::atomic::{compiler_fence, AtomicBool, AtomicU32, Ordering};
 
-use super::timer_queue::TimerEnqueueOperation;
-
 #[derive(Clone, Copy)]
 pub(crate) struct Token(());
 
@@ -16,7 +14,6 @@ pub(crate) fn locked<R>(f: impl FnOnce(Token) -> R) -> R {
 // Must be kept in sync with the layout of `State`!
 pub(crate) const STATE_SPAWNED: u32 = 1 << 0;
 pub(crate) const STATE_RUN_QUEUED: u32 = 1 << 8;
-pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 16;
 
 #[repr(C, align(4))]
 pub(crate) struct State {
@@ -24,9 +21,8 @@ pub(crate) struct State {
     spawned: AtomicBool,
     /// Task is in the executor run queue
     run_queued: AtomicBool,
-    /// Task is in the executor timer queue
-    timer_queued: AtomicBool,
     pad: AtomicBool,
+    pad2: AtomicBool,
 }
 
 impl State {
@@ -34,8 +30,8 @@ impl State {
         Self {
             spawned: AtomicBool::new(false),
             run_queued: AtomicBool::new(false),
-            timer_queued: AtomicBool::new(false),
             pad: AtomicBool::new(false),
+            pad2: AtomicBool::new(false),
         }
     }
 
@@ -101,32 +97,4 @@ impl State {
         self.run_queued.store(false, Ordering::Relaxed);
         r
     }
-
-    /// Mark the task as timer-queued. Return whether it can be enqueued.
-    #[inline(always)]
-    pub fn timer_enqueue(&self) -> TimerEnqueueOperation {
-        if self
-            .as_u32()
-            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |state| {
-                // If not started, ignore it
-                if state & STATE_SPAWNED == 0 {
-                    None
-                } else {
-                    // Mark it as enqueued
-                    Some(state | STATE_TIMER_QUEUED)
-                }
-            })
-            .is_ok()
-        {
-            TimerEnqueueOperation::Enqueue
-        } else {
-            TimerEnqueueOperation::Ignore
-        }
-    }
-
-    /// Unmark the task as timer-queued.
-    #[inline(always)]
-    pub fn timer_dequeue(&self) {
-        self.timer_queued.store(false, Ordering::Relaxed);
-    }
 }
diff --git a/embassy-executor/src/raw/state_critical_section.rs b/embassy-executor/src/raw/state_critical_section.rs
index 8e570b33c..29b10f6e3 100644
--- a/embassy-executor/src/raw/state_critical_section.rs
+++ b/embassy-executor/src/raw/state_critical_section.rs
@@ -3,14 +3,10 @@ use core::cell::Cell;
 pub(crate) use critical_section::{with as locked, CriticalSection as Token};
 use critical_section::{CriticalSection, Mutex};
 
-use super::timer_queue::TimerEnqueueOperation;
-
 /// Task is spawned (has a future)
 pub(crate) const STATE_SPAWNED: u32 = 1 << 0;
 /// Task is in the executor run queue
 pub(crate) const STATE_RUN_QUEUED: u32 = 1 << 1;
-/// Task is in the executor timer queue
-pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 2;
 
 pub(crate) struct State {
     state: Mutex<Cell<u32>>,
@@ -81,25 +77,4 @@ impl State {
             ok
         })
     }
-
-    /// Mark the task as timer-queued. Return whether it can be enqueued.
-    #[inline(always)]
-    pub fn timer_enqueue(&self) -> TimerEnqueueOperation {
-        self.update(|s| {
-            // FIXME: we need to split SPAWNED into two phases, to prevent enqueueing a task that is
-            // just being spawned, because its executor pointer may still be changing.
-            if *s & STATE_SPAWNED == STATE_SPAWNED {
-                *s |= STATE_TIMER_QUEUED;
-                TimerEnqueueOperation::Enqueue
-            } else {
-                TimerEnqueueOperation::Ignore
-            }
-        })
-    }
-
-    /// Unmark the task as timer-queued.
-    #[inline(always)]
-    pub fn timer_dequeue(&self) {
-        self.update(|s| *s &= !STATE_TIMER_QUEUED);
-    }
 }
diff --git a/embassy-executor/src/raw/timer_queue.rs b/embassy-executor/src/raw/timer_queue.rs
index 2ba0e00a9..e52453be4 100644
--- a/embassy-executor/src/raw/timer_queue.rs
+++ b/embassy-executor/src/raw/timer_queue.rs
@@ -4,6 +4,45 @@ use core::cell::Cell;
 
 use super::TaskRef;
 
+#[cfg(feature = "_timer-item-payload")]
+macro_rules! define_opaque {
+    ($size:tt) => {
+        /// An opaque data type.
+        #[repr(align($size))]
+        pub struct OpaqueData {
+            data: [u8; $size],
+        }
+
+        impl OpaqueData {
+            const fn new() -> Self {
+                Self { data: [0; $size] }
+            }
+
+            /// Access the data as a reference to a type `T`.
+            ///
+            /// Safety:
+            ///
+            /// The caller must ensure that the size of the type `T` is less than, or equal to
+            /// the size of the payload, and must ensure that the alignment of the type `T` is
+            /// less than, or equal to the alignment of the payload.
+            ///
+            /// The type must be valid when zero-initialized.
+            pub unsafe fn as_ref<T>(&self) -> &T {
+                &*(self.data.as_ptr() as *const T)
+            }
+        }
+    };
+}
+
+#[cfg(feature = "timer-item-payload-size-1")]
+define_opaque!(1);
+#[cfg(feature = "timer-item-payload-size-2")]
+define_opaque!(2);
+#[cfg(feature = "timer-item-payload-size-4")]
+define_opaque!(4);
+#[cfg(feature = "timer-item-payload-size-8")]
+define_opaque!(8);
+
 /// An item in the timer queue.
 pub struct TimerQueueItem {
     /// The next item in the queue.
@@ -14,6 +53,10 @@ pub struct TimerQueueItem {
 
     /// The time at which this item expires.
     pub expires_at: Cell<u64>,
+
+    /// Some implementation-defined, zero-initialized piece of data.
+    #[cfg(feature = "_timer-item-payload")]
+    pub payload: OpaqueData,
 }
 
 unsafe impl Sync for TimerQueueItem {}
@@ -23,17 +66,8 @@ impl TimerQueueItem {
         Self {
             next: Cell::new(None),
             expires_at: Cell::new(0),
+            #[cfg(feature = "_timer-item-payload")]
+            payload: OpaqueData::new(),
         }
     }
 }
-
-/// The operation to perform after `timer_enqueue` is called.
-#[derive(Debug, Copy, Clone, PartialEq)]
-#[cfg_attr(feature = "defmt", derive(defmt::Format))]
-#[must_use]
-pub enum TimerEnqueueOperation {
-    /// Enqueue the task (or update its expiration time).
-    Enqueue,
-    /// The task must not be enqueued in the timer queue.
-    Ignore,
-}
diff --git a/embassy-executor/tests/test.rs b/embassy-executor/tests/test.rs
index 992ab3da9..0ce1f1891 100644
--- a/embassy-executor/tests/test.rs
+++ b/embassy-executor/tests/test.rs
@@ -150,7 +150,3 @@ fn executor_task_cfg_args() {
         let (_, _, _) = (a, b, c);
     }
 }
-
-// We need this for the test to compile, even though we don't want to use timers at the moment.
-#[no_mangle]
-fn _embassy_time_schedule_wake(_at: u64, _waker: &core::task::Waker) {}
diff --git a/embassy-time-driver/src/lib.rs b/embassy-time-driver/src/lib.rs
index 57a9f7587..9f2795a01 100644
--- a/embassy-time-driver/src/lib.rs
+++ b/embassy-time-driver/src/lib.rs
@@ -46,6 +46,9 @@
 //!
 //! Then, you'll need to adapt the `schedule_wake` method to use this queue.
 //!
+//! Note that if you are using multiple queues, you will need to ensure that a single timer
+//! queue item is only ever enqueued into a single queue at a time.
+//!
 //! ```ignore
 //! use core::cell::RefCell;
 //! use core::task::Waker;
@@ -131,6 +134,7 @@ pub trait Driver: Send + Sync + 'static {
 
 extern "Rust" {
     fn _embassy_time_now() -> u64;
+    fn _embassy_time_schedule_wake(at: u64, waker: &Waker);
 }
 
 /// See [`Driver::now`]
@@ -138,6 +142,11 @@ pub fn now() -> u64 {
     unsafe { _embassy_time_now() }
 }
 
+/// Schedule the given waker to be woken at `at`.
+pub fn schedule_wake(at: u64, waker: &Waker) {
+    unsafe { _embassy_time_schedule_wake(at, waker) }
+}
+
 /// Set the time Driver implementation.
 ///
 /// See the module documentation for an example.
diff --git a/embassy-time-queue-driver/src/lib.rs b/embassy-time-queue-driver/src/lib.rs
index 97c81a124..333b6124d 100644
--- a/embassy-time-queue-driver/src/lib.rs
+++ b/embassy-time-queue-driver/src/lib.rs
@@ -10,8 +10,6 @@
 //! As a HAL implementer, you need to depend on this crate if you want to implement a time driver,
 //! but how you should do so is documented in `embassy-time-driver`.
 
-use core::task::Waker;
-
 #[cfg(feature = "_generic-queue")]
 pub mod queue_generic;
 #[cfg(not(feature = "_generic-queue"))]
@@ -21,29 +19,3 @@ pub mod queue_integrated;
 pub use queue_generic::Queue;
 #[cfg(not(feature = "_generic-queue"))]
 pub use queue_integrated::Queue;
-
-extern "Rust" {
-    fn _embassy_time_schedule_wake(at: u64, waker: &Waker);
-}
-
-/// Schedule the given waker to be woken at `at`.
-pub fn schedule_wake(at: u64, waker: &Waker) {
-    // This function is not implemented in embassy-time-driver because it needs access to executor
-    // internals. The function updates task state, then delegates to the implementation provided
-    // by the time driver.
-    #[cfg(not(feature = "_generic-queue"))]
-    {
-        use embassy_executor::raw::task_from_waker;
-        use embassy_executor::raw::timer_queue::TimerEnqueueOperation;
-        // The very first thing we must do, before we even access the timer queue, is to
-        // mark the task a TIMER_QUEUED. This ensures that the task that is being scheduled
-        // can not be respawn while we are accessing the timer queue.
-        let task = task_from_waker(waker);
-        if unsafe { task.timer_enqueue() } == TimerEnqueueOperation::Ignore {
-            // We are not allowed to enqueue the task in the timer queue. This is because the
-            // task is not spawned, and so it makes no sense to schedule it.
-            return;
-        }
-    }
-    unsafe { _embassy_time_schedule_wake(at, waker) }
-}
diff --git a/embassy-time-queue-driver/src/queue_integrated.rs b/embassy-time-queue-driver/src/queue_integrated.rs
index 6bb4c0c1a..246cf1d63 100644
--- a/embassy-time-queue-driver/src/queue_integrated.rs
+++ b/embassy-time-queue-driver/src/queue_integrated.rs
@@ -83,7 +83,6 @@ impl Queue {
                 // Remove it
                 prev.set(item.next.get());
                 item.next.set(None);
-                unsafe { p.timer_dequeue() };
             }
         }
     }
diff --git a/embassy-time/Cargo.toml b/embassy-time/Cargo.toml
index e3074119f..4f4ea0b14 100644
--- a/embassy-time/Cargo.toml
+++ b/embassy-time/Cargo.toml
@@ -24,8 +24,8 @@ target = "x86_64-unknown-linux-gnu"
 features = ["defmt", "std"]
 
 [features]
-std = ["tick-hz-1_000_000", "critical-section/std"]
-wasm = ["dep:wasm-bindgen", "dep:js-sys", "dep:wasm-timer", "tick-hz-1_000_000"]
+std = ["tick-hz-1_000_000", "critical-section/std", "dep:embassy-time-queue-driver"]
+wasm = ["dep:wasm-bindgen", "dep:js-sys", "dep:wasm-timer", "tick-hz-1_000_000", "dep:embassy-time-queue-driver"]
 
 ## Display the time since startup next to defmt log messages.
 ## At most 1 `defmt-timestamp-uptime-*` feature can be used.
@@ -40,7 +40,7 @@ defmt-timestamp-uptime-tms = ["defmt"]
 defmt-timestamp-uptime-tus = ["defmt"]
 
 ## Create a `MockDriver` that can be manually advanced for testing purposes.
-mock-driver = ["tick-hz-1_000_000"]
+mock-driver = ["tick-hz-1_000_000", "dep:embassy-time-queue-driver"]
 
 #! ### Tick Rate
 #!
@@ -384,7 +384,7 @@ tick-hz-5_242_880_000 = ["embassy-time-driver/tick-hz-5_242_880_000"]
 
 [dependencies]
 embassy-time-driver = { version = "0.1.0", path = "../embassy-time-driver" }
-embassy-time-queue-driver = { version = "0.1.0", path = "../embassy-time-queue-driver" }
+embassy-time-queue-driver = { version = "0.1.0", path = "../embassy-time-queue-driver", optional = true}
 
 defmt = { version = "0.3", optional = true }
 log = { version = "0.4.14", optional = true }
diff --git a/embassy-time/src/timer.rs b/embassy-time/src/timer.rs
index 4d7194b20..295ddbd9b 100644
--- a/embassy-time/src/timer.rs
+++ b/embassy-time/src/timer.rs
@@ -157,7 +157,7 @@ impl Future for Timer {
         if self.yielded_once && self.expires_at <= Instant::now() {
             Poll::Ready(())
         } else {
-            embassy_time_queue_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
+            embassy_time_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
             self.yielded_once = true;
             Poll::Pending
         }
@@ -238,7 +238,7 @@ impl Ticker {
                 self.expires_at += dur;
                 Poll::Ready(())
             } else {
-                embassy_time_queue_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
+                embassy_time_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
                 Poll::Pending
             }
         })
@@ -255,7 +255,7 @@ impl Stream for Ticker {
             self.expires_at += dur;
             Poll::Ready(Some(()))
         } else {
-            embassy_time_queue_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
+            embassy_time_driver::schedule_wake(self.expires_at.as_ticks(), cx.waker());
             Poll::Pending
         }
     }