5 files changed, 115 insertions, 55 deletions

diff --git a/embassy-sync/src/waitqueue/atomic_waker.rs b/embassy-sync/src/waitqueue/atomic_waker.rs
new file mode 100644
index 000000000..63fe04a6e
--- /dev/null
+++ b/embassy-sync/src/waitqueue/atomic_waker.rs
@@ -0,0 +1,41 @@
+use core::cell::Cell;
+use core::task::Waker;
+
+use crate::blocking_mutex::raw::CriticalSectionRawMutex;
+use crate::blocking_mutex::Mutex;
+
+/// Utility struct to register and wake a waker.
+pub struct AtomicWaker {
+    waker: Mutex<CriticalSectionRawMutex, Cell<Option<Waker>>>,
+}
+
+impl AtomicWaker {
+    /// Create a new `AtomicWaker`.
+    pub const fn new() -> Self {
+        Self {
+            waker: Mutex::const_new(CriticalSectionRawMutex::new(), 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);
+            cell.set(match cell.replace(None) {
+                Some(w2) if (w2.will_wake(w)) => Some(w2),
+                _ => Some(w.clone()),
+            })
+        })
+    }
+
+    /// Wake the registered waker, if any.
+    pub fn wake(&self) {
+        critical_section::with(|cs| {
+            let cell = self.waker.borrow(cs);
+            if let Some(w) = cell.replace(None) {
+                w.wake_by_ref();
+                cell.set(Some(w));
+            }
+        })
+    }
+}
diff --git a/embassy-sync/src/waitqueue/atomic_waker_turbo.rs b/embassy-sync/src/waitqueue/atomic_waker_turbo.rs
new file mode 100644
index 000000000..5c6a96ec8
--- /dev/null
+++ b/embassy-sync/src/waitqueue/atomic_waker_turbo.rs
@@ -0,0 +1,30 @@
+use core::ptr;
+use core::ptr::NonNull;
+use core::sync::atomic::{AtomicPtr, Ordering};
+use core::task::Waker;
+
+/// Utility struct to register and wake a waker.
+pub struct AtomicWaker {
+    waker: AtomicPtr<()>,
+}
+
+impl AtomicWaker {
+    /// Create a new `AtomicWaker`.
+    pub const fn new() -> Self {
+        Self {
+            waker: AtomicPtr::new(ptr::null_mut()),
+        }
+    }
+
+    /// Register a waker. Overwrites the previous waker, if any.
+    pub fn register(&self, w: &Waker) {
+        self.waker.store(w.as_turbo_ptr().as_ptr() as _, Ordering::Release);
+    }
+
+    /// Wake the registered waker, if any.
+    pub fn wake(&self) {
+        if let Some(ptr) = NonNull::new(self.waker.load(Ordering::Acquire)) {
+            unsafe { Waker::from_turbo_ptr(ptr) }.wake();
+        }
+    }
+}
diff --git a/embassy-sync/src/waitqueue/mod.rs b/embassy-sync/src/waitqueue/mod.rs
index 6661a6b61..6b0b0c64e 100644
--- a/embassy-sync/src/waitqueue/mod.rs
+++ b/embassy-sync/src/waitqueue/mod.rs
@@ -1,7 +1,11 @@
 //! Async low-level wait queues
 
-mod waker;
-pub use waker::*;
+#[cfg_attr(feature = "turbowakers", path = "atomic_waker_turbo.rs")]
+mod atomic_waker;
+pub use atomic_waker::*;
+
+mod waker_registration;
+pub use waker_registration::*;
 
 mod multi_waker;
 pub use multi_waker::*;
diff --git a/embassy-sync/src/waitqueue/multi_waker.rs b/embassy-sync/src/waitqueue/multi_waker.rs
index 325d2cb3a..824d192da 100644
--- a/embassy-sync/src/waitqueue/multi_waker.rs
+++ b/embassy-sync/src/waitqueue/multi_waker.rs
@@ -1,33 +1,58 @@
 use core::task::Waker;
 
