diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs index b47b90fb..d012c480 100644 --- a/embassy-executor/src/raw/mod.rs +++ b/embassy-executor/src/raw/mod.rs @@ -189,7 +189,7 @@ impl TaskStorage { /// An uninitialized [`TaskStorage`]. pub struct AvailableTask { - task: &'static TaskStorage, + task: Option<&'static TaskStorage>, } impl AvailableTask { @@ -201,21 +201,23 @@ impl AvailableTask { .state .compare_exchange(STATE_ELIGIBLE, STATE_CLAIMED, Ordering::AcqRel, Ordering::Acquire) .ok() - .map(|_| Self { task }) + .map(|_| Self { task: Some(task) }) } - fn initialize_impl(self, future: impl FnOnce() -> F) -> SpawnToken { + fn initialize_impl(mut self, future: impl FnOnce() -> F) -> SpawnToken { unsafe { - self.task.raw.poll_fn.set(Some(TaskStorage::::poll)); - self.task.future.write_in_place(future); + let task = self.task.take().unwrap(); - self.task - .raw + // initialize task body + task.raw.poll_fn.set(Some(TaskStorage::::poll)); + task.future.write_in_place(future); + + // set state now that task is valid + task.raw .state .store(STATE_SPAWNED | STATE_RUN_QUEUED, Ordering::Release); - let task = TaskRef::new(self.task); - + let task = TaskRef::new(task); SpawnToken::new(task) } } @@ -262,6 +264,18 @@ impl AvailableTask { } } +impl Drop for AvailableTask { + fn drop(&mut self) { + if let Some(task) = self.task { + // if AvailableTask is dropped before a task is spawned on it, + // set task back to STATE_ELIGIBLE so we can re-use the slot. + // This also ensures that we can re-use this slot even if the + // user callback in initialize_impl panics. + task.raw.state.store(STATE_ELIGIBLE, Ordering::Release); + } + } +} + /// Raw storage that can hold up to N tasks of the same type. /// /// This is essentially a `[TaskStorage; N]`.