3 files changed, 132 insertions, 0 deletions

diff --git a/embassy-util/src/waitqueue/mod.rs b/embassy-util/src/waitqueue/mod.rs
new file mode 100644
index 000000000..6661a6b61
--- /dev/null
+++ b/embassy-util/src/waitqueue/mod.rs
@@ -0,0 +1,7 @@
+//! Async low-level wait queues
+
+mod waker;
+pub use waker::*;
+
+mod multi_waker;
+pub use multi_waker::*;
diff --git a/embassy-util/src/waitqueue/multi_waker.rs b/embassy-util/src/waitqueue/multi_waker.rs
new file mode 100644
index 000000000..325d2cb3a
--- /dev/null
+++ b/embassy-util/src/waitqueue/multi_waker.rs
@@ -0,0 +1,33 @@
+use core::task::Waker;
+
+use super::WakerRegistration;
+
+/// Utility struct to register and wake multiple wakers.
+pub struct MultiWakerRegistration<const N: usize> {
+    wakers: [WakerRegistration; 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] }
+    }
+
+    /// 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)
+        }
+    }
+
+    /// Wake all registered wakers. This clears the buffer
+    pub fn wake(&mut self) {
+        for waker_slot in self.wakers.iter_mut() {
+            waker_slot.wake()
+        }
+    }
+}
diff --git a/embassy-util/src/waitqueue/waker.rs b/embassy-util/src/waitqueue/waker.rs
new file mode 100644
index 000000000..64e300eb8
--- /dev/null
+++ b/embassy-util/src/waitqueue/waker.rs
@@ -0,0 +1,92 @@
+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)]
+pub struct WakerRegistration {
+    waker: Option<Waker>,
+}
+
+impl WakerRegistration {
+    /// Create a new `WakerRegistration`.
+    pub const fn new() -> Self {
+        Self { waker: None }
+    }
+
+    /// Register a waker. Overwrites the previous waker, if any.
+    pub fn register(&mut self, w: &Waker) {
+        match self.waker {
+            // Optimization: If both the old and new Wakers wake the same task, we can simply
+            // keep the old waker, skipping the clone. (In most executor implementations,
+            // cloning a waker is somewhat expensive, comparable to cloning an Arc).
+            Some(ref w2) if (w2.will_wake(w)) => {}
+            _ => {
+                // clone the new waker and store it
+                if let Some(old_waker) = mem::replace(&mut self.waker, Some(w.clone())) {
+                    // We had a waker registered for another task. Wake it, so the other task can
+                    // reregister itself if it's still interested.
+                    //
+                    // If two tasks are waiting on the same thing concurrently, this will cause them
+                    // to wake each other in a loop fighting over this WakerRegistration. This wastes
+                    // CPU but things will still work.
+                    //
+                    // If the user wants to have two tasks waiting on the same thing they should use
+                    // a more appropriate primitive that can store multiple wakers.
+                    old_waker.wake()
+                }
+            }
+        }
+    }
+
+    /// Wake the registered waker, if any.
+    pub fn wake(&mut self) {
+        if let Some(w) = self.waker.take() {
+            w.wake()
+        }
+    }
+
+    /// Returns true if a waker is currently registered
+    pub fn occupied(&self) -> bool {
+        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));
+            }
+        })
+    }
+}