From f389ba37219d842d7db0ab94cd421c69645a5757 Mon Sep 17 00:00:00 2001
From: Dániel Buga <bugadani@gmail.com>
Date: Fri, 13 Dec 2024 20:35:40 +0100
Subject: Only lock once to wake a task

---
 embassy-executor/src/raw/state_critical_section.rs | 42 +++++++++++++---------
 1 file changed, 25 insertions(+), 17 deletions(-)

(limited to 'embassy-executor/src/raw/state_critical_section.rs')

diff --git a/embassy-executor/src/raw/state_critical_section.rs b/embassy-executor/src/raw/state_critical_section.rs
index 367162ba2..8e570b33c 100644
--- a/embassy-executor/src/raw/state_critical_section.rs
+++ b/embassy-executor/src/raw/state_critical_section.rs
@@ -1,6 +1,7 @@
 use core::cell::Cell;
 
-use critical_section::Mutex;
+pub(crate) use critical_section::{with as locked, CriticalSection as Token};
+use critical_section::{CriticalSection, Mutex};
 
 use super::timer_queue::TimerEnqueueOperation;
 
@@ -23,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 spawned + run_queued and return true.
@@ -51,17 +54,22 @@ impl State {
         self.update(|s| *s &= !STATE_SPAWNED);
     }
 
-    /// Mark the task as run-queued if it's spawned and isn't already run-queued. Return true on success.
+    /// Mark the task as run-queued if it's spawned and isn't already run-queued. Run the given
+    /// function if the task was successfully marked.
     #[inline(always)]
-    pub fn run_enqueue(&self) -> bool {
-        self.update(|s| {
-            if (*s & STATE_RUN_QUEUED != 0) || (*s & STATE_SPAWNED == 0) {
-                false
-            } else {
-                *s |= STATE_RUN_QUEUED;
-                true
+    pub fn run_enqueue(&self, f: impl FnOnce(Token)) {
+        critical_section::with(|cs| {
+            if self.update_with_cs(cs, |s| {
+                if (*s & STATE_RUN_QUEUED != 0) || (*s & STATE_SPAWNED == 0) {
+                    false
+                } else {
+                    *s |= STATE_RUN_QUEUED;
+                    true
+                }
+            }) {
+                f(cs);
             }
-        })
+        });
     }
 
     /// Unmark the task as run-queued. Return whether the task is spawned.
-- 
cgit