task: add Stream wrapper for JoinSet (#8189)

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"]
 

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
+        )*
+    }
+}

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;

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<T> {
+    inner: JoinSet<T>,
+}
+
+impl<T> JoinSetStream<T> {
+    /// Create a new `JoinSetStream`.
+    pub fn new(join_set: JoinSet<T>) -> Self {
+        Self { inner: join_set }
+    }
+
+    /// Get back the inner `JoinSet`.
+    pub fn into_inner(self) -> JoinSet<T> {
+        self.inner
+    }
+}
+
+impl<T: 'static> Stream for JoinSetStream<T> {
+    type Item = Result<T, JoinError>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        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<usize>) {
+        let size = self.inner.len();
+        (size, Some(size))
+    }
+}
+
+impl<T> AsRef<JoinSet<T>> for JoinSetStream<T> {
+    fn as_ref(&self) -> &JoinSet<T> {
+        &self.inner
+    }
+}
+
+impl<T> AsMut<JoinSet<T>> for JoinSetStream<T> {
+    fn as_mut(&mut self) -> &mut JoinSet<T> {
+        &mut self.inner
+    }
+}
+
+impl<T> From<JoinSet<T>> for JoinSetStream<T> {
+    fn from(join_set: JoinSet<T>) -> Self {
+        Self::new(join_set)
+    }
+}

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()));
+}