Prevent task from respawning while in the timer queue

1 files changed, 38 insertions, 2 deletions

diff --git a/embassy-executor/src/raw/state_atomics_arm.rs b/embassy-executor/src/raw/state_atomics_arm.rs
index b673c7359..f6f2e8f08 100644
--- a/embassy-executor/src/raw/state_atomics_arm.rs
+++ b/embassy-executor/src/raw/state_atomics_arm.rs
@@ -1,9 +1,14 @@
 use core::arch::asm;
 use core::sync::atomic::{compiler_fence, AtomicBool, AtomicU32, Ordering};
 
+#[cfg(feature = "integrated-timers")]
+use super::timer_queue::TimerEnqueueOperation;
+
 // 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;
+#[cfg(feature = "integrated-timers")]
+pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 16;
 
 #[repr(C, align(4))]
 pub(crate) struct State {
@@ -11,8 +16,9 @@ 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 {
@@ -20,8 +26,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),
         }
     }
 
@@ -85,4 +91,34 @@ impl State {
         self.run_queued.store(false, Ordering::Relaxed);
         r
     }
+
+    /// Mark the task as timer-queued. Return whether it can be enqueued.
+    #[cfg(feature = "integrated-timers")]
+    #[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.
+    #[cfg(feature = "integrated-timers")]
+    #[inline(always)]
+    pub fn timer_dequeue(&self) {
+        self.timer_queued.store(false, Ordering::Relaxed);
+    }
 }