1 files changed, 36 insertions, 0 deletions

diff --git a/embassy-executor/src/raw/state_atomics.rs b/embassy-executor/src/raw/state_atomics.rs
index e4127897e..d03c61ade 100644
--- a/embassy-executor/src/raw/state_atomics.rs
+++ b/embassy-executor/src/raw/state_atomics.rs
@@ -1,9 +1,15 @@
 use core::sync::atomic::{AtomicU32, Ordering};
 
+#[cfg(feature = "integrated-timers")]
+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
+#[cfg(feature = "integrated-timers")]
+pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 2;
 
 pub(crate) struct State {
     state: AtomicU32,
@@ -52,4 +58,34 @@ 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.
+    #[cfg(feature = "integrated-timers")]
+    #[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.
+    #[cfg(feature = "integrated-timers")]
+    #[inline(always)]
+    pub fn timer_dequeue(&self) {
+        self.state.fetch_and(!STATE_TIMER_QUEUED, Ordering::Relaxed);
+    }
 }