4 files changed, 49 insertions, 11 deletions

diff --git a/embassy-sync/src/waitqueue/atomic_waker.rs b/embassy-sync/src/waitqueue/atomic_waker.rs
index 63fe04a6e..d2bf890e5 100644
--- a/embassy-sync/src/waitqueue/atomic_waker.rs
+++ b/embassy-sync/src/waitqueue/atomic_waker.rs
@@ -1,26 +1,28 @@
 use core::cell::Cell;
 use core::task::Waker;
 
-use crate::blocking_mutex::raw::CriticalSectionRawMutex;
 use crate::blocking_mutex::Mutex;
+use crate::blocking_mutex::raw::{CriticalSectionRawMutex, RawMutex};
 
 /// Utility struct to register and wake a waker.
-pub struct AtomicWaker {
-    waker: Mutex<CriticalSectionRawMutex, Cell<Option<Waker>>>,
+/// If a waker is registered, registering another waker will replace the previous one without waking it.
+/// Intended to wake a task from an interrupt. Therefore, it is generally not expected,
+/// that multiple tasks register try to register a waker simultaneously.
+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 +32,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 +40,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();
+    }
+}
diff --git a/embassy-sync/src/waitqueue/atomic_waker_turbo.rs b/embassy-sync/src/waitqueue/atomic_waker_turbo.rs
index 5c6a96ec8..a45adeab8 100644
--- a/embassy-sync/src/waitqueue/atomic_waker_turbo.rs
+++ b/embassy-sync/src/waitqueue/atomic_waker_turbo.rs
@@ -4,6 +4,10 @@ use core::sync::atomic::{AtomicPtr, Ordering};
 use core::task::Waker;
 
 /// Utility struct to register and wake a waker.
+/// If a waker is registered, registering another waker will replace the previous one without waking it.
+/// The intended use case is to wake tasks from interrupts. Therefore, it is generally not expected,
+/// that multiple tasks register try to register a waker simultaneously.
+#[derive(Debug)]
 pub struct AtomicWaker {
     waker: AtomicPtr<()>,
 }
diff --git a/embassy-sync/src/waitqueue/multi_waker.rs b/embassy-sync/src/waitqueue/multi_waker.rs
index 0e520bf40..56c0cd1b2 100644
--- a/embassy-sync/src/waitqueue/multi_waker.rs
+++ b/embassy-sync/src/waitqueue/multi_waker.rs
@@ -3,6 +3,9 @@ use core::task::Waker;
 use heapless::Vec;
 
 /// Utility struct to register and wake multiple wakers.
+/// Queue of wakers with a maximum length of `N`.
+/// Intended for waking multiple tasks.
+#[derive(Debug)]
 pub struct MultiWakerRegistration<const N: usize> {
     wakers: Vec<Waker, N>,
 }
@@ -13,7 +16,9 @@ impl<const N: usize> MultiWakerRegistration<N> {
         Self { wakers: Vec::new() }
     }
 
-    /// Register a waker. If the buffer is full the function returns it in the error
+    /// Register a waker.
+    ///
+    /// If the buffer is full, [wakes all the wakers](Self::wake), clears its buffer and registers the waker.
     pub fn register(&mut self, w: &Waker) {
         // If we already have some waker that wakes the same task as `w`, do nothing.
         // This avoids cloning wakers, and avoids unnecessary mass-wakes.
diff --git a/embassy-sync/src/waitqueue/waker_registration.rs b/embassy-sync/src/waitqueue/waker_registration.rs
index 9b666e7c4..7f24f8fb6 100644
--- a/embassy-sync/src/waitqueue/waker_registration.rs
+++ b/embassy-sync/src/waitqueue/waker_registration.rs
@@ -2,6 +2,10 @@ use core::mem;
 use core::task::Waker;
 
 /// Utility struct to register and wake a waker.
+/// If a waker is registered, registering another waker will replace the previous one.
+/// The previous waker will be woken in this case, giving it a chance to reregister itself.
+/// Although it is possible to wake multiple tasks this way,
+/// this will cause them to wake each other in a loop registering themselves.
 #[derive(Debug, Default)]
 pub struct WakerRegistration {
     waker: Option<Waker>,