remove dependency of async_stream

tokio-rs · Jan 31, 2021 · 34eab73 · 34eab73
1 parent ef12286
commit 34eab73
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 36 deletions.
diff --git a/tokio-stream/Cargo.toml b/tokio-stream/Cargo.toml
@@ -28,9 +28,9 @@ fs = ["tokio/fs"]
 
 [dependencies]
 futures-core = { version = "0.3.0" }
-async-stream = "0.3"
 pin-project-lite = "0.2.0"
 tokio = { version = "1.0", features = ["sync"] }
+tokio-util = { path = "../tokio-util" }
 
 [dev-dependencies]
 tokio = { version = "1.0", features = ["full", "test-util"] }

diff --git a/tokio-stream/src/wrappers/broadcast.rs b/tokio-stream/src/wrappers/broadcast.rs
@@ -1,9 +1,10 @@
-use crate::Stream;
-use async_stream::stream;
 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};
 
@@ -12,7 +13,7 @@ use std::task::{Context, Poll};
 /// [`tokio::sync::broadcast::Receiver`]: struct@tokio::sync::broadcast::Receiver
 /// [`Stream`]: trait@crate::Stream
 pub struct BroadcastStream<T> {
-    inner: Pin<Box<dyn Stream<Item = Result<T, BroadcastStreamRecvError>> + Send + Sync>>,
+    inner: ReusableBoxFuture<Result<(T, Receiver<T>), RecvError>>,
 }
 
 /// An error returned from the inner stream of a [`BroadcastStream`].
@@ -25,31 +26,35 @@ pub enum BroadcastStreamRecvError {
     Lagged(u64),
 }
 
+async fn make_future<T: Clone + Send + Sync>(
+    mut rx: Receiver<T>,
+) -> Result<(T, Receiver<T>), RecvError> {
+    let item = rx.recv().await?;
+    Ok((item, rx))
+}
+
 impl<T: Clone + Unpin + 'static + Send + Sync> BroadcastStream<T> {
     /// Create a new `BroadcastStream`.
-    pub fn new(mut rx: Receiver<T>) -> Self {
-        let stream = stream! {
-            loop {
-                match rx.recv().await {
-                    Ok(item) => yield Ok(item),
-                    Err(err) => match err {
-                        RecvError::Closed => break,
-                        RecvError::Lagged(n) => yield Err(BroadcastStreamRecvError::Lagged(n)),
-                    },
-                }
-            }
-        };
+    pub fn new(rx: Receiver<T>) -> Self {
         Self {
-            inner: Box::pin(stream),
+            inner: ReusableBoxFuture::new(make_future(rx)),
         }
     }
 }
 
-impl<T: Clone> Stream for BroadcastStream<T> {
+impl<T: Clone + 'static + Send + Sync> Stream for BroadcastStream<T> {
     type Item = Result<T, BroadcastStreamRecvError>;
-
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        Pin::new(&mut self.inner).poll_next(cx)
+        match ready!(self.inner.poll(cx)) {
+            Ok((item, rx)) => {
+                self.inner.set(make_future(rx));
+                Poll::Ready(Some(Ok(item)))
+            }
+            Err(err) => match err {
+                RecvError::Closed => Poll::Ready(None),
+                RecvError::Lagged(n) => Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(n)))),
+            },
+        }
     }
 }
 

diff --git a/tokio-stream/src/wrappers/watch.rs b/tokio-stream/src/wrappers/watch.rs
@@ -1,41 +1,49 @@
-use crate::Stream;
-use async_stream::stream;
 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: Pin<Box<dyn Stream<Item = 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> WatchStream<T> {
+impl<T: 'static + Clone + Unpin + Send + Sync> WatchStream<T> {
     /// Create a new `WatchStream`.
-    pub fn new(mut rx: Receiver<T>) -> Self {
-        let stream = stream! {
-            loop {
-                match rx.changed().await {
-                    Ok(_) => yield (*rx.borrow()).clone(),
-                    Err(_) => break,
-                }
-            }
-        };
+    pub fn new(rx: Receiver<T>) -> Self {
         Self {
-            inner: Box::pin(stream),
+            inner: ReusableBoxFuture::new(make_future(rx)),
         }
     }
 }
 
-impl<T> Stream for WatchStream<T> {
+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>> {
-        Pin::new(&mut self.inner).poll_next(cx)
+        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),
+        }
     }
 }