Take critical section instead of unsafe

embassy-executor/src/raw/run_queue_critical_section.rs

@@ -1,6 +1,6 @@
 use core::cell::Cell;
 
-use critical_section::{CriticalSection, Mutex};
+use critical_section::Mutex;
 
 use super::TaskRef;
 
@@ -65,10 +65,10 @@ 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();
-            task.header().state.run_dequeue();
+            critical_section::with(|cs| {
+                next = task.header().run_queue_item.next.borrow(cs).get();
+                task.header().state.run_dequeue(cs);
+            });
 
             on_task(task);
         }

embassy-executor/src/raw/state_critical_section.rs

@@ -1,6 +1,6 @@
 use core::cell::Cell;
 
-use critical_section::Mutex;
+use critical_section::{CriticalSection, Mutex};
 
 /// Task is claimed (it is being spawned)
 pub(crate) const STATE_CLAIMED: u32 = 1 << 0;
@@ -24,13 +24,15 @@ impl State {
     }
 
     fn update<R>(&self, f: impl FnOnce(&mut u32) -> R) -> R {
-        critical_section::with(|cs| {
-            let s = self.state.borrow(cs);
-            let mut val = s.get();
-            let r = f(&mut val);
-            s.set(val);
-            r
-        })
+        critical_section::with(|cs| self.update_with_cs(cs, f))
+    }
+
+    fn update_with_cs<R>(&self, cs: CriticalSection<'_>, f: impl FnOnce(&mut u32) -> R) -> R {
+        let s = self.state.borrow(cs);
+        let mut val = s.get();
+        let r = f(&mut val);
+        s.set(val);
+        r
     }
 
     /// If task is idle, mark it as claimed and return true.
@@ -83,8 +85,8 @@ impl State {
 
     /// Unmark the task as run-queued.
     #[inline(always)]
-    pub fn run_dequeue(&self) {
-        self.update(|s| *s &= !STATE_RUN_QUEUED);
+    pub fn run_dequeue(&self, cs: CriticalSection<'_>) {
+        self.update_with_cs(cs, |s| *s &= !STATE_RUN_QUEUED);
     }
 
     /// Mark the task as timer-queued. Return whether it was newly queued (i.e. not queued before)