Don't force a wake to despawn

2 files changed, 4 insertions, 5 deletions

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 39d2d73ab..4a4ecf603 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -215,8 +215,6 @@ impl<F: Future + 'static> TaskStorage<F> {
         let mut cx = Context::from_waker(&waker);
         match future.poll(&mut cx) {
             Poll::Ready(_) => {
-                waker.wake_by_ref();
-
                 // As the future has finished and this function will not be called
                 // again, we can safely drop the future here.
                 this.future.drop_in_place();
@@ -224,6 +222,10 @@ impl<F: Future + 'static> TaskStorage<F> {
                 // We replace the poll_fn with a despawn function, so that the task is cleaned up
                 // when the executor polls it next.
                 this.raw.poll_fn.set(Some(poll_to_despawn));
+
+                // Make sure we despawn last, so that other threads can only spawn the task
+                // after we're done with it.
+                this.raw.state.despawn();
             }
             Poll::Pending => {}
         }
diff --git a/embassy-executor/tests/test.rs b/embassy-executor/tests/test.rs
index d8c5a6ae3..78c49c071 100644
--- a/embassy-executor/tests/test.rs
+++ b/embassy-executor/tests/test.rs
@@ -69,7 +69,6 @@ fn executor_task() {
         &[
             "pend",       // spawning a task pends the executor
             "poll task1", // poll only once.
-            "pend",       // task is done, wakes itself to exit
         ]
     )
 }
@@ -180,7 +179,6 @@ fn waking_after_completion_does_not_poll() {
             "pend",       // manual wake, single pend for two wakes
             "pend",       // respawning a task pends the executor
             "poll task1", //
-            "pend",       // task is done, wakes itself to exit
         ]
     )
 }
@@ -268,7 +266,6 @@ fn waking_with_old_waker_after_respawn() {
             "yield_now",  //
             "pend",       // manual wake, gets cleared by poll
             "poll task1", //
-            "pend",       // task is done, wakes itself to exit
         ]
     );
 }