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/// # Example
14///
15/// ```
16/// use tokio::sync::broadcast;
17/// use tokio_stream::wrappers::BroadcastStream;
18/// use tokio_stream::StreamExt;
19///
20/// # #[tokio::main(flavor = "current_thread")]
21/// # async fn main() -> Result<(), tokio::sync::broadcast::error::SendError<u8>> {
22/// let (tx, rx) = broadcast::channel(16);
23/// tx.send(10)?;
24/// tx.send(20)?;
25/// # // prevent the doc test from hanging
26/// drop(tx);
27///
28/// let mut stream = BroadcastStream::new(rx);
29/// assert_eq!(stream.next().await, Some(Ok(10)));
30/// assert_eq!(stream.next().await, Some(Ok(20)));
31/// assert_eq!(stream.next().await, None);
32/// # Ok(())
33/// # }
34/// ```
35///
36/// [`tokio::sync::broadcast::Receiver`]: struct@tokio::sync::broadcast::Receiver
37/// [`Stream`]: trait@futures_core::Stream
38#[cfg_attr(docsrs, doc(cfg(feature = "sync")))]
39pub struct BroadcastStream<T> {
40    inner: ReusableBoxFuture<'static, (Result<T, RecvError>, Receiver<T>)>,
41}
42
43/// An error returned from the inner stream of a [`BroadcastStream`].
44#[derive(Debug, PartialEq, Eq, Clone)]
45pub enum BroadcastStreamRecvError {
46    /// The receiver lagged too far behind. Attempting to receive again will
47    /// return the oldest message still retained by the channel.
48    ///
49    /// Includes the number of skipped messages.
50    Lagged(u64),
51}
52
53impl fmt::Display for BroadcastStreamRecvError {
54    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55        match self {
56            BroadcastStreamRecvError::Lagged(amt) => write!(f, "channel lagged by {}", amt),
57        }
58    }
59}
60
61impl std::error::Error for BroadcastStreamRecvError {}
62
63async fn make_future<T: Clone>(mut rx: Receiver<T>) -> (Result<T, RecvError>, Receiver<T>) {
64    let result = rx.recv().await;
65    (result, rx)
66}
67
68impl<T: 'static + Clone + Send> BroadcastStream<T> {
69    /// Create a new `BroadcastStream`.
70    pub fn new(rx: Receiver<T>) -> Self {
71        Self {
72            inner: ReusableBoxFuture::new(make_future(rx)),
73        }
74    }
75}
76
77impl<T: 'static + Clone + Send> Stream for BroadcastStream<T> {
78    type Item = Result<T, BroadcastStreamRecvError>;
79    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
80        let (result, rx) = ready!(self.inner.poll(cx));
81        self.inner.set(make_future(rx));
82        match result {
83            Ok(item) => Poll::Ready(Some(Ok(item))),
84            Err(RecvError::Closed) => Poll::Ready(None),
85            Err(RecvError::Lagged(n)) => {
86                Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(n))))
87            }
88        }
89    }
90}
91
92impl<T> fmt::Debug for BroadcastStream<T> {
93    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
94        f.debug_struct("BroadcastStream").finish()
95    }
96}
97
98impl<T: 'static + Clone + Send> From<Receiver<T>> for BroadcastStream<T> {
99    fn from(recv: Receiver<T>) -> Self {
100        Self::new(recv)
101    }
102}
tokio_stream/wrappers/broadcast.rs

tokio_stream/wrappers/
broadcast.rs
