aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--embassy/src/channel/mod.rs2
-rw-r--r--embassy/src/channel/pubsub/mod.rs542
-rw-r--r--embassy/src/channel/pubsub/publisher.rs183
-rw-r--r--embassy/src/channel/pubsub/subscriber.rs153
-rw-r--r--embassy/src/waitqueue/mod.rs3
-rw-r--r--embassy/src/waitqueue/multi_waker.rs33
-rw-r--r--embassy/src/waitqueue/waker.rs5
-rw-r--r--embassy/src/waitqueue/waker_agnostic.rs5
-rw-r--r--examples/nrf/src/bin/pubsub.rs106
9 files changed, 1031 insertions, 1 deletions
diff --git a/embassy/src/channel/mod.rs b/embassy/src/channel/mod.rs
index 05edc55d1..5df1f5c5c 100644
--- a/embassy/src/channel/mod.rs
+++ b/embassy/src/channel/mod.rs
@@ -1,5 +1,5 @@
1//! Async channels 1//! Async channels
2 2
3pub mod mpmc; 3pub mod mpmc;
4 4pub mod pubsub;
5pub mod signal; 5pub mod signal;
diff --git a/embassy/src/channel/pubsub/mod.rs b/embassy/src/channel/pubsub/mod.rs
new file mode 100644
index 000000000..9bfb845e0
--- /dev/null
+++ b/embassy/src/channel/pubsub/mod.rs
@@ -0,0 +1,542 @@
1//! Implementation of [PubSubChannel], a queue where published messages get received by all subscribers.
2
3#![deny(missing_docs)]
4
5use core::cell::RefCell;
6use core::fmt::Debug;
7use core::task::{Context, Poll, Waker};
8
9use heapless::Deque;
10
11use self::publisher::{ImmediatePub, Pub};
12use self::subscriber::Sub;
13use crate::blocking_mutex::raw::RawMutex;
14use crate::blocking_mutex::Mutex;
15use crate::waitqueue::MultiWakerRegistration;
16
17pub mod publisher;
18pub mod subscriber;
19
20pub use publisher::{DynImmediatePublisher, DynPublisher, ImmediatePublisher, Publisher};
21pub use subscriber::{DynSubscriber, Subscriber};
22
23/// A broadcast channel implementation where multiple publishers can send messages to multiple subscribers
24///
25/// Any published message can be read by all subscribers.
26/// A publisher can choose how it sends its message.
27///
28/// - With [Publisher::publish] the publisher has to wait until there is space in the internal message queue.
29/// - With [Publisher::publish_immediate] the publisher doesn't await and instead lets the oldest message
30/// in the queue drop if necessary. This will cause any [Subscriber] that missed the message to receive
31/// an error to indicate that it has lagged.
32///
33/// ## Example
34///
35/// ```
36/// # use embassy::blocking_mutex::raw::NoopRawMutex;
37/// # use embassy::channel::pubsub::WaitResult;
38/// # use embassy::channel::pubsub::PubSubChannel;
39/// # use futures_executor::block_on;
40/// # let test = async {
41/// // Create the channel. This can be static as well
42/// let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
43///
44/// // This is a generic subscriber with a direct reference to the channel
45/// let mut sub0 = channel.subscriber().unwrap();
46/// // This is a dynamic subscriber with a dynamic (trait object) reference to the channel
47/// let mut sub1 = channel.dyn_subscriber().unwrap();
48///
49/// let pub0 = channel.publisher().unwrap();
50///
51/// // Publish a message, but wait if the queue is full
52/// pub0.publish(42).await;
53///
54/// // Publish a message, but if the queue is full, just kick out the oldest message.
55/// // This may cause some subscribers to miss a message
56/// pub0.publish_immediate(43);
57///
58/// // Wait for a new message. If the subscriber missed a message, the WaitResult will be a Lag result
59/// assert_eq!(sub0.next_message().await, WaitResult::Message(42));
60/// assert_eq!(sub1.next_message().await, WaitResult::Message(42));
61///
62/// // Wait again, but this time ignore any Lag results
63/// assert_eq!(sub0.next_message_pure().await, 43);
64/// assert_eq!(sub1.next_message_pure().await, 43);
65///
66/// // There's also a polling interface
67/// assert_eq!(sub0.try_next_message(), None);
68/// assert_eq!(sub1.try_next_message(), None);
69/// # };
70/// #
71/// # block_on(test);
72/// ```
73///
74pub struct PubSubChannel<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> {
75 inner: Mutex<M, RefCell<PubSubState<T, CAP, SUBS, PUBS>>>,
76}
77
78impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize>
79 PubSubChannel<M, T, CAP, SUBS, PUBS>
80{
81 /// Create a new channel
82 pub const fn new() -> Self {
83 Self {
84 inner: Mutex::const_new(M::INIT, RefCell::new(PubSubState::new())),
85 }
86 }
87
88 /// Create a new subscriber. It will only receive messages that are published after its creation.
89 ///
90 /// If there are no subscriber slots left, an error will be returned.
91 pub fn subscriber(&self) -> Result<Subscriber<M, T, CAP, SUBS, PUBS>, Error> {
92 self.inner.lock(|inner| {
93 let mut s = inner.borrow_mut();
94
95 if s.subscriber_count >= SUBS {
96 Err(Error::MaximumSubscribersReached)
97 } else {
98 s.subscriber_count += 1;
99 Ok(Subscriber(Sub::new(s.next_message_id, self)))
100 }
101 })
102 }
103
104 /// Create a new subscriber. It will only receive messages that are published after its creation.
105 ///
106 /// If there are no subscriber slots left, an error will be returned.
107 pub fn dyn_subscriber<'a>(&'a self) -> Result<DynSubscriber<'a, T>, Error> {
108 self.inner.lock(|inner| {
109 let mut s = inner.borrow_mut();
110
111 if s.subscriber_count >= SUBS {
112 Err(Error::MaximumSubscribersReached)
113 } else {
114 s.subscriber_count += 1;
115 Ok(DynSubscriber(Sub::new(s.next_message_id, self)))
116 }
117 })
118 }
119
120 /// Create a new publisher
121 ///
122 /// If there are no publisher slots left, an error will be returned.
123 pub fn publisher(&self) -> Result<Publisher<M, T, CAP, SUBS, PUBS>, Error> {
124 self.inner.lock(|inner| {
125 let mut s = inner.borrow_mut();
126
127 if s.publisher_count >= PUBS {
128 Err(Error::MaximumPublishersReached)
129 } else {
130 s.publisher_count += 1;
131 Ok(Publisher(Pub::new(self)))
132 }
133 })
134 }
135
136 /// Create a new publisher
137 ///
138 /// If there are no publisher slots left, an error will be returned.
139 pub fn dyn_publisher<'a>(&'a self) -> Result<DynPublisher<'a, T>, Error> {
140 self.inner.lock(|inner| {
141 let mut s = inner.borrow_mut();
142
143 if s.publisher_count >= PUBS {
144 Err(Error::MaximumPublishersReached)
145 } else {
146 s.publisher_count += 1;
147 Ok(DynPublisher(Pub::new(self)))
148 }
149 })
150 }
151
152 /// Create a new publisher that can only send immediate messages.
153 /// This kind of publisher does not take up a publisher slot.
154 pub fn immediate_publisher(&self) -> ImmediatePublisher<M, T, CAP, SUBS, PUBS> {
155 ImmediatePublisher(ImmediatePub::new(self))
156 }
157
158 /// Create a new publisher that can only send immediate messages.
159 /// This kind of publisher does not take up a publisher slot.
160 pub fn dyn_immediate_publisher(&self) -> DynImmediatePublisher<T> {
161 DynImmediatePublisher(ImmediatePub::new(self))
162 }
163}
164
165impl<M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubBehavior<T>
166 for PubSubChannel<M, T, CAP, SUBS, PUBS>
167{
168 fn get_message_with_context(&self, next_message_id: &mut u64, cx: Option<&mut Context<'_>>) -> Poll<WaitResult<T>> {
169 self.inner.lock(|s| {
170 let mut s = s.borrow_mut();
171
172 // Check if we can read a message
173 match s.get_message(*next_message_id) {
174 // Yes, so we are done polling
175 Some(WaitResult::Message(message)) => {
176 *next_message_id += 1;
177 Poll::Ready(WaitResult::Message(message))
178 }
179 // No, so we need to reregister our waker and sleep again
180 None => {
181 if let Some(cx) = cx {
182 s.register_subscriber_waker(cx.waker());
183 }
184 Poll::Pending
185 }
186 // We missed a couple of messages. We must do our internal bookkeeping and return that we lagged
187 Some(WaitResult::Lagged(amount)) => {
188 *next_message_id += amount;
189 Poll::Ready(WaitResult::Lagged(amount))
190 }
191 }
192 })
193 }
194
195 fn publish_with_context(&self, message: T, cx: Option<&mut Context<'_>>) -> Result<(), T> {
196 self.inner.lock(|s| {
197 let mut s = s.borrow_mut();
198 // Try to publish the message
199 match s.try_publish(message) {
200 // We did it, we are ready
201 Ok(()) => Ok(()),
202 // The queue is full, so we need to reregister our waker and go to sleep
203 Err(message) => {
204 if let Some(cx) = cx {
205 s.register_publisher_waker(cx.waker());
206 }
207 Err(message)
208 }
209 }
210 })
211 }
212
213 fn publish_immediate(&self, message: T) {
214 self.inner.lock(|s| {
215 let mut s = s.borrow_mut();
216 s.publish_immediate(message)
217 })
218 }
219
220 fn unregister_subscriber(&self, subscriber_next_message_id: u64) {
221 self.inner.lock(|s| {
222 let mut s = s.borrow_mut();
223 s.unregister_subscriber(subscriber_next_message_id)
224 })
225 }
226
227 fn unregister_publisher(&self) {
228 self.inner.lock(|s| {
229 let mut s = s.borrow_mut();
230 s.unregister_publisher()
231 })
232 }
233}
234
235/// Internal state for the PubSub channel
236struct PubSubState<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> {
237 /// The queue contains the last messages that have been published and a countdown of how many subscribers are yet to read it
238 queue: Deque<(T, usize), CAP>,
239 /// Every message has an id.
240 /// Don't worry, we won't run out.
241 /// If a million messages were published every second, then the ID's would run out in about 584942 years.
242 next_message_id: u64,
243 /// Collection of wakers for Subscribers that are waiting.
244 subscriber_wakers: MultiWakerRegistration<SUBS>,
245 /// Collection of wakers for Publishers that are waiting.
246 publisher_wakers: MultiWakerRegistration<PUBS>,
247 /// The amount of subscribers that are active
248 subscriber_count: usize,
249 /// The amount of publishers that are active
250 publisher_count: usize,
251}
252
253impl<T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> PubSubState<T, CAP, SUBS, PUBS> {
254 /// Create a new internal channel state
255 const fn new() -> Self {
256 Self {
257 queue: Deque::new(),
258 next_message_id: 0,
259 subscriber_wakers: MultiWakerRegistration::new(),
260 publisher_wakers: MultiWakerRegistration::new(),
261 subscriber_count: 0,
262 publisher_count: 0,
263 }
264 }
265
266 fn try_publish(&mut self, message: T) -> Result<(), T> {
267 if self.subscriber_count == 0 {
268 // We don't need to publish anything because there is no one to receive it
269 return Ok(());
270 }
271
272 if self.queue.is_full() {
273 return Err(message);
274 }
275 // We just did a check for this
276 self.queue.push_back((message, self.subscriber_count)).ok().unwrap();
277
278 self.next_message_id += 1;
279
280 // Wake all of the subscribers
281 self.subscriber_wakers.wake();
282
283 Ok(())
284 }
285
286 fn publish_immediate(&mut self, message: T) {
287 // Make space in the queue if required
288 if self.queue.is_full() {
289 self.queue.pop_front();
290 }
291
292 // This will succeed because we made sure there is space
293 self.try_publish(message).ok().unwrap();
294 }
295
296 fn get_message(&mut self, message_id: u64) -> Option<WaitResult<T>> {
297 let start_id = self.next_message_id - self.queue.len() as u64;
298
299 if message_id < start_id {
300 return Some(WaitResult::Lagged(start_id - message_id));
301 }
302
303 let current_message_index = (message_id - start_id) as usize;
304
305 if current_message_index >= self.queue.len() {
306 return None;
307 }
308
309 // We've checked that the index is valid
310 let queue_item = self.queue.iter_mut().nth(current_message_index).unwrap();
311
312 // We're reading this item, so decrement the counter
313 queue_item.1 -= 1;
314 let message = queue_item.0.clone();
315
316 if current_message_index == 0 && queue_item.1 == 0 {
317 self.queue.pop_front();
318 self.publisher_wakers.wake();
319 }
320
321 Some(WaitResult::Message(message))
322 }
323
324 fn register_subscriber_waker(&mut self, waker: &Waker) {
325 match self.subscriber_wakers.register(waker) {
326 Ok(()) => {}
327 Err(_) => {
328 // All waker slots were full. This can only happen when there was a subscriber that now has dropped.
329 // We need to throw it away. It's a bit inefficient, but we can wake everything.
330 // Any future that is still active will simply reregister.
331 // This won't happen a lot, so it's ok.
332 self.subscriber_wakers.wake();
333 self.subscriber_wakers.register(waker).unwrap();
334 }
335 }
336 }
337
338 fn register_publisher_waker(&mut self, waker: &Waker) {
339 match self.publisher_wakers.register(waker) {
340 Ok(()) => {}
341 Err(_) => {
342 // All waker slots were full. This can only happen when there was a publisher that now has dropped.
343 // We need to throw it away. It's a bit inefficient, but we can wake everything.
344 // Any future that is still active will simply reregister.
345 // This won't happen a lot, so it's ok.
346 self.publisher_wakers.wake();
347 self.publisher_wakers.register(waker).unwrap();
348 }
349 }
350 }
351
352 fn unregister_subscriber(&mut self, subscriber_next_message_id: u64) {
353 self.subscriber_count -= 1;
354
355 // All messages that haven't been read yet by this subscriber must have their counter decremented
356 let start_id = self.next_message_id - self.queue.len() as u64;
357 if subscriber_next_message_id >= start_id {
358 let current_message_index = (subscriber_next_message_id - start_id) as usize;
359 self.queue
360 .iter_mut()
361 .skip(current_message_index)
362 .for_each(|(_, counter)| *counter -= 1);
363 }
364 }
365
366 fn unregister_publisher(&mut self) {
367 self.publisher_count -= 1;
368 }
369}
370
371/// Error type for the [PubSubChannel]
372#[derive(Debug, PartialEq, Clone)]
373#[cfg_attr(feature = "defmt", derive(defmt::Format))]
374pub enum Error {
375 /// All subscriber slots are used. To add another subscriber, first another subscriber must be dropped or
376 /// the capacity of the channels must be increased.
377 MaximumSubscribersReached,
378 /// All publisher slots are used. To add another publisher, first another publisher must be dropped or
379 /// the capacity of the channels must be increased.
380 MaximumPublishersReached,
381}
382
383/// 'Middle level' behaviour of the pubsub channel.
384/// This trait is used so that Sub and Pub can be generic over the channel.
385pub trait PubSubBehavior<T> {
386 /// Try to get a message from the queue with the given message id.
387 ///
388 /// If the message is not yet present and a context is given, then its waker is registered in the subsriber wakers.
389 fn get_message_with_context(&self, next_message_id: &mut u64, cx: Option<&mut Context<'_>>) -> Poll<WaitResult<T>>;
390
391 /// Try to publish a message to the queue.
392 ///
393 /// If the queue is full and a context is given, then its waker is registered in the publisher wakers.
394 fn publish_with_context(&self, message: T, cx: Option<&mut Context<'_>>) -> Result<(), T>;
395
396 /// Publish a message immediately
397 fn publish_immediate(&self, message: T);
398
399 /// Let the channel know that a subscriber has dropped
400 fn unregister_subscriber(&self, subscriber_next_message_id: u64);
401
402 /// Let the channel know that a publisher has dropped
403 fn unregister_publisher(&self);
404}
405
406/// The result of the subscriber wait procedure
407#[derive(Debug, Clone, PartialEq)]
408#[cfg_attr(feature = "defmt", derive(defmt::Format))]
409pub enum WaitResult<T> {
410 /// The subscriber did not receive all messages and lagged by the given amount of messages.
411 /// (This is the amount of messages that were missed)
412 Lagged(u64),
413 /// A message was received
414 Message(T),
415}
416
417#[cfg(test)]
418mod tests {
419 use super::*;
420 use crate::blocking_mutex::raw::NoopRawMutex;
421
422 #[futures_test::test]
423 async fn dyn_pub_sub_works() {
424 let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
425
426 let mut sub0 = channel.dyn_subscriber().unwrap();
427 let mut sub1 = channel.dyn_subscriber().unwrap();
428 let pub0 = channel.dyn_publisher().unwrap();
429
430 pub0.publish(42).await;
431
432 assert_eq!(sub0.next_message().await, WaitResult::Message(42));
433 assert_eq!(sub1.next_message().await, WaitResult::Message(42));
434
435 assert_eq!(sub0.try_next_message(), None);
436 assert_eq!(sub1.try_next_message(), None);
437 }
438
439 #[futures_test::test]
440 async fn all_subscribers_receive() {
441 let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
442
443 let mut sub0 = channel.subscriber().unwrap();
444 let mut sub1 = channel.subscriber().unwrap();
445 let pub0 = channel.publisher().unwrap();
446
447 pub0.publish(42).await;
448
449 assert_eq!(sub0.next_message().await, WaitResult::Message(42));
450 assert_eq!(sub1.next_message().await, WaitResult::Message(42));
451
452 assert_eq!(sub0.try_next_message(), None);
453 assert_eq!(sub1.try_next_message(), None);
454 }
455
456 #[futures_test::test]
457 async fn lag_when_queue_full_on_immediate_publish() {
458 let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
459
460 let mut sub0 = channel.subscriber().unwrap();
461 let pub0 = channel.publisher().unwrap();
462
463 pub0.publish_immediate(42);
464 pub0.publish_immediate(43);
465 pub0.publish_immediate(44);
466 pub0.publish_immediate(45);
467 pub0.publish_immediate(46);
468 pub0.publish_immediate(47);
469
470 assert_eq!(sub0.try_next_message(), Some(WaitResult::Lagged(2)));
471 assert_eq!(sub0.next_message().await, WaitResult::Message(44));
472 assert_eq!(sub0.next_message().await, WaitResult::Message(45));
473 assert_eq!(sub0.next_message().await, WaitResult::Message(46));
474 assert_eq!(sub0.next_message().await, WaitResult::Message(47));
475 assert_eq!(sub0.try_next_message(), None);
476 }
477
478 #[test]
479 fn limited_subs_and_pubs() {
480 let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
481
482 let sub0 = channel.subscriber();
483 let sub1 = channel.subscriber();
484 let sub2 = channel.subscriber();
485 let sub3 = channel.subscriber();
486 let sub4 = channel.subscriber();
487
488 assert!(sub0.is_ok());
489 assert!(sub1.is_ok());
490 assert!(sub2.is_ok());
491 assert!(sub3.is_ok());
492 assert_eq!(sub4.err().unwrap(), Error::MaximumSubscribersReached);
493
494 drop(sub0);
495
496 let sub5 = channel.subscriber();
497 assert!(sub5.is_ok());
498
499 // publishers
500
501 let pub0 = channel.publisher();
502 let pub1 = channel.publisher();
503 let pub2 = channel.publisher();
504 let pub3 = channel.publisher();
505 let pub4 = channel.publisher();
506
507 assert!(pub0.is_ok());
508 assert!(pub1.is_ok());
509 assert!(pub2.is_ok());
510 assert!(pub3.is_ok());
511 assert_eq!(pub4.err().unwrap(), Error::MaximumPublishersReached);
512
513 drop(pub0);
514
515 let pub5 = channel.publisher();
516 assert!(pub5.is_ok());
517 }
518
519 #[test]
520 fn publisher_wait_on_full_queue() {
521 let channel = PubSubChannel::<NoopRawMutex, u32, 4, 4, 4>::new();
522
523 let pub0 = channel.publisher().unwrap();
524
525 // There are no subscribers, so the queue will never be full
526 assert_eq!(pub0.try_publish(0), Ok(()));
527 assert_eq!(pub0.try_publish(0), Ok(()));
528 assert_eq!(pub0.try_publish(0), Ok(()));
529 assert_eq!(pub0.try_publish(0), Ok(()));
530 assert_eq!(pub0.try_publish(0), Ok(()));
531
532 let sub0 = channel.subscriber().unwrap();
533
534 assert_eq!(pub0.try_publish(0), Ok(()));
535 assert_eq!(pub0.try_publish(0), Ok(()));
536 assert_eq!(pub0.try_publish(0), Ok(()));
537 assert_eq!(pub0.try_publish(0), Ok(()));
538 assert_eq!(pub0.try_publish(0), Err(0));
539
540 drop(sub0);
541 }
542}
diff --git a/embassy/src/channel/pubsub/publisher.rs b/embassy/src/channel/pubsub/publisher.rs
new file mode 100644
index 000000000..89a0b9247
--- /dev/null
+++ b/embassy/src/channel/pubsub/publisher.rs
@@ -0,0 +1,183 @@
1//! Implementation of anything directly publisher related
2
3use core::marker::PhantomData;
4use core::ops::{Deref, DerefMut};
5use core::pin::Pin;
6use core::task::{Context, Poll};
7
8use futures::Future;
9
10use super::{PubSubBehavior, PubSubChannel};
11use crate::blocking_mutex::raw::RawMutex;
12
13/// A publisher to a channel
14pub struct Pub<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> {
15 /// The channel we are a publisher for
16 channel: &'a PSB,
17 _phantom: PhantomData<T>,
18}
19
20impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Pub<'a, PSB, T> {
21 pub(super) fn new(channel: &'a PSB) -> Self {
22 Self {
23 channel,
24 _phantom: Default::default(),
25 }
26 }
27
28 /// Publish a message right now even when the queue is full.
29 /// This may cause a subscriber to miss an older message.
30 pub fn publish_immediate(&self, message: T) {
31 self.channel.publish_immediate(message)
32 }
33
34 /// Publish a message. But if the message queue is full, wait for all subscribers to have read the last message
35 pub fn publish<'s>(&'s self, message: T) -> PublisherWaitFuture<'s, 'a, PSB, T> {
36 PublisherWaitFuture {
37 message: Some(message),
38 publisher: self,
39 }
40 }
41
42 /// Publish a message if there is space in the message queue
43 pub fn try_publish(&self, message: T) -> Result<(), T> {
44 self.channel.publish_with_context(message, None)
45 }
46}
47
48impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Drop for Pub<'a, PSB, T> {
49 fn drop(&mut self) {
50 self.channel.unregister_publisher()
51 }
52}
53
54/// A publisher that holds a dynamic reference to the channel
55pub struct DynPublisher<'a, T: Clone>(pub(super) Pub<'a, dyn PubSubBehavior<T> + 'a, T>);
56
57impl<'a, T: Clone> Deref for DynPublisher<'a, T> {
58 type Target = Pub<'a, dyn PubSubBehavior<T> + 'a, T>;
59
60 fn deref(&self) -> &Self::Target {
61 &self.0
62 }
63}
64
65impl<'a, T: Clone> DerefMut for DynPublisher<'a, T> {
66 fn deref_mut(&mut self) -> &mut Self::Target {
67 &mut self.0
68 }
69}
70
71/// A publisher that holds a generic reference to the channel
72pub struct Publisher<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize>(
73 pub(super) Pub<'a, PubSubChannel<M, T, CAP, SUBS, PUBS>, T>,
74);
75
76impl<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> Deref
77 for Publisher<'a, M, T, CAP, SUBS, PUBS>
78{
79 type Target = Pub<'a, PubSubChannel<M, T, CAP, SUBS, PUBS>, T>;
80
81 fn deref(&self) -> &Self::Target {
82 &self.0
83 }
84}
85
86impl<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> DerefMut
87 for Publisher<'a, M, T, CAP, SUBS, PUBS>
88{
89 fn deref_mut(&mut self) -> &mut Self::Target {
90 &mut self.0
91 }
92}
93
94/// A publisher that can only use the `publish_immediate` function, but it doesn't have to be registered with the channel.
95/// (So an infinite amount is possible)
96pub struct ImmediatePub<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> {
97 /// The channel we are a publisher for
98 channel: &'a PSB,
99 _phantom: PhantomData<T>,
100}
101
102impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> ImmediatePub<'a, PSB, T> {
103 pub(super) fn new(channel: &'a PSB) -> Self {
104 Self {
105 channel,
106 _phantom: Default::default(),
107 }
108 }
109 /// Publish the message right now even when the queue is full.
110 /// This may cause a subscriber to miss an older message.
111 pub fn publish_immediate(&mut self, message: T) {
112 self.channel.publish_immediate(message)
113 }
114
115 /// Publish a message if there is space in the message queue
116 pub fn try_publish(&self, message: T) -> Result<(), T> {
117 self.channel.publish_with_context(message, None)
118 }
119}
120
121/// An immediate publisher that holds a dynamic reference to the channel
122pub struct DynImmediatePublisher<'a, T: Clone>(pub(super) ImmediatePub<'a, dyn PubSubBehavior<T> + 'a, T>);
123
124impl<'a, T: Clone> Deref for DynImmediatePublisher<'a, T> {
125 type Target = ImmediatePub<'a, dyn PubSubBehavior<T> + 'a, T>;
126
127 fn deref(&self) -> &Self::Target {
128 &self.0
129 }
130}
131
132impl<'a, T: Clone> DerefMut for DynImmediatePublisher<'a, T> {
133 fn deref_mut(&mut self) -> &mut Self::Target {
134 &mut self.0
135 }
136}
137
138/// An immediate publisher that holds a generic reference to the channel
139pub struct ImmediatePublisher<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize>(
140 pub(super) ImmediatePub<'a, PubSubChannel<M, T, CAP, SUBS, PUBS>, T>,
141);
142
143impl<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> Deref
144 for ImmediatePublisher<'a, M, T, CAP, SUBS, PUBS>
145{
146 type Target = ImmediatePub<'a, PubSubChannel<M, T, CAP, SUBS, PUBS>, T>;
147
148 fn deref(&self) -> &Self::Target {
149 &self.0
150 }
151}
152
153impl<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> DerefMut
154 for ImmediatePublisher<'a, M, T, CAP, SUBS, PUBS>
155{
156 fn deref_mut(&mut self) -> &mut Self::Target {
157 &mut self.0
158 }
159}
160
161/// Future for the publisher wait action
162pub struct PublisherWaitFuture<'s, 'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> {
163 /// The message we need to publish
164 message: Option<T>,
165 publisher: &'s Pub<'a, PSB, T>,
166}
167
168impl<'s, 'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Future for PublisherWaitFuture<'s, 'a, PSB, T> {
169 type Output = ();
170
171 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
172 let message = self.message.take().unwrap();
173 match self.publisher.channel.publish_with_context(message, Some(cx)) {
174 Ok(()) => Poll::Ready(()),
175 Err(message) => {
176 self.message = Some(message);
177 Poll::Pending
178 }
179 }
180 }
181}
182
183impl<'s, 'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Unpin for PublisherWaitFuture<'s, 'a, PSB, T> {}
diff --git a/embassy/src/channel/pubsub/subscriber.rs b/embassy/src/channel/pubsub/subscriber.rs
new file mode 100644
index 000000000..23c4938d9
--- /dev/null
+++ b/embassy/src/channel/pubsub/subscriber.rs
@@ -0,0 +1,153 @@
1//! Implementation of anything directly subscriber related
2
3use core::marker::PhantomData;
4use core::ops::{Deref, DerefMut};
5use core::pin::Pin;
6use core::task::{Context, Poll};
7
8use futures::Future;
9
10use super::{PubSubBehavior, PubSubChannel, WaitResult};
11use crate::blocking_mutex::raw::RawMutex;
12
13/// A subscriber to a channel
14pub struct Sub<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> {
15 /// The message id of the next message we are yet to receive
16 next_message_id: u64,
17 /// The channel we are a subscriber to
18 channel: &'a PSB,
19 _phantom: PhantomData<T>,
20}
21
22impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Sub<'a, PSB, T> {
23 pub(super) fn new(next_message_id: u64, channel: &'a PSB) -> Self {
24 Self {
25 next_message_id,
26 channel,
27 _phantom: Default::default(),
28 }
29 }
30
31 /// Wait for a published message
32 pub fn next_message<'s>(&'s mut self) -> SubscriberWaitFuture<'s, 'a, PSB, T> {
33 SubscriberWaitFuture { subscriber: self }
34 }
35
36 /// Wait for a published message (ignoring lag results)
37 pub async fn next_message_pure(&mut self) -> T {
38 loop {
39 match self.next_message().await {
40 WaitResult::Lagged(_) => continue,
41 WaitResult::Message(message) => break message,
42 }
43 }
44 }
45
46 /// Try to see if there's a published message we haven't received yet.
47 ///
48 /// This function does not peek. The message is received if there is one.
49 pub fn try_next_message(&mut self) -> Option<WaitResult<T>> {
50 match self.channel.get_message_with_context(&mut self.next_message_id, None) {
51 Poll::Ready(result) => Some(result),
52 Poll::Pending => None,
53 }
54 }
55
56 /// Try to see if there's a published message we haven't received yet (ignoring lag results).
57 ///
58 /// This function does not peek. The message is received if there is one.
59 pub fn try_next_message_pure(&mut self) -> Option<T> {
60 loop {
61 match self.try_next_message() {
62 Some(WaitResult::Lagged(_)) => continue,
63 Some(WaitResult::Message(message)) => break Some(message),
64 None => break None,
65 }
66 }
67 }
68}
69
70impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Drop for Sub<'a, PSB, T> {
71 fn drop(&mut self) {
72 self.channel.unregister_subscriber(self.next_message_id)
73 }
74}
75
76impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Unpin for Sub<'a, PSB, T> {}
77
78/// Warning: The stream implementation ignores lag results and returns all messages.
79/// This might miss some messages without you knowing it.
80impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> futures::Stream for Sub<'a, PSB, T> {
81 type Item = T;
82
83 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
84 match self
85 .channel
86 .get_message_with_context(&mut self.next_message_id, Some(cx))
87 {
88 Poll::Ready(WaitResult::Message(message)) => Poll::Ready(Some(message)),
89 Poll::Ready(WaitResult::Lagged(_)) => {
90 cx.waker().wake_by_ref();
91 Poll::Pending
92 }
93 Poll::Pending => Poll::Pending,
94 }
95 }
96}
97
98/// A subscriber that holds a dynamic reference to the channel
99pub struct DynSubscriber<'a, T: Clone>(pub(super) Sub<'a, dyn PubSubBehavior<T> + 'a, T>);
100
101impl<'a, T: Clone> Deref for DynSubscriber<'a, T> {
102 type Target = Sub<'a, dyn PubSubBehavior<T> + 'a, T>;
103
104 fn deref(&self) -> &Self::Target {
105 &self.0
106 }
107}
108
109impl<'a, T: Clone> DerefMut for DynSubscriber<'a, T> {
110 fn deref_mut(&mut self) -> &mut Self::Target {
111 &mut self.0
112 }
113}
114
115/// A subscriber that holds a generic reference to the channel
116pub struct Subscriber<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize>(
117 pub(super) Sub<'a, PubSubChannel<M, T, CAP, SUBS, PUBS>, T>,
118);
119
120impl<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> Deref
121 for Subscriber<'a, M, T, CAP, SUBS, PUBS>
122{
123 type Target = Sub<'a, PubSubChannel<M, T, CAP, SUBS, PUBS>, T>;
124
125 fn deref(&self) -> &Self::Target {
126 &self.0
127 }
128}
129
130impl<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> DerefMut
131 for Subscriber<'a, M, T, CAP, SUBS, PUBS>
132{
133 fn deref_mut(&mut self) -> &mut Self::Target {
134 &mut self.0
135 }
136}
137
138/// Future for the subscriber wait action
139pub struct SubscriberWaitFuture<'s, 'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> {
140 subscriber: &'s mut Sub<'a, PSB, T>,
141}
142
143impl<'s, 'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Future for SubscriberWaitFuture<'s, 'a, PSB, T> {
144 type Output = WaitResult<T>;
145
146 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
147 self.subscriber
148 .channel
149 .get_message_with_context(&mut self.subscriber.next_message_id, Some(cx))
150 }
151}
152
153impl<'s, 'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Unpin for SubscriberWaitFuture<'s, 'a, PSB, T> {}
diff --git a/embassy/src/waitqueue/mod.rs b/embassy/src/waitqueue/mod.rs
index a2bafad99..5c4e1bc3b 100644
--- a/embassy/src/waitqueue/mod.rs
+++ b/embassy/src/waitqueue/mod.rs
@@ -3,3 +3,6 @@
3#[cfg_attr(feature = "executor-agnostic", path = "waker_agnostic.rs")] 3#[cfg_attr(feature = "executor-agnostic", path = "waker_agnostic.rs")]
4mod waker; 4mod waker;
5pub use waker::*; 5pub use waker::*;
6
7mod multi_waker;
8pub use multi_waker::*;
diff --git a/embassy/src/waitqueue/multi_waker.rs b/embassy/src/waitqueue/multi_waker.rs
new file mode 100644
index 000000000..325d2cb3a
--- /dev/null
+++ b/embassy/src/waitqueue/multi_waker.rs
@@ -0,0 +1,33 @@
1use core::task::Waker;
2
3use super::WakerRegistration;
4
5/// Utility struct to register and wake multiple wakers.
6pub struct MultiWakerRegistration<const N: usize> {
7 wakers: [WakerRegistration; N],
8}
9
10impl<const N: usize> MultiWakerRegistration<N> {
11 /// Create a new empty instance
12 pub const fn new() -> Self {
13 const WAKER: WakerRegistration = WakerRegistration::new();
14 Self { wakers: [WAKER; N] }
15 }
16
17 /// Register a waker. If the buffer is full the function returns it in the error
18 pub fn register<'a>(&mut self, w: &'a Waker) -> Result<(), &'a Waker> {
19 if let Some(waker_slot) = self.wakers.iter_mut().find(|waker_slot| !waker_slot.occupied()) {
20 waker_slot.register(w);
21 Ok(())
22 } else {
23 Err(w)
24 }
25 }
26
27 /// Wake all registered wakers. This clears the buffer
28 pub fn wake(&mut self) {
29 for waker_slot in self.wakers.iter_mut() {
30 waker_slot.wake()
31 }
32 }
33}
diff --git a/embassy/src/waitqueue/waker.rs b/embassy/src/waitqueue/waker.rs
index da907300a..a90154cce 100644
--- a/embassy/src/waitqueue/waker.rs
+++ b/embassy/src/waitqueue/waker.rs
@@ -50,6 +50,11 @@ impl WakerRegistration {
50 unsafe { wake_task(w) } 50 unsafe { wake_task(w) }
51 } 51 }
52 } 52 }
53
54 /// Returns true if a waker is currently registered
55 pub fn occupied(&self) -> bool {
56 self.waker.is_some()
57 }
53} 58}
54 59
55// SAFETY: `WakerRegistration` effectively contains an `Option<Waker>`, 60// SAFETY: `WakerRegistration` effectively contains an `Option<Waker>`,
diff --git a/embassy/src/waitqueue/waker_agnostic.rs b/embassy/src/waitqueue/waker_agnostic.rs
index 89430aa4c..62e3adb79 100644
--- a/embassy/src/waitqueue/waker_agnostic.rs
+++ b/embassy/src/waitqueue/waker_agnostic.rs
@@ -47,6 +47,11 @@ impl WakerRegistration {
47 w.wake() 47 w.wake()
48 } 48 }
49 } 49 }
50
51 /// Returns true if a waker is currently registered
52 pub fn occupied(&self) -> bool {
53 self.waker.is_some()
54 }
50} 55}
51 56
52/// Utility struct to register and wake a waker. 57/// Utility struct to register and wake a waker.
diff --git a/examples/nrf/src/bin/pubsub.rs b/examples/nrf/src/bin/pubsub.rs
new file mode 100644
index 000000000..2c3a355c2
--- /dev/null
+++ b/examples/nrf/src/bin/pubsub.rs
@@ -0,0 +1,106 @@
1#![no_std]
2#![no_main]
3#![feature(type_alias_impl_trait)]
4
5use defmt::unwrap;
6use embassy::blocking_mutex::raw::ThreadModeRawMutex;
7use embassy::channel::pubsub::{DynSubscriber, PubSubChannel, Subscriber};
8use embassy::executor::Spawner;
9use embassy::time::{Duration, Timer};
10use {defmt_rtt as _, panic_probe as _};
11
12/// Create the message bus. It has a queue of 4, supports 3 subscribers and 1 publisher
13static MESSAGE_BUS: PubSubChannel<ThreadModeRawMutex, Message, 4, 3, 1> = PubSubChannel::new();
14
15#[derive(Clone, defmt::Format)]
16enum Message {
17 A,
18 B,
19 C,
20}
21
22#[embassy::main]
23async fn main(spawner: Spawner, _p: embassy_nrf::Peripherals) {
24 defmt::info!("Hello World!");
25
26 // It's good to set up the subscribers before publishing anything.
27 // A subscriber will only yield messages that have been published after its creation.
28
29 spawner.must_spawn(fast_logger(unwrap!(MESSAGE_BUS.subscriber())));
30 spawner.must_spawn(slow_logger(unwrap!(MESSAGE_BUS.dyn_subscriber())));
31 spawner.must_spawn(slow_logger_pure(unwrap!(MESSAGE_BUS.dyn_subscriber())));
32
33 // Get a publisher
34 let message_publisher = unwrap!(MESSAGE_BUS.publisher());
35 // We can't get more (normal) publishers
36 // We can have an infinite amount of immediate publishers. They can't await a publish, only do an immediate publish
37 defmt::assert!(MESSAGE_BUS.publisher().is_err());
38
39 let mut index = 0;
40 loop {
41 Timer::after(Duration::from_millis(500)).await;
42
43 let message = match index % 3 {
44 0 => Message::A,
45 1 => Message::B,
46 2..=u32::MAX => Message::C,
47 };
48
49 // We publish immediately and don't await anything.
50 // If the queue is full, it will cause the oldest message to not be received by some/all subscribers
51 message_publisher.publish_immediate(message);
52
53 // Try to comment out the last one and uncomment this line below.
54 // The behaviour will change:
55 // - The subscribers won't miss any messages any more
56 // - Trying to publish now has some wait time when the queue is full
57
58 // message_publisher.publish(message).await;
59
60 index += 1;
61 }
62}
63
64/// A logger task that just awaits the messages it receives
65///
66/// This takes the generic `Subscriber`. This is most performant, but requires you to write down all of the generics
67#[embassy::task]
68async fn fast_logger(mut messages: Subscriber<'static, ThreadModeRawMutex, Message, 4, 3, 1>) {
69 loop {
70 let message = messages.next_message().await;
71 defmt::info!("Received message at fast logger: {:?}", message);
72 }
73}
74
75/// A logger task that awaits the messages, but also does some other work.
76/// Because of this, depeding on how the messages were published, the subscriber might miss some messages
77///
78/// This takes the dynamic `DynSubscriber`. This is not as performant as the generic version, but let's you ignore some of the generics
79#[embassy::task]
80async fn slow_logger(mut messages: DynSubscriber<'static, Message>) {
81 loop {
82 // Do some work
83 Timer::after(Duration::from_millis(2000)).await;
84
85 // If the publisher has used the `publish_immediate` function, then we may receive a lag message here
86 let message = messages.next_message().await;
87 defmt::info!("Received message at slow logger: {:?}", message);
88
89 // If the previous one was a lag message, then we should receive the next message here immediately
90 let message = messages.next_message().await;
91 defmt::info!("Received message at slow logger: {:?}", message);
92 }
93}
94
95/// Same as `slow_logger` but it ignores lag results
96#[embassy::task]
97async fn slow_logger_pure(mut messages: DynSubscriber<'static, Message>) {
98 loop {
99 // Do some work
100 Timer::after(Duration::from_millis(2000)).await;
101
102 // Instead of receiving lags here, we just ignore that and read the next message
103 let message = messages.next_message_pure().await;
104 defmt::info!("Received message at slow logger pure: {:?}", message);
105 }
106}