diff options
Diffstat (limited to 'embassy-sync/src/waitqueue')
-rw-r--r-- | embassy-sync/src/waitqueue/mod.rs | 7 | ||||
-rw-r--r-- | embassy-sync/src/waitqueue/multi_waker.rs | 33 | ||||
-rw-r--r-- | embassy-sync/src/waitqueue/waker.rs | 92 |
3 files changed, 132 insertions, 0 deletions
diff --git a/embassy-sync/src/waitqueue/mod.rs b/embassy-sync/src/waitqueue/mod.rs new file mode 100644 index 00000000..6661a6b6 --- /dev/null +++ b/embassy-sync/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-sync/src/waitqueue/multi_waker.rs b/embassy-sync/src/waitqueue/multi_waker.rs new file mode 100644 index 00000000..325d2cb3 --- /dev/null +++ b/embassy-sync/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-sync/src/waitqueue/waker.rs b/embassy-sync/src/waitqueue/waker.rs new file mode 100644 index 00000000..64e300eb --- /dev/null +++ b/embassy-sync/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)); + } + }) + } +} |