diff --git a/tokio-stream/Cargo.toml b/tokio-stream/Cargo.toml index 77f20761a..72f0a7ae4 100644 --- a/tokio-stream/Cargo.toml +++ b/tokio-stream/Cargo.toml @@ -24,6 +24,7 @@ full = [ "net", "io-util", "fs", + "rt", "sync", "signal" ] @@ -32,6 +33,7 @@ time = ["tokio/time"] net = ["tokio/net"] io-util = ["tokio/io-util"] fs = ["tokio/fs"] +rt = ["tokio/rt"] sync = ["tokio/sync", "tokio-util"] signal = ["tokio/signal"] diff --git a/tokio-stream/src/macros.rs b/tokio-stream/src/macros.rs index 5aa797b4f..8d96da8f0 100644 --- a/tokio-stream/src/macros.rs +++ b/tokio-stream/src/macros.rs @@ -57,3 +57,13 @@ macro_rules! cfg_signal { )* } } + +macro_rules! cfg_rt { + ($($item:item)*) => { + $( + #[cfg(feature = "rt")] + #[cfg_attr(docsrs, doc(cfg(feature = "rt")))] + $item + )* + } +} diff --git a/tokio-stream/src/wrappers.rs b/tokio-stream/src/wrappers.rs index c79b6c768..0604e7d2f 100644 --- a/tokio-stream/src/wrappers.rs +++ b/tokio-stream/src/wrappers.rs @@ -13,6 +13,11 @@ pub use mpsc_bounded::ReceiverStream; mod mpsc_unbounded; pub use mpsc_unbounded::UnboundedReceiverStream; +cfg_rt! { + mod task; + pub use task::JoinSetStream; +} + cfg_sync! { mod broadcast; pub use broadcast::BroadcastStream; diff --git a/tokio-stream/src/wrappers/task.rs b/tokio-stream/src/wrappers/task.rs new file mode 100644 index 000000000..d844a38a5 --- /dev/null +++ b/tokio-stream/src/wrappers/task.rs @@ -0,0 +1,78 @@ +use crate::Stream; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::task::{JoinError, JoinSet}; + +/// A wrapper around [`tokio::task::JoinSet`] that implements [`Stream`]. +/// +/// # Example +/// +/// ``` +/// use tokio::task::JoinSet; +/// use tokio_stream::wrappers::JoinSetStream; +/// use tokio_stream::StreamExt; +/// +/// # #[tokio::main(flavor = "current_thread")] +/// # async fn main() -> Result<(), tokio::task::JoinError> { +/// let set: JoinSet<_> = (0..2).map(|i| async move { i }).collect(); +/// +/// let mut stream = JoinSetStream::new(set); +/// assert_eq!(stream.next().await.transpose()?, Some(0)); +/// assert_eq!(stream.next().await.transpose()?, Some(1)); +/// assert_eq!(stream.next().await.transpose()?, None); +/// # Ok(()) +/// # } +/// ``` +/// +/// [`tokio::task::JoinSet`]: struct@tokio::task::JoinSet +/// [`Stream`]: trait@crate::Stream +#[derive(Debug)] +pub struct JoinSetStream { + inner: JoinSet, +} + +impl JoinSetStream { + /// Create a new `JoinSetStream`. + pub fn new(join_set: JoinSet) -> Self { + Self { inner: join_set } + } + + /// Get back the inner `JoinSet`. + pub fn into_inner(self) -> JoinSet { + self.inner + } +} + +impl Stream for JoinSetStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_join_next(cx) + } + + /// Returns the bounds of the stream based on the underlying `JoinSet`. + /// + /// It returns `(set.len(), Some(set.len()))`. + fn size_hint(&self) -> (usize, Option) { + let size = self.inner.len(); + (size, Some(size)) + } +} + +impl AsRef> for JoinSetStream { + fn as_ref(&self) -> &JoinSet { + &self.inner + } +} + +impl AsMut> for JoinSetStream { + fn as_mut(&mut self) -> &mut JoinSet { + &mut self.inner + } +} + +impl From> for JoinSetStream { + fn from(join_set: JoinSet) -> Self { + Self::new(join_set) + } +} diff --git a/tokio-stream/tests/join_set_stream.rs b/tokio-stream/tests/join_set_stream.rs new file mode 100644 index 000000000..f9ebf0a56 --- /dev/null +++ b/tokio-stream/tests/join_set_stream.rs @@ -0,0 +1,39 @@ +#![cfg(feature = "rt")] + +use futures::{Stream, StreamExt}; +use std::collections::HashSet; +use tokio::task::JoinSet; +use tokio_stream::wrappers::JoinSetStream; + +#[tokio::test] +async fn size_hint_stream() { + let set: JoinSet<_> = (0..2).map(|i| async move { i }).collect(); + let mut stream = JoinSetStream::new(set); + + assert_eq!(stream.size_hint(), (2, Some(2))); + stream.next().await; + assert_eq!(stream.size_hint(), (1, Some(1))); + stream.next().await; + assert_eq!(stream.size_hint(), (0, Some(0))); +} + +#[tokio::test] +async fn join_set_as_stream() { + let set: JoinSet<_> = (0..2).map(|i| async move { i }).collect(); + let stream = JoinSetStream::new(set); + + let values: HashSet<_> = stream.map(|result| result.unwrap()).collect().await; + assert_eq!(values, HashSet::from([0, 1])); +} + +// Cannot run this test when “unwind” is disabled +// since `JoinSet` use it to catch futures that panics. +#[cfg(panic = "unwind")] +#[tokio::test] +async fn join_set_as_stream_panics_with_error() { + let set: JoinSet<_> = std::iter::once(async move { panic!("boom!") }).collect(); + let mut stream = JoinSetStream::new(set); + + let result = stream.next().await.transpose(); + assert!(matches!(result, Err(e) if e.is_panic())); +}