Prevent task from respawning while in the timer queue

2 files changed, 29 insertions, 5 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) }
 }
 
diff --git a/embassy-time-queue-driver/src/queue_integrated.rs b/embassy-time-queue-driver/src/queue_integrated.rs
index cb0f79356..b905c00c3 100644
--- a/embassy-time-queue-driver/src/queue_integrated.rs
+++ b/embassy-time-queue-driver/src/queue_integrated.rs
@@ -24,16 +24,21 @@ impl TimerQueue {
         if item.next.get().is_none() {
             // If not in the queue, add it and update.
             let prev = self.head.replace(Some(p));
-            item.next.set(prev);
+            item.next.set(if prev.is_none() {
+                Some(unsafe { TaskRef::dangling() })
+            } else {
+                prev
+            });
+            item.expires_at.set(at);
+            true
         } else if at <= item.expires_at.get() {
             // If expiration is sooner than previously set, update.
+            item.expires_at.set(at);
+            true
         } else {
             // Task does not need to be updated.
-            return false;
+            false
         }
-
-        item.expires_at.set(at);
-        true
     }
 
     /// Dequeues expired timers and returns the next alarm time.
@@ -64,6 +69,10 @@ impl TimerQueue {
     fn retain(&self, mut f: impl FnMut(TaskRef) -> bool) {
         let mut prev = &self.head;
         while let Some(p) = prev.get() {
+            if unsafe { p == TaskRef::dangling() } {
+                // prev was the last item, stop
+                break;
+            }
             let item = p.timer_queue_item();
             if f(p) {
                 // Skip to next
@@ -72,6 +81,7 @@ impl TimerQueue {
                 // Remove it
                 prev.set(item.next.get());
                 item.next.set(None);
+                unsafe { p.timer_dequeue() };
             }
         }
     }