max/embassy
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added a function to WakerRegistration to check if it's occupied.

Browse Source 
Created a MultiWakerRegistration that can hold multiple wakers.
Got rid of some options and the pub/sub_index
This commit is contained in:
Dion Dokter 2022-06-17 13:54:34 +02:00
parent 2a4cdd05fa
commit eb304c2448
5 changed files with 120 additions and 100 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

172
embassy/src/channel/pubsub.rs
View File

@ -10,7 +10,7 @@ use heapless::Deque;
use crate::blocking_mutex::raw::RawMutex; use crate::blocking_mutex::raw::RawMutex;
use crate::blocking_mutex::Mutex; use crate::blocking_mutex::Mutex;
use crate::waitqueue::WakerRegistration; use crate::waitqueue::MultiWakerRegistration;
/// A broadcast channel implementation where multiple publishers can send messages to multiple subscribers /// A broadcast channel implementation where multiple publishers can send messages to multiple subscribers
/// ///
@ -42,21 +42,15 @@ impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usi
self.inner.lock(|inner| { self.inner.lock(|inner| {
let mut s = inner.borrow_mut(); let mut s = inner.borrow_mut();
// Search for an empty subscriber spot if s.subscriber_count >= SUBS {
for (i, sub_spot) in s.subscriber_wakers.iter_mut().enumerate() { Err(Error::MaximumSubscribersReached)
if sub_spot.is_none() { } else {
// We've found a spot, so now fill it and create the subscriber s.subscriber_count += 1;
*sub_spot = Some(WakerRegistration::new()); Ok(Subscriber {
return Ok(Subscriber {
subscriber_index: i,
next_message_id: s.next_message_id, next_message_id: s.next_message_id,
channel: self, channel: self,
}); })
} }
}
// No spot was found, we're full
Err(Error::MaximumSubscribersReached)
}) })
} }
@ -67,20 +61,12 @@ impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usi
self.inner.lock(|inner| { self.inner.lock(|inner| {
let mut s = inner.borrow_mut(); let mut s = inner.borrow_mut();
// Search for an empty publisher spot if s.publisher_count >= PUBS {
for (i, pub_spot) in s.publisher_wakers.iter_mut().enumerate() {
if pub_spot.is_none() {
// We've found a spot, so now fill it and create the subscriber
*pub_spot = Some(WakerRegistration::new());
return Ok(Publisher {
publisher_index: i,
channel: self,
});
}
}
// No spot was found, we're full
Err(Error::MaximumPublishersReached) Err(Error::MaximumPublishersReached)
} else {
s.publisher_count += 1;
Ok(Publisher { channel: self })
}
}) })
} }
@ -94,12 +80,7 @@ impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usi
impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubBehavior<T> impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubBehavior<T>
for PubSubChannel<M, T, CAP, SUBS, PUBS> for PubSubChannel<M, T, CAP, SUBS, PUBS>
{ {
fn get_message_with_context( fn get_message_with_context(&self, next_message_id: &mut u64, cx: Option<&mut Context<'_>>) -> Poll<WaitResult<T>> {
&self,
next_message_id: &mut u64,
subscriber_index: usize,
cx: Option<&mut Context<'_>>,
) -> Poll<WaitResult<T>> {
self.inner.lock(|s| { self.inner.lock(|s| {
let mut s = s.borrow_mut(); let mut s = s.borrow_mut();
@ -113,7 +94,7 @@ impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usi
// No, so we need to reregister our waker and sleep again // No, so we need to reregister our waker and sleep again
None => { None => {
if let Some(cx) = cx { if let Some(cx) = cx {
s.register_subscriber_waker(subscriber_index, cx.waker()); s.register_subscriber_waker(cx.waker());
} }
Poll::Pending Poll::Pending
} }
@ -126,7 +107,7 @@ impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usi
}) })
} }
fn publish_with_context(&self, message: T, publisher_index: usize, cx: Option<&mut Context<'_>>) -> Result<(), T> { fn publish_with_context(&self, message: T, cx: Option<&mut Context<'_>>) -> Result<(), T> {
self.inner.lock(|s| { self.inner.lock(|s| {
let mut s = s.borrow_mut(); let mut s = s.borrow_mut();
// Try to publish the message // Try to publish the message
@ -136,7 +117,7 @@ impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usi
// The queue is full, so we need to reregister our waker and go to sleep // The queue is full, so we need to reregister our waker and go to sleep
Err(message) => { Err(message) => {
if let Some(cx) = cx { if let Some(cx) = cx {
s.register_publisher_waker(publisher_index, cx.waker()); s.register_publisher_waker(cx.waker());
} }
Err(message) Err(message)
} }
@ -151,17 +132,17 @@ impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usi
}) })
} }
fn unregister_subscriber(&self, subscriber_index: usize, subscriber_next_message_id: u64) { fn unregister_subscriber(&self, subscriber_next_message_id: u64) {
self.inner.lock(|s| { self.inner.lock(|s| {
let mut s = s.borrow_mut(); let mut s = s.borrow_mut();
s.unregister_subscriber(subscriber_index, subscriber_next_message_id) s.unregister_subscriber(subscriber_next_message_id)
}) })
} }
fn unregister_publisher(&self, publisher_index: usize) { fn unregister_publisher(&self) {
self.inner.lock(|s| { self.inner.lock(|s| {
let mut s = s.borrow_mut(); let mut s = s.borrow_mut();
s.unregister_publisher(publisher_index) s.unregister_publisher()
}) })
} }
} }
@ -175,29 +156,30 @@ struct PubSubState<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: us
/// If a million messages were published every second, then the ID's would run out in about 584942 years. /// If a million messages were published every second, then the ID's would run out in about 584942 years.
next_message_id: u64, next_message_id: u64,
/// Collection of wakers for Subscribers that are waiting. /// Collection of wakers for Subscribers that are waiting.
/// The [Subscriber::subscriber_index] field indexes into this array. subscriber_wakers: MultiWakerRegistration<SUBS>,
subscriber_wakers: [Option<WakerRegistration>; SUBS],
/// Collection of wakers for Publishers that are waiting. /// Collection of wakers for Publishers that are waiting.
/// The [Publisher::publisher_index] field indexes into this array. publisher_wakers: MultiWakerRegistration<PUBS>,
publisher_wakers: [Option<WakerRegistration>; PUBS], /// The amount of subscribers that are active
subscriber_count: usize,
/// The amount of publishers that are active
publisher_count: usize,
} }
impl<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubState<T, CAP, SUBS, PUBS> { impl<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubState<T, CAP, SUBS, PUBS> {
/// Create a new internal channel state /// Create a new internal channel state
const fn new() -> Self { const fn new() -> Self {
const WAKER_INIT: Option<WakerRegistration> = None;
Self { Self {
queue: Deque::new(), queue: Deque::new(),
next_message_id: 0, next_message_id: 0,
subscriber_wakers: [WAKER_INIT; SUBS], subscriber_wakers: MultiWakerRegistration::new(),
publisher_wakers: [WAKER_INIT; PUBS], publisher_wakers: MultiWakerRegistration::new(),
subscriber_count: 0,
publisher_count: 0,
} }
} }
fn try_publish(&mut self, message: T) -> Result<(), T> { fn try_publish(&mut self, message: T) -> Result<(), T> {
let active_subscriber_count = self.subscriber_wakers.iter().flatten().count(); if self.subscriber_count == 0 {
if active_subscriber_count == 0 {
// We don't need to publish anything because there is no one to receive it // We don't need to publish anything because there is no one to receive it
return Ok(()); return Ok(());
} }
@ -206,14 +188,12 @@ impl<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubSta
return Err(message); return Err(message);
} }
// We just did a check for this // We just did a check for this
self.queue.push_back((message, active_subscriber_count)).ok().unwrap(); self.queue.push_back((message, self.subscriber_count)).ok().unwrap();
self.next_message_id += 1; self.next_message_id += 1;
// Wake all of the subscribers // Wake all of the subscribers
for active_subscriber in self.subscriber_wakers.iter_mut().flatten() { self.subscriber_wakers.wake();
active_subscriber.wake()
}
Ok(()) Ok(())
} }
@ -250,26 +230,42 @@ impl<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubSta
if current_message_index == 0 && queue_item.1 == 0 { if current_message_index == 0 && queue_item.1 == 0 {
self.queue.pop_front(); self.queue.pop_front();
self.publisher_wakers.iter_mut().flatten().for_each(|w| w.wake()); self.publisher_wakers.wake();
} }
Some(WaitResult::Message(message)) Some(WaitResult::Message(message))
} }
fn register_subscriber_waker(&mut self, subscriber_index: usize, waker: &Waker) { fn register_subscriber_waker(&mut self, waker: &Waker) {
self.subscriber_wakers[subscriber_index] match self.subscriber_wakers.register(waker) {
.as_mut() Ok(()) => {}
.unwrap() Err(_) => {
.register(waker); // 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, publisher_index: usize, waker: &Waker) { fn register_publisher_waker(&mut self, waker: &Waker) {
self.publisher_wakers[publisher_index].as_mut().unwrap().register(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_index: usize, subscriber_next_message_id: u64) { fn unregister_subscriber(&mut self, subscriber_next_message_id: u64) {
// Remove the subscriber from the wakers self.subscriber_count -= 1;
self.subscriber_wakers[subscriber_index] = None;
// All messages that haven't been read yet by this subscriber must have their counter decremented // All messages that haven't been read yet by this subscriber must have their counter decremented
let start_id = self.next_message_id - self.queue.len() as u64; let start_id = self.next_message_id - self.queue.len() as u64;
@ -282,9 +278,8 @@ impl<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubSta
} }
} }
fn unregister_publisher(&mut self, publisher_index: usize) { fn unregister_publisher(&mut self) {
// Remove the publisher from the wakers self.publisher_count -= 1;
self.publisher_wakers[publisher_index] = None;
} }
} }
@ -293,8 +288,6 @@ impl<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubSta
/// This instance carries a reference to the channel, but uses a trait object for it so that the channel's /// This instance carries a reference to the channel, but uses a trait object for it so that the channel's
/// generics are erased on this subscriber /// generics are erased on this subscriber
pub struct Subscriber<'a, T: Clone> { pub struct Subscriber<'a, T: Clone> {
/// Our index into the channel
subscriber_index: usize,
/// The message id of the next message we are yet to receive /// The message id of the next message we are yet to receive
next_message_id: u64, next_message_id: u64,
/// The channel we are a subscriber to /// The channel we are a subscriber to
@ -321,10 +314,7 @@ impl<'a, T: Clone> Subscriber<'a, T> {
/// ///
/// This function does not peek. The message is received if there is one. /// This function does not peek. The message is received if there is one.
pub fn try_next_message(&mut self) -> Option<WaitResult<T>> { pub fn try_next_message(&mut self) -> Option<WaitResult<T>> {
match self match self.channel.get_message_with_context(&mut self.next_message_id, None) {
.channel
.get_message_with_context(&mut self.next_message_id, self.subscriber_index, None)
{
Poll::Ready(result) => Some(result), Poll::Ready(result) => Some(result),
Poll::Pending => None, Poll::Pending => None,
} }
@ -346,8 +336,7 @@ impl<'a, T: Clone> Subscriber<'a, T> {
impl<'a, T: Clone> Drop for Subscriber<'a, T> { impl<'a, T: Clone> Drop for Subscriber<'a, T> {
fn drop(&mut self) { fn drop(&mut self) {
self.channel self.channel.unregister_subscriber(self.next_message_id)
.unregister_subscriber(self.subscriber_index, self.next_message_id)
} }
} }
@ -357,10 +346,9 @@ impl<'a, T: Clone> futures::Stream for Subscriber<'a, T> {
type Item = T; type Item = T;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let sub_index = self.subscriber_index;
match self match self
.channel .channel
.get_message_with_context(&mut self.next_message_id, sub_index, Some(cx)) .get_message_with_context(&mut self.next_message_id, Some(cx))
{ {
Poll::Ready(WaitResult::Message(message)) => Poll::Ready(Some(message)), Poll::Ready(WaitResult::Message(message)) => Poll::Ready(Some(message)),
Poll::Ready(WaitResult::Lagged(_)) => { Poll::Ready(WaitResult::Lagged(_)) => {
@ -377,8 +365,6 @@ impl<'a, T: Clone> futures::Stream for Subscriber<'a, T> {
/// This instance carries a reference to the channel, but uses a trait object for it so that the channel's /// This instance carries a reference to the channel, but uses a trait object for it so that the channel's
/// generics are erased on this subscriber /// generics are erased on this subscriber
pub struct Publisher<'a, T: Clone> { pub struct Publisher<'a, T: Clone> {
/// Our index into the channel
publisher_index: usize,
/// The channel we are a publisher for /// The channel we are a publisher for
channel: &'a dyn PubSubBehavior<T>, channel: &'a dyn PubSubBehavior<T>,
} }
@ -400,13 +386,13 @@ impl<'a, T: Clone> Publisher<'a, T> {
/// Publish a message if there is space in the message queue /// Publish a message if there is space in the message queue
pub fn try_publish(&self, message: T) -> Result<(), T> { pub fn try_publish(&self, message: T) -> Result<(), T> {
self.channel.publish_with_context(message, self.publisher_index, None) self.channel.publish_with_context(message, None)
} }
} }
impl<'a, T: Clone> Drop for Publisher<'a, T> { impl<'a, T: Clone> Drop for Publisher<'a, T> {
fn drop(&mut self) { fn drop(&mut self) {
self.channel.unregister_publisher(self.publisher_index) self.channel.unregister_publisher()
} }
} }
@ -426,7 +412,7 @@ impl<'a, T: Clone> ImmediatePublisher<'a, T> {
/// Publish a message if there is space in the message queue /// Publish a message if there is space in the message queue
pub fn try_publish(&self, message: T) -> Result<(), T> { pub fn try_publish(&self, message: T) -> Result<(), T> {
self.channel.publish_with_context(message, usize::MAX, None) self.channel.publish_with_context(message, None)
} }
} }
@ -442,20 +428,15 @@ pub enum Error {
} }
trait PubSubBehavior<T> { trait PubSubBehavior<T> {
fn get_message_with_context( fn get_message_with_context(&self, next_message_id: &mut u64, cx: Option<&mut Context<'_>>) -> Poll<WaitResult<T>>;
&self,
next_message_id: &mut u64,
subscriber_index: usize,
cx: Option<&mut Context<'_>>,
) -> Poll<WaitResult<T>>;
fn publish_with_context(&self, message: T, publisher_index: usize, cx: Option<&mut Context<'_>>) -> Result<(), T>; fn publish_with_context(&self, message: T, cx: Option<&mut Context<'_>>) -> Result<(), T>;
fn publish_immediate(&self, message: T); fn publish_immediate(&self, message: T);
fn unregister_subscriber(&self, subscriber_index: usize, subscriber_next_message_id: u64); fn unregister_subscriber(&self, subscriber_next_message_id: u64);
fn unregister_publisher(&self, publisher_index: usize); fn unregister_publisher(&self);
} }
/// Future for the subscriber wait action /// Future for the subscriber wait action
@ -467,10 +448,9 @@ impl<'s, 'a, T: Clone> Future for SubscriberWaitFuture<'s, 'a, T> {
type Output = WaitResult<T>; type Output = WaitResult<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let sub_index = self.subscriber.subscriber_index;
self.subscriber self.subscriber
.channel .channel
.get_message_with_context(&mut self.subscriber.next_message_id, sub_index, Some(cx)) .get_message_with_context(&mut self.subscriber.next_message_id, Some(cx))
} }
} }
@ -488,11 +468,7 @@ impl<'s, 'a, T: Clone> Future for PublisherWaitFuture<'s, 'a, T> {
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let message = self.message.take().unwrap(); let message = self.message.take().unwrap();
match self match self.publisher.channel.publish_with_context(message, Some(cx)) {
.publisher
.channel
.publish_with_context(message, self.publisher.publisher_index, Some(cx))
{
Ok(()) => Poll::Ready(()), Ok(()) => Poll::Ready(()),
Err(message) => { Err(message) => {
self.message = Some(message); self.message = Some(message);

3
embassy/src/waitqueue/mod.rs
View File

@ -3,3 +3,6 @@
#[cfg_attr(feature = "executor-agnostic", path = "waker_agnostic.rs")] #[cfg_attr(feature = "executor-agnostic", path = "waker_agnostic.rs")]
mod waker; mod waker;
pub use waker::*; pub use waker::*;
mod multi_waker;
pub use multi_waker::*;

31
embassy/src/waitqueue/multi_waker.rs Normal file
View File

@ -0,0 +1,31 @@
use core::task::Waker;
use super::WakerRegistration;
pub struct MultiWakerRegistration<const N: usize> {
wakers: [WakerRegistration; N],
}
impl<const N: usize> MultiWakerRegistration<N> {
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()
}
}
}

5
embassy/src/waitqueue/waker.rs
View File

@ -50,6 +50,11 @@ impl WakerRegistration {
unsafe { wake_task(w) } unsafe { wake_task(w) }
} }
} }
/// Returns true if a waker is currently registered
pub fn occupied(&self) -> bool {
self.waker.is_some()
}
} }
// SAFETY: `WakerRegistration` effectively contains an `Option<Waker>`, // SAFETY: `WakerRegistration` effectively contains an `Option<Waker>`,

5
embassy/src/waitqueue/waker_agnostic.rs
View File

@ -47,6 +47,11 @@ impl WakerRegistration {
w.wake() 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. /// Utility struct to register and wake a waker.