1 use std::pin::Pin;
2 use tokio::sync::broadcast::error::RecvError;
3 use tokio::sync::broadcast::Receiver;
4
5 use futures_core::Stream;
6 use tokio_util::sync::ReusableBoxFuture;
7
8 use std::fmt;
9 use std::task::{ready, Context, Poll};
10
11 /// A wrapper around [`tokio::sync::broadcast::Receiver`] that implements [`Stream`].
12 ///
13 /// [`tokio::sync::broadcast::Receiver`]: struct@tokio::sync::broadcast::Receiver
14 /// [`Stream`]: trait@crate::Stream
15 #[cfg_attr(docsrs, doc(cfg(feature = "sync")))]
16 pub struct BroadcastStream<T> {
17 inner: ReusableBoxFuture<'static, (Result<T, RecvError>, Receiver<T>)>,
18 }
19
20 /// An error returned from the inner stream of a [`BroadcastStream`].
21 #[derive(Debug, PartialEq, Eq, Clone)]
22 pub enum BroadcastStreamRecvError {
23 /// The receiver lagged too far behind. Attempting to receive again will
24 /// return the oldest message still retained by the channel.
25 ///
26 /// Includes the number of skipped messages.
27 Lagged(u64),
28 }
29
30 impl fmt::Display for BroadcastStreamRecvError {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result31 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32 match self {
33 BroadcastStreamRecvError::Lagged(amt) => write!(f, "channel lagged by {}", amt),
34 }
35 }
36 }
37
38 impl std::error::Error for BroadcastStreamRecvError {}
39
make_future<T: Clone>(mut rx: Receiver<T>) -> (Result<T, RecvError>, Receiver<T>)40 async fn make_future<T: Clone>(mut rx: Receiver<T>) -> (Result<T, RecvError>, Receiver<T>) {
41 let result = rx.recv().await;
42 (result, rx)
43 }
44
45 impl<T: 'static + Clone + Send> BroadcastStream<T> {
46 /// Create a new `BroadcastStream`.
new(rx: Receiver<T>) -> Self47 pub fn new(rx: Receiver<T>) -> Self {
48 Self {
49 inner: ReusableBoxFuture::new(make_future(rx)),
50 }
51 }
52 }
53
54 impl<T: 'static + Clone + Send> Stream for BroadcastStream<T> {
55 type Item = Result<T, BroadcastStreamRecvError>;
poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>56 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
57 let (result, rx) = ready!(self.inner.poll(cx));
58 self.inner.set(make_future(rx));
59 match result {
60 Ok(item) => Poll::Ready(Some(Ok(item))),
61 Err(RecvError::Closed) => Poll::Ready(None),
62 Err(RecvError::Lagged(n)) => {
63 Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(n))))
64 }
65 }
66 }
67 }
68
69 impl<T> fmt::Debug for BroadcastStream<T> {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result70 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71 f.debug_struct("BroadcastStream").finish()
72 }
73 }
74
75 impl<T: 'static + Clone + Send> From<Receiver<T>> for BroadcastStream<T> {
from(recv: Receiver<T>) -> Self76 fn from(recv: Receiver<T>) -> Self {
77 Self::new(recv)
78 }
79 }
80