-use super::WakerRegistration;
+use heapless::Vec;
 
 /// Utility struct to register and wake multiple wakers.
 pub struct MultiWakerRegistration<const N: usize> {
-    wakers: [WakerRegistration; N],
+    wakers: Vec<Waker, N>,
 }
 
 impl<const N: usize> MultiWakerRegistration<N> {
     /// Create a new empty instance
     pub const fn new() -> Self {
-        const WAKER: WakerRegistration = WakerRegistration::new();
-        Self { wakers: [WAKER; N] }
+        Self { wakers: Vec::new() }
     }
 
     /// Register a waker. If the buffer is full the function returns it in the error
-    pub fn register<'a>(&mut self, w: &'a Waker) -> Result<(), &'a Waker> {
-        if let Some(waker_slot) = self.wakers.iter_mut().find(|waker_slot| !waker_slot.occupied()) {
-            waker_slot.register(w);
-            Ok(())
-        } else {
-            Err(w)
+    pub fn register<'a>(&mut self, w: &'a 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.
+        for w2 in &self.wakers {
+            if w.will_wake(w2) {
+                return;
+            }
+        }
+
+        if self.wakers.is_full() {
+            // All waker slots were full. It's a bit inefficient, but we can wake everything.
+            // Any future that is still active will simply reregister.
+            // This won't happen a lot, so it's ok.
+            self.wake();
+        }
+
+        if self.wakers.push(w.clone()).is_err() {
+            // This can't happen unless N=0
+            // (Either `wakers` wasn't full, or it was in which case `wake()` empied it)
+            panic!("tried to push a waker to a zero-length MultiWakerRegistration")
         }
     }
 
     /// Wake all registered wakers. This clears the buffer
     pub fn wake(&mut self) {
-        for waker_slot in self.wakers.iter_mut() {
-            waker_slot.wake()
+        // heapless::Vec has no `drain()`, do it unsafely ourselves...
+
+        // First set length to 0, without dropping the contents.
+        // This is necessary for soundness: if wake() panics and we're using panic=unwind.
+        // Setting len=0 upfront ensures other code can't observe the vec in an inconsistent state.
+        // (it'll leak wakers, but that's not UB)
+        let len = self.wakers.len();
+        unsafe { self.wakers.set_len(0) }
+
+        for i in 0..len {
+            // Move a waker out of the vec.
+            let waker = unsafe { self.wakers.as_mut_ptr().add(i).read() };
+            // Wake it by value, which consumes (drops) it.
+            waker.wake();
         }
     }
 }
diff --git a/embassy-sync/src/waitqueue/waker.rs b/embassy-sync/src/waitqueue/waker_registration.rs
index 64e300eb8..9b666e7c4 100644
--- a/embassy-sync/src/waitqueue/waker.rs
+++ b/embassy-sync/src/waitqueue/waker_registration.rs
@@ -1,12 +1,8 @@
-use core::cell::Cell;
 use core::mem;
 use core::task::Waker;
 
-use crate::blocking_mutex::raw::CriticalSectionRawMutex;
-use crate::blocking_mutex::Mutex;
-
 /// Utility struct to register and wake a waker.
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct WakerRegistration {
     waker: Option<Waker>,
 }
@@ -54,39 +50,3 @@ impl WakerRegistration {
         self.waker.is_some()
     }
 }
-
-/// Utility struct to register and wake a waker.
-pub struct AtomicWaker {
-    waker: Mutex<CriticalSectionRawMutex, Cell<Option<Waker>>>,
-}
-
-impl AtomicWaker {
-    /// Create a new `AtomicWaker`.
-    pub const fn new() -> Self {
-        Self {
-            waker: Mutex::const_new(CriticalSectionRawMutex::new(), 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);
-            cell.set(match cell.replace(None) {
-                Some(w2) if (w2.will_wake(w)) => Some(w2),
-                _ => Some(w.clone()),
-            })
-        })
-    }
-
-    /// Wake the registered waker, if any.
-    pub fn wake(&self) {
-        critical_section::with(|cs| {
-            let cell = self.waker.borrow(cs);
-            if let Some(w) = cell.replace(None) {
-                w.wake_by_ref();
-                cell.set(Some(w));
-            }
-        })
-    }
-}