Take critical section instead of unsafe

3 files changed, 22 insertions, 14 deletions

diff --git a/embassy-executor/src/raw/mod.rs b/embassy-executor/src/raw/mod.rs
index 6503b556f..c79fdae60 100644
--- a/embassy-executor/src/raw/mod.rs
+++ b/embassy-executor/src/raw/mod.rs
@@ -411,15 +411,6 @@ impl SyncExecutor {
         self.run_queue.dequeue_all(|p| {
             let task = p.header();
 
-            if !task.state.run_dequeue() {
-                // If task is not running, ignore it. This can happen in the following scenario:
-                //   - Task gets dequeued, poll starts
-                //   - While task is being polled, it gets woken. It gets placed in the queue.
-                //   - Task poll finishes, returning done=true
-                //   - RUNNING bit is cleared, but the task is already in the queue.
-                return;
-            }
-
             #[cfg(feature = "trace")]
             trace::task_exec_begin(self, &p);
 
diff --git a/embassy-executor/src/raw/run_queue_atomics.rs b/embassy-executor/src/raw/run_queue_atomics.rs
index efdafdff0..aad90d767 100644
--- a/embassy-executor/src/raw/run_queue_atomics.rs
+++ b/embassy-executor/src/raw/run_queue_atomics.rs
@@ -81,7 +81,16 @@ impl RunQueue {
             // safety: there are no concurrent accesses to `next`
             next = unsafe { task.header().run_queue_item.next.get() };
 
-            on_task(task);
+            let run_task = task.header().state.run_dequeue();
+
+            if run_task {
+                // If task is not running, ignore it. This can happen in the following scenario:
+                //   - Task gets dequeued, poll starts
+                //   - While task is being polled, it gets woken. It gets placed in the queue.
+                //   - Task poll finishes, returning done=true
+                //   - RUNNING bit is cleared, but the task is already in the queue.
+                on_task(task);
+            }
         }
     }
 }
diff --git a/embassy-executor/src/raw/run_queue_critical_section.rs b/embassy-executor/src/raw/run_queue_critical_section.rs
index 90f09e8c8..4f1b2855a 100644
--- a/embassy-executor/src/raw/run_queue_critical_section.rs
+++ b/embassy-executor/src/raw/run_queue_critical_section.rs
@@ -63,11 +63,19 @@ impl RunQueue {
             // If the task re-enqueues itself, the `next` pointer will get overwritten.
             // Therefore, first read the next pointer, and only then process the task.
 
-            // safety: we know if the task is enqueued, no one else will touch the `next` pointer.
-            let cs = unsafe { CriticalSection::new() };
-            next = task.header().run_queue_item.next.borrow(cs).get();
+            let run_task = critical_section::with(|cs| {
+                next = task.header().run_queue_item.next.borrow(cs).get();
+                task.header().state.run_dequeue(cs)
+            });
 
-            on_task(task);
+            if run_task {
+                // If task is not running, ignore it. This can happen in the following scenario:
+                //   - Task gets dequeued, poll starts
+                //   - While task is being polled, it gets woken. It gets placed in the queue.
+                //   - Task poll finishes, returning done=true
+                //   - RUNNING bit is cleared, but the task is already in the queue.
+                on_task(task);
+            }
         }
     }
 }