#![deny(missing_docs)]
use core::cell::RefCell;
use core::fmt::Debug;
use core::task::{Context, Poll};
use heapless::Deque;
use self::publisher::{ImmediatePub, Pub};
use self::subscriber::Sub;
use crate::blocking_mutex::raw::RawMutex;
use crate::blocking_mutex::Mutex;
use crate::waitqueue::MultiWakerRegistration;
pub mod publisher;
pub mod subscriber;
pub use publisher::{DynImmediatePublisher, DynPublisher, ImmediatePublisher, Publisher};
pub use subscriber::{DynSubscriber, Subscriber};
pub struct PubSubChannel<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> {
inner: Mutex<M, RefCell<PubSubState<T, CAP, SUBS, PUBS>>>,
}
impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize>
PubSubChannel<M, T, CAP, SUBS, PUBS>
{
pub const fn new() -> Self {
Self {
inner: Mutex::const_new(M::INIT, RefCell::new(PubSubState::new())),
}
}
pub fn subscriber(&self) -> Result<Subscriber<M, T, CAP, SUBS, PUBS>, Error> {
self.inner.lock(|inner| {
let mut s = inner.borrow_mut();
if s.subscriber_count >= SUBS {
Err(Error::MaximumSubscribersReached)
} else {
s.subscriber_count += 1;
Ok(Subscriber(Sub::new(s.next_message_id, self)))
}
})
}
pub fn dyn_subscriber(&self) -> Result<DynSubscriber<'_, T>, Error> {
self.inner.lock(|inner| {
let mut s = inner.borrow_mut();
if s.subscriber_count >= SUBS {
Err(Error::MaximumSubscribersReached)
} else {
s.subscriber_count += 1;
Ok(DynSubscriber(Sub::new(s.next_message_id, self)))
}
})
}
pub fn publisher(&self) -> Result<Publisher<M, T, CAP, SUBS, PUBS>, Error> {
self.inner.lock(|inner| {
let mut s = inner.borrow_mut();
if s.publisher_count >= PUBS {
Err(Error::MaximumPublishersReached)
} else {
s.publisher_count += 1;
Ok(Publisher(Pub::new(self)))
}
})
}
pub fn dyn_publisher(&self) -> Result<DynPublisher<'_, T>, Error> {
self.inner.lock(|inner| {
let mut s = inner.borrow_mut();
if s.publisher_count >= PUBS {
Err(Error::MaximumPublishersReached)
} else {
s.publisher_count += 1;
Ok(DynPublisher(Pub::new(self)))
}
})
}
pub fn immediate_publisher(&self) -> ImmediatePublisher<M, T, CAP, SUBS, PUBS> {
ImmediatePublisher(ImmediatePub::new(self))
}
pub fn dyn_immediate_publisher(&self) -> DynImmediatePublisher<T> {
DynImmediatePublisher(ImmediatePub::new(self))
}
}
impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubBehavior<T>
for PubSubChannel<M, T, CAP, SUBS, PUBS>
{
fn get_message_with_context(&self, next_message_id: &mut u64, cx: Option<&mut Context<'_>>) -> Poll<WaitResult<T>> {
self.inner.lock(|s| {
let mut s = s.borrow_mut();
match s.get_message(*next_message_id) {
Some(WaitResult::Message(message)) => {
*next_message_id += 1;
Poll::Ready(WaitResult::Message(message))
}
None => {
if let Some(cx) = cx {
s.subscriber_wakers.register(cx.waker());
}
Poll::Pending
}
Some(WaitResult::Lagged(amount)) => {
*next_message_id += amount;
Poll::Ready(WaitResult::Lagged(amount))
}
}
})
}
fn available(&self, next_message_id: u64) -> u64 {
self.inner.lock(|s| s.borrow().next_message_id - next_message_id)
}
fn publish_with_context(&self, message: T, cx: Option<&mut Context<'_>>) -> Result<(), T> {
self.inner.lock(|s| {
let mut s = s.borrow_mut();
match s.try_publish(message) {
Ok(()) => Ok(()),
Err(message) => {
if let Some(cx) = cx {
s.publisher_wakers.register(cx.waker());
}
Err(message)
}
}
})
}
fn publish_immediate(&self, message: T) {
self.inner.lock(|s| {
let mut s = s.borrow_mut();
s.publish_immediate(message)
})
}
fn space(&self) -> usize {
self.inner.lock(|s| {
let s = s.borrow();
s.queue.capacity() - s.queue.len()
})
}
fn unregister_subscriber(&self, subscriber_next_message_id: u64) {
self.inner.lock(|s| {
let mut s = s.borrow_mut();
s.unregister_subscriber(subscriber_next_message_id)
})
}
fn unregister_publisher(&self) {
self.inner.lock(|s| {
let mut s = s.borrow_mut();
s.unregister_publisher()
})
}
}
struct PubSubState<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> {
queue: Deque<(T, usize), CAP>,
next_message_id: u64,
subscriber_wakers: MultiWakerRegistration<SUBS>,
publisher_wakers: MultiWakerRegistration<PUBS>,
subscriber_count: usize,
publisher_count: usize,
}
impl<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubState<T, CAP, SUBS, PUBS> {
const fn new() -> Self {
Self {
queue: Deque::new(),
next_message_id: 0,
subscriber_wakers: MultiWakerRegistration::new(),
publisher_wakers: MultiWakerRegistration::new(),
subscriber_count: 0,
publisher_count: 0,
}
}
fn try_publish(&mut self, message: T) -> Result<(), T> {
if self.subscriber_count == 0 {
return Ok(());
}
if self.queue.is_full() {
return Err(message);
}
self.queue.push_back((message, self.subscriber_count)).ok().unwrap();
self.next_message_id += 1;
self.subscriber_wakers.wake();
Ok(())
}
fn publish_immediate(&mut self, message: T) {
if self.queue.is_full() {
self.queue.pop_front();
}
self.try_publish(message).ok().unwrap();
}
fn get_message(&mut self, message_id: u64) -> Option<WaitResult<T>> {
let start_id = self.next_message_id - self.queue.len() as u64;
if message_id < start_id {
return Some(WaitResult::Lagged(start_id - message_id));
}
let current_message_index = (message_id - start_id) as usize;
if current_message_index >= self.queue.len() {
return None;
}
let queue_item = self.queue.iter_mut().nth(current_message_index).unwrap();
queue_item.1 -= 1;
let message = if current_message_index == 0 && queue_item.1 == 0 {
let (message, _) = self.queue.pop_front().unwrap();
self.publisher_wakers.wake();
message
} else {
queue_item.0.clone()
};
Some(WaitResult::Message(message))
}
fn unregister_subscriber(&mut self, subscriber_next_message_id: u64) {
self.subscriber_count -= 1;
let start_id = self.next_message_id - self.queue.len() as u64;
if subscriber_next_message_id >= start_id {
let current_message_index = (subscriber_next_message_id - start_id) as usize;
self.queue
.iter_mut()
.skip(current_message_index)
.for_each(|(_, counter)| *counter -= 1);
let mut wake_publishers = false;
while let Some((_, count)) = self.queue.front() {
if *count == 0 {
self.queue.pop_front().unwrap();
wake_publishers = true;
} else {
break;
}
}
if wake_publishers {
self.publisher_wakers.wake();
}
}
}
fn unregister_publisher(&mut self) {
self.publisher_count -= 1;
}
}
#[derive(Debug, PartialEq, Eq, Clone)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub enum Error {
MaximumSubscribersReached,
MaximumPublishersReached,
}
pub trait PubSubBehavior<T> {
fn get_message_with_context(&self, next_message_id: &mut u64, cx: Option<&mut Context<'_>>) -> Poll<WaitResult<T>>;
fn available(&self, next_message_id: u64) -> u64;
fn publish_with_context(&self, message: T, cx: Option<&mut Context<'_>>) -> Result<(), T>;
fn publish_immediate(&self, message: T);
fn space(&self) -> usize;
fn unregister_subscriber(&self, subscriber_next_message_id: u64);
fn unregister_publisher(&self);
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub enum WaitResult<T> {
Lagged(u64),
Message(T),
}
#[cfg(test)]
mod tests {
use super::*;
use crate::blocking_mutex::raw::NoopRawMutex;
#[futures_test::test]
async fn dyn_pub_sub_works() {
let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
let mut sub0 = channel.dyn_subscriber().unwrap();
let mut sub1 = channel.dyn_subscriber().unwrap();
let pub0 = channel.dyn_publisher().unwrap();
pub0.publish(42).await;
assert_eq!(sub0.next_message().await, WaitResult::Message(42));
assert_eq!(sub1.next_message().await, WaitResult::Message(42));
assert_eq!(sub0.try_next_message(), None);
assert_eq!(sub1.try_next_message(), None);
}
#[futures_test::test]
async fn all_subscribers_receive() {
let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
let mut sub0 = channel.subscriber().unwrap();
let mut sub1 = channel.subscriber().unwrap();
let pub0 = channel.publisher().unwrap();
pub0.publish(42).await;
assert_eq!(sub0.next_message().await, WaitResult::Message(42));
assert_eq!(sub1.next_message().await, WaitResult::Message(42));
assert_eq!(sub0.try_next_message(), None);
assert_eq!(sub1.try_next_message(), None);
}
#[futures_test::test]
async fn lag_when_queue_full_on_immediate_publish() {
let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
let mut sub0 = channel.subscriber().unwrap();
let pub0 = channel.publisher().unwrap();
pub0.publish_immediate(42);
pub0.publish_immediate(43);
pub0.publish_immediate(44);
pub0.publish_immediate(45);
pub0.publish_immediate(46);
pub0.publish_immediate(47);
assert_eq!(sub0.try_next_message(), Some(WaitResult::Lagged(2)));
assert_eq!(sub0.next_message().await, WaitResult::Message(44));
assert_eq!(sub0.next_message().await, WaitResult::Message(45));
assert_eq!(sub0.next_message().await, WaitResult::Message(46));
assert_eq!(sub0.next_message().await, WaitResult::Message(47));
assert_eq!(sub0.try_next_message(), None);
}
#[test]
fn limited_subs_and_pubs() {
let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
let sub0 = channel.subscriber();
let sub1 = channel.subscriber();
let sub2 = channel.subscriber();
let sub3 = channel.subscriber();
let sub4 = channel.subscriber();
assert!(sub0.is_ok());
assert!(sub1.is_ok());
assert!(sub2.is_ok());
assert!(sub3.is_ok());
assert_eq!(sub4.err().unwrap(), Error::MaximumSubscribersReached);
drop(sub0);
let sub5 = channel.subscriber();
assert!(sub5.is_ok());
let pub0 = channel.publisher();
let pub1 = channel.publisher();
let pub2 = channel.publisher();
let pub3 = channel.publisher();
let pub4 = channel.publisher();
assert!(pub0.is_ok());
assert!(pub1.is_ok());
assert!(pub2.is_ok());
assert!(pub3.is_ok());
assert_eq!(pub4.err().unwrap(), Error::MaximumPublishersReached);
drop(pub0);
let pub5 = channel.publisher();
assert!(pub5.is_ok());
}
#[test]
fn publisher_wait_on_full_queue() {
let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
let pub0 = channel.publisher().unwrap();
assert_eq!(pub0.try_publish(0), Ok(()));
assert_eq!(pub0.try_publish(0), Ok(()));
assert_eq!(pub0.try_publish(0), Ok(()));
assert_eq!(pub0.try_publish(0), Ok(()));
assert_eq!(pub0.try_publish(0), Ok(()));
let sub0 = channel.subscriber().unwrap();
assert_eq!(pub0.try_publish(0), Ok(()));
assert_eq!(pub0.try_publish(0), Ok(()));
assert_eq!(pub0.try_publish(0), Ok(()));
assert_eq!(pub0.try_publish(0), Ok(()));
assert_eq!(pub0.try_publish(0), Err(0));
drop(sub0);
}
#[futures_test::test]
async fn correct_available() {
let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
let sub0 = channel.subscriber().unwrap();
let mut sub1 = channel.subscriber().unwrap();
let pub0 = channel.publisher().unwrap();
assert_eq!(sub0.available(), 0);
assert_eq!(sub1.available(), 0);
pub0.publish(42).await;
assert_eq!(sub0.available(), 1);
assert_eq!(sub1.available(), 1);
sub1.next_message().await;
assert_eq!(sub1.available(), 0);
pub0.publish(42).await;
assert_eq!(sub0.available(), 2);
assert_eq!(sub1.available(), 1);
}
#[futures_test::test]
async fn correct_space() {
let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
let mut sub0 = channel.subscriber().unwrap();
let mut sub1 = channel.subscriber().unwrap();
let pub0 = channel.publisher().unwrap();
assert_eq!(pub0.space(), 4);
pub0.publish(42).await;
assert_eq!(pub0.space(), 3);
pub0.publish(42).await;
assert_eq!(pub0.space(), 2);
sub0.next_message().await;
sub0.next_message().await;
assert_eq!(pub0.space(), 2);
sub1.next_message().await;
assert_eq!(pub0.space(), 3);
sub1.next_message().await;
assert_eq!(pub0.space(), 4);
}
#[futures_test::test]
async fn empty_channel_when_last_subscriber_is_dropped() {
let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
let pub0 = channel.publisher().unwrap();
let mut sub0 = channel.subscriber().unwrap();
let mut sub1 = channel.subscriber().unwrap();
assert_eq!(4, pub0.space());
pub0.publish(1).await;
pub0.publish(2).await;
assert_eq!(2, channel.space());
assert_eq!(1, sub0.try_next_message_pure().unwrap());
assert_eq!(2, sub0.try_next_message_pure().unwrap());
assert_eq!(2, channel.space());
drop(sub0);
assert_eq!(2, channel.space());
assert_eq!(1, sub1.try_next_message_pure().unwrap());
assert_eq!(3, channel.space());
drop(sub1);
assert_eq!(4, channel.space());
}
struct CloneCallCounter(usize);
impl Clone for CloneCallCounter {
fn clone(&self) -> Self {
Self(self.0 + 1)
}
}
#[futures_test::test]
async fn skip_clone_for_last_message() {
let channel = PubSubChannel::<NoopRawMutex, CloneCallCounter, 1, 2, 1>::new();
let pub0 = channel.publisher().unwrap();
let mut sub0 = channel.subscriber().unwrap();
let mut sub1 = channel.subscriber().unwrap();
pub0.publish(CloneCallCounter(0)).await;
assert_eq!(1, sub0.try_next_message_pure().unwrap().0);
assert_eq!(0, sub1.try_next_message_pure().unwrap().0);
}
}