From ec96395d084d5edc8be25ddaea8547e2ebd447a6 Mon Sep 17 00:00:00 2001
From: Dániel Buga <bugadani@gmail.com>
Date: Mon, 9 Dec 2024 08:43:57 +0100
Subject: Prevent task from respawning while in the timer queue

---
 embassy-executor/src/raw/mod.rs | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

(limited to 'embassy-executor/src/raw/mod.rs')

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index f9c6509f1..14d689900 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -50,7 +50,7 @@ pub(crate) struct TaskHeader {
 }
 
 /// This is essentially a `&'static TaskStorage<F>` where the type of the future has been erased.
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, PartialEq)]
 pub struct TaskRef {
     ptr: NonNull<TaskHeader>,
 }
@@ -72,6 +72,16 @@ impl TaskRef {
         }
     }
 
+    /// # Safety
+    ///
+    /// The result of this function must only be compared
+    /// for equality, or stored, but not used.
+    pub const unsafe fn dangling() -> Self {
+        Self {
+            ptr: NonNull::dangling(),
+        }
+    }
+
     pub(crate) fn header(self) -> &'static TaskHeader {
         unsafe { self.ptr.as_ref() }
     }
@@ -88,6 +98,30 @@ impl TaskRef {
         &self.header().timer_queue_item
     }
 
+    /// Mark the task as timer-queued. Return whether it was newly queued (i.e. not queued before)
+    ///
+    /// 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 implementation, before
+    /// enqueueing the timer item.
+    #[cfg(feature = "integrated-timers")]
+    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.
+    #[cfg(feature = "integrated-timers")]
+    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()
-- 
cgit