tokio_stream/wrappers/
broadcast.rs

1use std::pin::Pin;
2use tokio::sync::broadcast::error::RecvError;
3use tokio::sync::broadcast::Receiver;
4
5use futures_core::Stream;
6use tokio_util::sync::ReusableBoxFuture;
7
8use std::fmt;
9use 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@futures_core::Stream
15#[cfg_attr(docsrs, doc(cfg(feature = "sync")))]
16pub 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)]
22pub 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
30impl fmt::Display for BroadcastStreamRecvError {
31    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
38impl std::error::Error for BroadcastStreamRecvError {}
39
40async 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
45impl<T: 'static + Clone + Send> BroadcastStream<T> {
46    /// Create a new `BroadcastStream`.
47    pub fn new(rx: Receiver<T>) -> Self {
48        Self {
49            inner: ReusableBoxFuture::new(make_future(rx)),
50        }
51    }
52}
53
54impl<T: 'static + Clone + Send> Stream for BroadcastStream<T> {
55    type Item = Result<T, BroadcastStreamRecvError>;
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
69impl<T> fmt::Debug for BroadcastStream<T> {
70    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71        f.debug_struct("BroadcastStream").finish()
72    }
73}
74
75impl<T: 'static + Clone + Send> From<Receiver<T>> for BroadcastStream<T> {
76    fn from(recv: Receiver<T>) -> Self {
77        Self::new(recv)
78    }
79}
tokio_stream/wrappers/broadcast.rs

tokio_stream/wrappers/
broadcast.rs
