diff --git a/embassy-sync/src/pubsub/mod.rs b/embassy-sync/src/pubsub/mod.rs index 59e701c58..6afd54af5 100644 --- a/embassy-sync/src/pubsub/mod.rs +++ b/embassy-sync/src/pubsub/mod.rs @@ -4,7 +4,7 @@ use core::cell::RefCell; use core::fmt::Debug; -use core::task::{Context, Poll, Waker}; +use core::task::{Context, Poll}; use heapless::Deque; @@ -179,7 +179,7 @@ impl { if let Some(cx) = cx { - s.register_subscriber_waker(cx.waker()); + s.subscriber_wakers.register(cx.waker()); } Poll::Pending } @@ -206,7 +206,7 @@ impl { if let Some(cx) = cx { - s.register_publisher_waker(cx.waker()); + s.publisher_wakers.register(cx.waker()); } Err(message) } @@ -335,34 +335,6 @@ impl PubSubSta Some(WaitResult::Message(message)) } - fn register_subscriber_waker(&mut self, waker: &Waker) { - match self.subscriber_wakers.register(waker) { - Ok(()) => {} - Err(_) => { - // All waker slots were full. This can only happen when there was a subscriber that now has dropped. - // We need to throw it away. 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.subscriber_wakers.wake(); - self.subscriber_wakers.register(waker).unwrap(); - } - } - } - - fn register_publisher_waker(&mut self, waker: &Waker) { - match self.publisher_wakers.register(waker) { - Ok(()) => {} - Err(_) => { - // All waker slots were full. This can only happen when there was a publisher that now has dropped. - // We need to throw it away. 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.publisher_wakers.wake(); - self.publisher_wakers.register(waker).unwrap(); - } - } - } - fn unregister_subscriber(&mut self, subscriber_next_message_id: u64) { self.subscriber_count -= 1; 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 { - wakers: [WakerRegistration; N], + wakers: Vec, } impl MultiWakerRegistration { /// 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(); } } }