1 files changed, 32 insertions, 10 deletions

diff --git a/embassy-sync/src/waitqueue/atomic_waker.rs b/embassy-sync/src/waitqueue/atomic_waker.rs
index 63fe04a6e..231902c5a 100644
--- a/embassy-sync/src/waitqueue/atomic_waker.rs
+++ b/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 {
-    waker: Mutex<CriticalSectionRawMutex, Cell<Option<Waker>>>,
+pub struct GenericAtomicWaker<M: RawMutex> {
+    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| {
-            let cell = self.waker.borrow(cs);
+        self.waker.lock(|cell| {
             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| {
-            let cell = self.waker.borrow(cs);
+        self.waker.lock(|cell| {
             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();
+    }
+}