Prevent task from respawning while in the timer queue

1 files changed, 14 insertions, 0 deletions

diff --git a/embassy-time-queue-driver/src/lib.rs b/embassy-time-queue-driver/src/lib.rs
index 0c78921ed..2d5fd449a 100644
--- a/embassy-time-queue-driver/src/lib.rs
+++ b/embassy-time-queue-driver/src/lib.rs
@@ -73,6 +73,20 @@ extern "Rust" {
 
 /// Schedule the given waker to be woken at `at`.
 pub fn schedule_wake(at: u64, waker: &Waker) {
+    #[cfg(feature = "integrated-timers")]
+    {
+        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) }
 }