tokio-stream: add wrapper for broadcast and watch

examples/Cargo.toml

@@ -8,7 +8,7 @@ edition = "2018"
 # [dependencies] instead.
 [dev-dependencies]
 tokio = { version = "1.0.0", features = ["full", "tracing"] }
-tokio-util = { version = "0.6.1", features = ["full"] }
+tokio-util = { version = "0.6.3", features = ["full"] }
 tokio-stream = { version = "0.1" }
 
 async-stream = "0.3"

tokio-stream/Cargo.toml

@@ -30,11 +30,11 @@ fs = ["tokio/fs"]
 futures-core = { version = "0.3.0" }
 pin-project-lite = "0.2.0"
 tokio = { version = "1.0", features = ["sync"] }
+tokio-util = { version = "0.6.3" }
 
 [dev-dependencies]
 tokio = { version = "1.0", features = ["full", "test-util"] }
 tokio-test = { path = "../tokio-test" }
-async-stream = "0.3"
 futures = { version = "0.3", default-features = false }
 
 proptest = "0.10.0"

tokio-stream/src/wrappers.rs

@@ -6,6 +6,13 @@ pub use mpsc_bounded::ReceiverStream;
 mod mpsc_unbounded;
 pub use mpsc_unbounded::UnboundedReceiverStream;
 
+mod broadcast;
+pub use broadcast::BroadcastStream;
+pub use broadcast::BroadcastStreamRecvError;
+
+mod watch;
+pub use watch::WatchStream;
+
 cfg_time! {
     mod interval;
     pub use interval::IntervalStream;

tokio-stream/src/wrappers/broadcast.rs

@@ -0,0 +1,77 @@
+use std::pin::Pin;
+use tokio::sync::broadcast::error::RecvError;
+use tokio::sync::broadcast::Receiver;
+
+use futures_core::Stream;
+use tokio_util::sync::ReusableBoxFuture;
+
+use std::fmt;
+use std::task::{Context, Poll};
+
+/// A wrapper around [`tokio::sync::broadcast::Receiver`] that implements [`Stream`].
+///
+/// [`tokio::sync::broadcast::Receiver`]: struct@tokio::sync::broadcast::Receiver
+/// [`Stream`]: trait@crate::Stream
+pub struct BroadcastStream<T> {
+    inner: ReusableBoxFuture<Result<(T, Receiver<T>), WrappedRecvError<T>>>,
+}
+
+/// An error returned from the inner stream of a [`BroadcastStream`].
+#[derive(Debug, PartialEq)]
+pub enum BroadcastStreamRecvError {
+    /// The receiver lagged too far behind. Attempting to receive again will
+    /// return the oldest message still retained by the channel.
+    ///
+    /// Includes the number of skipped messages.
+    Lagged(u64),
+}
+
+#[derive(Debug)]
+enum WrappedRecvError<T> {
+    Lagged(u64, Receiver<T>),
+    Closed,
+}
+
+async fn make_future<T: Clone>(
+    mut rx: Receiver<T>,
+) -> Result<(T, Receiver<T>), WrappedRecvError<T>> {
+    match rx.recv().await {
+        Ok(item) => Ok((item, rx)),
+        Err(RecvError::Lagged(n)) => Err(WrappedRecvError::Lagged(n, rx)),
+        Err(RecvError::Closed) => Err(WrappedRecvError::Closed),
+    }
+}
+
+impl<T: 'static + Clone + Send> BroadcastStream<T> {
+    /// Create a new `BroadcastStream`.
+    pub fn new(rx: Receiver<T>) -> Self {
+        Self {
+            inner: ReusableBoxFuture::new(make_future(rx)),
+        }
+    }
+}
+
+impl<T: 'static + Clone + Send> Stream for BroadcastStream<T> {
+    type Item = Result<T, BroadcastStreamRecvError>;
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match ready!(self.inner.poll(cx)) {
+            Ok((item, rx)) => {
+                self.inner.set(make_future(rx));
+                Poll::Ready(Some(Ok(item)))
+            }
+            Err(err) => match err {
+                WrappedRecvError::Closed => Poll::Ready(None),
+                WrappedRecvError::Lagged(n, rx) => {
+                    self.inner.set(make_future(rx));
+                    Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(n))))
+                }
+            },
+        }
+    }
+}
+
+impl<T: Clone> fmt::Debug for BroadcastStream<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("BroadcastStream").finish()
+    }
+}

tokio-stream/src/wrappers/watch.rs

@@ -0,0 +1,56 @@
+use std::pin::Pin;
+use tokio::sync::watch::Receiver;
+
+use futures_core::Stream;
+use tokio_util::sync::ReusableBoxFuture;
+
+use std::fmt;
+use std::task::{Context, Poll};
+use tokio::sync::watch::error::RecvError;
+
+/// A wrapper around [`tokio::sync::watch::Receiver`] that implements [`Stream`].
+///
+/// [`tokio::sync::watch::Receiver`]: struct@tokio::sync::watch::Receiver
+/// [`Stream`]: trait@crate::Stream
+pub struct WatchStream<T> {
+    inner: ReusableBoxFuture<Result<((), Receiver<T>), RecvError>>,
+}
+
+async fn make_future<T: Clone + Send + Sync>(
+    mut rx: Receiver<T>,
+) -> Result<((), Receiver<T>), RecvError> {
+    let signal = rx.changed().await?;
+    Ok((signal, rx))
+}
+
+impl<T: 'static + Clone + Unpin + Send + Sync> WatchStream<T> {
+    /// Create a new `WatchStream`.
+    pub fn new(rx: Receiver<T>) -> Self {
+        Self {
+            inner: ReusableBoxFuture::new(make_future(rx)),
+        }
+    }
+}
+
+impl<T: Clone + 'static + Send + Sync> Stream for WatchStream<T> {
+    type Item = T;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match ready!(self.inner.poll(cx)) {
+            Ok((_, rx)) => {
+                let received = (*rx.borrow()).clone();
+                self.inner.set(make_future(rx));
+                Poll::Ready(Some(received))
+            }
+            Err(_) => Poll::Ready(None),
+        }
+    }
+}
+
+impl<T> Unpin for WatchStream<T> {}
+
+impl<T> fmt::Debug for WatchStream<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("WatchStream").finish()
+    }
+}