Separate claimed and spawned states

embassy-executor/src/raw/mod.rs

@@ -189,9 +189,9 @@ pub struct AvailableTask<F: Future + 'static> {
 impl<F: Future + 'static> AvailableTask<F> {
     /// Try to claim a [`TaskStorage`].
     ///
-    /// This function returns `None` if a task has already been spawned and has not finished running.
+    /// This function returns `None` if a task has already been claimed and has not finished running.
     pub fn claim(task: &'static TaskStorage<F>) -> Option<Self> {
-        task.raw.state.spawn().then(|| Self { task })
+        task.raw.state.claim().then(|| Self { task })
     }
 
     fn initialize_impl<S>(self, future: impl FnOnce() -> F) -> SpawnToken<S> {
@@ -368,6 +368,7 @@ impl SyncExecutor {
 
     pub(super) unsafe fn spawn(&'static self, task: TaskRef) {
         task.header().executor.set(Some(self));
+        task.header().state.mark_spawned();
 
         #[cfg(feature = "rtos-trace")]
         trace::task_new(task.as_ptr() as u32);

embassy-executor/src/raw/state_atomics.rs

@@ -1,12 +1,14 @@
 use core::sync::atomic::{AtomicU32, Ordering};
 
+/// Task is claimed (its storage is in use)
+pub(crate) const STATE_CLAIMED: u32 = 1 << 0;
 /// Task is spawned (has a future)
-pub(crate) const STATE_SPAWNED: u32 = 1 << 0;
+pub(crate) const STATE_SPAWNED: u32 = 1 << 1;
 /// Task is in the executor run queue
-pub(crate) const STATE_RUN_QUEUED: u32 = 1 << 1;
+pub(crate) const STATE_RUN_QUEUED: u32 = 1 << 2;
 /// Task is in the executor timer queue
 #[cfg(feature = "integrated-timers")]
-pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 2;
+pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 3;
 
 pub(crate) struct State {
     state: AtomicU32,
@@ -19,18 +21,29 @@ impl State {
         }
     }
 
-    /// If task is idle, mark it as spawned + run_queued and return true.
+    /// If task is idle, mark it as claimed and return true.
     #[inline(always)]
-    pub fn spawn(&self) -> bool {
+    pub fn claim(&self) -> bool {
         self.state
-            .compare_exchange(0, STATE_SPAWNED | STATE_RUN_QUEUED, Ordering::AcqRel, Ordering::Acquire)
+            .compare_exchange(0, STATE_CLAIMED, Ordering::AcqRel, Ordering::Acquire)
             .is_ok()
     }
 
+    /// Mark a claimed task ready to run.
+    ///
+    /// # Safety
+    ///
+    /// The task must be claimed, its executor must be configured. This function must
+    /// not be called when the task is already spawned.
+    #[inline(always)]
+    pub unsafe fn mark_spawned(&self) {
+        self.state.store(STATE_SPAWNED | STATE_RUN_QUEUED, Ordering::Release);
+    }
+
     /// Unmark the task as spawned.
     #[inline(always)]
     pub fn despawn(&self) {
-        self.state.fetch_and(!STATE_SPAWNED, Ordering::AcqRel);
+        self.state.fetch_and(!(STATE_SPAWNED | STATE_CLAIMED), Ordering::AcqRel);
     }
 
     /// Mark the task as run-queued if it's spawned and isn't already run-queued. Return true on success.

embassy-executor/src/raw/state_atomics_arm.rs

@@ -2,56 +2,66 @@ use core::arch::asm;
 use core::sync::atomic::{compiler_fence, AtomicBool, AtomicU32, Ordering};
 
 // 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;
+pub(crate) const STATE_CLAIMED: u32 = 1 << 0;
+pub(crate) const STATE_SPAWNED: u32 = 1 << 8;
+pub(crate) const STATE_RUN_QUEUED: u32 = 1 << 16;
 
 #[repr(C, align(4))]
 pub(crate) struct State {
+    /// Task is claimed (its storage is in use)
+    claimed: AtomicBool,
     /// Task is spawned (has a future)
     spawned: AtomicBool,
     /// Task is in the executor run queue
     run_queued: AtomicBool,
     /// Task is in the executor timer queue
     timer_queued: AtomicBool,
-    pad: AtomicBool,
 }
 
 impl State {
     pub const fn new() -> State {
         Self {
+            claimed: AtomicBool::new(false),
             spawned: AtomicBool::new(false),
             run_queued: AtomicBool::new(false),
             timer_queued: AtomicBool::new(false),
-            pad: AtomicBool::new(false),
         }
     }
 
     fn as_u32(&self) -> &AtomicU32 {
         unsafe { &*(self as *const _ as *const AtomicU32) }
     }
 
-    /// If task is idle, mark it as spawned + run_queued and return true.
+    /// If task is idle, mark it as claimed and return true.
     #[inline(always)]
-    pub fn spawn(&self) -> bool {
+    pub fn claim(&self) -> bool {
         compiler_fence(Ordering::Release);
+
         let r = self
-            .as_u32()
-            .compare_exchange(
-                0,
-                STATE_SPAWNED | STATE_RUN_QUEUED,
-                Ordering::Relaxed,
-                Ordering::Relaxed,
-            )
+            .claimed
+            .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
             .is_ok();
+
         compiler_fence(Ordering::Acquire);
         r
     }
 
+    /// Mark a claimed task ready to run.
+    ///
+    /// # Safety
+    ///
+    /// The task must be claimed, its executor must be configured. This function must
+    /// not be called when the task is already spawned.
+    #[inline(always)]
+    pub unsafe fn mark_spawned(&self) {
+        self.as_u32().store(STATE_SPAWNED | STATE_RUN_QUEUED, Ordering::Release);
+    }
+
     /// Unmark the task as spawned.
     #[inline(always)]
     pub fn despawn(&self) {
-        compiler_fence(Ordering::Release);
-        self.spawned.store(false, Ordering::Relaxed);
+        self.as_u32()
+            .fetch_and(!(STATE_SPAWNED | STATE_CLAIMED), Ordering::AcqRel);
     }
 
     /// Mark the task as run-queued if it's spawned and isn't already run-queued. Return true on success.

embassy-executor/src/raw/state_critical_section.rs

@@ -2,13 +2,15 @@ use core::cell::Cell;
 
 use critical_section::Mutex;
 
+/// Task is claimed (its storage is in use)
+pub(crate) const STATE_CLAIMED: u32 = 1 << 0;
 /// Task is spawned (has a future)
-pub(crate) const STATE_SPAWNED: u32 = 1 << 0;
+pub(crate) const STATE_SPAWNED: u32 = 1 << 1;
 /// Task is in the executor run queue
-pub(crate) const STATE_RUN_QUEUED: u32 = 1 << 1;
+pub(crate) const STATE_RUN_QUEUED: u32 = 1 << 2;
 /// Task is in the executor timer queue
 #[cfg(feature = "integrated-timers")]
-pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 2;
+pub(crate) const STATE_TIMER_QUEUED: u32 = 1 << 3;
 
 pub(crate) struct State {
     state: Mutex<Cell<u32>>,
@@ -31,23 +33,34 @@ impl State {
         })
     }
 
-    /// If task is idle, mark it as spawned + run_queued and return true.
+    /// If task is idle, mark it as claimed and return true.
     #[inline(always)]
-    pub fn spawn(&self) -> bool {
+    pub fn claim(&self) -> bool {
         self.update(|s| {
             if *s == 0 {
-                *s = STATE_SPAWNED | STATE_RUN_QUEUED;
+                *s = STATE_CLAIMED;
                 true
             } else {
                 false
             }
         })
     }
 
+    /// Mark a claimed task ready to run.
+    ///
+    /// # Safety
+    ///
+    /// The task must be claimed, its executor must be configured. This function must
+    /// not be called when the task is already spawned.
+    #[inline(always)]
+    pub unsafe fn mark_spawned(&self) {
+        self.update(|s| *s = STATE_SPAWNED | STATE_RUN_QUEUED);
+    }
+
     /// Unmark the task as spawned.
     #[inline(always)]
     pub fn despawn(&self) {
-        self.update(|s| *s &= !STATE_SPAWNED);
+        self.update(|s| *s &= !(STATE_SPAWNED | STATE_CLAIMED));
     }
 
     /// Mark the task as run-queued if it's spawned and isn't already run-queued. Return true on success.