Merge pull request #3559 from bugadani/gaw

Generalize AtomicWaker

embassy-sync/CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `clear`, `len`, `is_empty` and `is_full` functions to `zerocopy_channel`.
 - Add `capacity`, `free_capacity`, `clear`, `len`, `is_empty` and `is_full` functions to `channel::{Sender, Receiver}`.
 - Add `capacity`, `free_capacity`, `clear`, `len`, `is_empty` and `is_full` functions to `priority_channel::{Sender, Receiver}`.
+- Add `GenericAtomicWaker` utility.
 
 ## 0.6.0 - 2024-05-29
 

embassy-sync/src/waitqueue/atomic_waker.rs

@@ -1,26 +1,25 @@
 use core::cell::Cell;
 use core::task::Waker;
 
-use crate::blocking_mutex::raw::CriticalSectionRawMutex;
+use crate::blocking_mutex::raw::{CriticalSectionRawMutex, RawMutex};
 use crate::blocking_mutex::Mutex;
 
 /// Utility struct to register and wake a waker.
-pub struct AtomicWaker {
+pub struct GenericAtomicWaker<M: RawMutex> {
-    waker: Mutex<CriticalSectionRawMutex, Cell<Option<Waker>>>,
+    waker: Mutex<M, Cell<Option<Waker>>>,
 }
 
-impl AtomicWaker {
+impl<M: RawMutex> GenericAtomicWaker<M> {
     /// Create a new `AtomicWaker`.
-    pub const fn new() -> Self {
+    pub const fn new(mutex: M) -> Self {
         Self {
-            waker: Mutex::const_new(CriticalSectionRawMutex::new(), Cell::new(None)),
+            waker: Mutex::const_new(mutex, Cell::new(None)),
         }
     }
 
     /// Register a waker. Overwrites the previous waker, if any.
     pub fn register(&self, w: &Waker) {
-        critical_section::with(|cs| {
+        self.waker.lock(|cell| {
-            let cell = self.waker.borrow(cs);
             cell.set(match cell.replace(None) {
                 Some(w2) if (w2.will_wake(w)) => Some(w2),
                 _ => Some(w.clone()),
@@ -30,8 +29,7 @@ impl AtomicWaker {
 
     /// Wake the registered waker, if any.
     pub fn wake(&self) {
-        critical_section::with(|cs| {
+        self.waker.lock(|cell| {
-            let cell = self.waker.borrow(cs);
             if let Some(w) = cell.replace(None) {
                 w.wake_by_ref();
                 cell.set(Some(w));
@@ -39,3 +37,27 @@ impl AtomicWaker {
         })
     }
 }
+
+/// Utility struct to register and wake a waker.
+pub struct AtomicWaker {
+    waker: GenericAtomicWaker<CriticalSectionRawMutex>,
+}
+
+impl AtomicWaker {
+    /// Create a new `AtomicWaker`.
+    pub const fn new() -> Self {
+        Self {
+            waker: GenericAtomicWaker::new(CriticalSectionRawMutex::new()),
+        }
+    }
+
+    /// Register a waker. Overwrites the previous waker, if any.
+    pub fn register(&self, w: &Waker) {
+        self.waker.register(w);
+    }
+
+    /// Wake the registered waker, if any.
+    pub fn wake(&self) {
+        self.waker.wake();
+    }
+}