add taskset api

tokio-stream/Cargo.toml

@@ -26,6 +26,7 @@ io-util = ["tokio/io-util"]
 fs = ["tokio/fs"]
 sync = ["tokio/sync", "tokio-util"]
 signal = ["tokio/signal"]
+task = ["tokio-util/task"]
 
 [dependencies]
 futures-core = { version = "0.3.0" }

tokio-stream/src/macros.rs

@@ -18,6 +18,16 @@ macro_rules! cfg_io_util {
     }
 }
 
+macro_rules! cfg_task {
+    ($($item:item)*) => {
+        $(
+            #[cfg(feature = "task")]
+            #[cfg_attr(docsrs, doc(cfg(feature = "task")))]
+            $item
+        )*
+    }
+}
+
 macro_rules! cfg_net {
     ($($item:item)*) => {
         $(

tokio-stream/src/wrappers.rs

@@ -48,6 +48,11 @@ cfg_net! {
     pub use unix_listener::UnixListenerStream;
 }
 
+cfg_task! {
+    mod task;
+    pub use task::TaskSetStream;
+}
+
 cfg_io_util! {
     mod split;
     pub use split::SplitStream;

tokio-stream/src/wrappers/task.rs

@@ -0,0 +1,47 @@
+use crate::Stream;
+use std::task::Context;
+use tokio::macros::support::{Pin, Poll};
+use tokio::task::JoinError;
+use tokio_util::taskset::TaskSet;
+
+/// A wrapper around [`TaskSet`] that implements [`Stream`].
+///
+/// [`TcpListener`]: struct@tokio_util::task::TaskSet
+/// [`Stream`]: trait@crate::Stream
+#[derive(Debug, Default)]
+#[cfg_attr(docsrs, doc(cfg(feature = "task")))]
+pub struct TaskSetStream<T> {
+    inner: TaskSet<T>,
+}
+
+impl<T> TaskSetStream<T> {
+    /// Create a new `TaskSetStream`.
+    pub fn new(task_set: TaskSet<T>) -> Self {
+        Self { inner: task_set }
+    }
+
+    /// Get back the inner `TaskSet`.
+    pub fn into_inner(self) -> TaskSet<T> {
+        self.inner
+    }
+}
+
+impl<T> Stream for TaskSetStream<T> {
+    type Item = Result<T, JoinError>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.inner.poll_next_finished(cx)
+    }
+}
+
+impl<T> AsRef<TaskSet<T>> for TaskSetStream<T> {
+    fn as_ref(&self) -> &TaskSet<T> {
+        &self.inner
+    }
+}
+
+impl<T> AsMut<TaskSet<T>> for TaskSetStream<T> {
+    fn as_mut(&mut self) -> &mut TaskSet<T> {
+        &mut self.inner
+    }
+}

tokio-util/Cargo.toml

@@ -23,7 +23,7 @@ categories = ["asynchronous"]
 default = []
 
 # Shorthand for enabling everything
-full = ["codec", "compat", "io-util", "time", "net", "rt"]
+full = ["codec", "compat", "io-util", "time", "net", "rt", "task"]
 
 net = ["tokio/net"]
 compat = ["futures-io",]
@@ -32,6 +32,7 @@ time = ["tokio/time","slab"]
 io = []
 io-util = ["io", "tokio/rt", "tokio/io-util"]
 rt = ["tokio/rt"]
+task = ["tokio/rt", "futures-util"]
 
 __docs_rs = ["futures-util"]
 
@@ -43,6 +44,7 @@ futures-core = "0.3.0"
 futures-sink = "0.3.0"
 futures-io = { version = "0.3.0", optional = true }
 futures-util = { version = "0.3.0", optional = true }
+futures = { version = "0.3.0", optional = true }
 log = "0.4"
 pin-project-lite = "0.2.0"
 slab = { version = "0.4.1", optional = true } # Backs `DelayQueue`
@@ -52,6 +54,7 @@ tokio = { version = "1.0.0", path = "../tokio", features = ["full"] }
 tokio-test = { version = "0.4.0", path = "../tokio-test" }
 tokio-stream = { version = "0.1", path = "../tokio-stream" }
 
+scopeguard = "1.1.0"
 async-stream = "0.3.0"
 futures = "0.3.0"
 futures-test = "0.3.5"

tokio-util/src/cfg.rs

@@ -8,6 +8,16 @@ macro_rules! cfg_codec {
     }
 }
 
+macro_rules! cfg_task {
+    ($($item:item)*) => {
+        $(
+            #[cfg(feature = "task")]
+            #[cfg_attr(docsrs, doc(cfg(feature = "task")))]
+            $item
+        )*
+    }
+}
+
 macro_rules! cfg_compat {
     ($($item:item)*) => {
         $(

tokio-util/src/lib.rs

@@ -45,6 +45,10 @@ cfg_rt! {
     pub mod context;
 }
 
+cfg_task! {
+    pub mod taskset;
+}
+
 cfg_time! {
     pub mod time;
 }