tokio-stream: add wrapper for broadcast and watch (#3384)

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,12 @@ 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"
+tokio-test = { path = "../tokio-test" }
 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,62 @@
+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, RecvError>, Receiver<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),
+}
+
+async fn make_future<T: Clone>(mut rx: Receiver<T>) -> (Result<T, RecvError>, Receiver<T>) {
+    let result = rx.recv().await;
+    (result, rx)
+}
+
+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>> {
+        let (result, rx) = ready!(self.inner.poll(cx));
+        self.inner.set(make_future(rx));
+        match result {
+            Ok(item) => Poll::Ready(Some(Ok(item))),
+            Err(RecvError::Closed) => Poll::Ready(None),
+            Err(RecvError::Lagged(n)) => {
+                Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(n))))
+            }
+        }
+    }
+}
+
+impl<T> 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,60 @@
+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<(), RecvError>, Receiver<T>)>,
+}
+
+async fn make_future<T: Clone + Send + Sync>(
+    mut rx: Receiver<T>,
+) -> (Result<(), RecvError>, Receiver<T>) {
+    let result = rx.changed().await;
+    (result, 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>> {
+        let (result, rx) = ready!(self.inner.poll(cx));
+        match result {
+            Ok(_) => {
+                let received = (*rx.borrow()).clone();
+                self.inner.set(make_future(rx));
+                Poll::Ready(Some(received))
+            }
+            Err(_) => {
+                self.inner.set(make_future(rx));
+                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()
+    }
+}