Implement a shared Stream

Author: tmccombs

futures-util/src/future/future/shared.rs

@@ -1,5 +1,5 @@
 use crate::task::waker_ref;
-use crate::wakerset::{WakerSet, WakerKey};
+use crate::wakerset::{WakerKey, WakerSet};
 use futures_core::future::{FusedFuture, Future};
 use futures_core::task::{Context, Poll};
 use std::cell::UnsafeCell;
@@ -243,11 +243,7 @@ where
 
         inner.notifier.record_waker(&mut this.waker_key, cx);
 
-        match inner
-            .state
-            .compare_exchange(IDLE, POLLING, SeqCst, SeqCst)
-            .unwrap_or_else(|x| x)
-        {
+        match inner.state.compare_exchange(IDLE, POLLING, SeqCst, SeqCst).unwrap_or_else(|x| x) {
             IDLE => {
                 // Lock acquired, fall through
             }
@@ -296,8 +292,7 @@ where
 
             match poll_result {
                 Poll::Pending => {
-                    if inner.state.compare_exchange(POLLING, IDLE, SeqCst, SeqCst).is_ok()
-                    {
+                    if inner.state.compare_exchange(POLLING, IDLE, SeqCst, SeqCst).is_ok() {
                         // Success
                         drop(reset);
                         this.inner = Some(inner);
@@ -352,6 +347,6 @@ impl<Fut: Future> WeakShared<Fut> {
     /// Returns [`None`] if all clones of the [`Shared`] have been dropped or polled
     /// to completion.
     pub fn upgrade(&self) -> Option<Shared<Fut>> {
-        Some(Shared { inner: Some(self.0.upgrade()?), waker_key: NULL_WAKER_KEY })
+        Some(Shared { inner: Some(self.0.upgrade()?), waker_key: WakerKey::NULL })
     }
 }

futures-util/src/stream/mod.rs

@@ -25,7 +25,7 @@ pub use self::stream::{
 };
 
 #[cfg(feature = "std")]
-pub use self::stream::CatchUnwind;
+pub use self::stream::{CatchUnwind, Shared};
 
 #[cfg(feature = "alloc")]
 pub use self::stream::Chunks;

futures-util/src/stream/stream/mod.rs

@@ -256,6 +256,12 @@ mod catch_unwind;
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
 pub use self::catch_unwind::CatchUnwind;
 
+#[cfg(feature = "std")]
+mod shared;
+#[cfg(feature = "std")]
+#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
+pub use self::shared::Shared;
+
 impl<T: ?Sized> StreamExt for T where T: Stream {}
 
 /// An extension trait for `Stream`s that provides a variety of convenient
@@ -1501,6 +1507,72 @@ pub trait StreamExt: Stream {
         assert_stream::<Self::Item, _>(Box::pin(self))
     }
 
+    /// Create a cloneable handle to this stream where all handles will resolve
+    /// to the same result.
+    ///
+    /// The shared() method provides a method to convert any stream into a
+    /// cloneable stream. It enables a stream to be polled by multiple threads.
+    ///
+    /// This method is only available when the `std` feature of this library is
+    /// activiated, and it is activated by default.
+    ///
+    /// # Panics
+    /// If the capacity is zero. It must have space for at least one item.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use futures::executor::block_on;
+    /// use futures::stream::{self, StreamExt};
+    ///
+    /// let stream = stream::iter(1..=3);
+    /// let shared1 = stream.shared(4);
+    /// let shared2 = shared1.clone();
+    ///
+    /// assert_eq!(vec![1,2,3], block_on(shared1.collect::<Vec<_>>()));
+    /// assert_eq!(vec![1,2,3], block_on(shared2.collect::<Vec<_>>()));
+    /// ```
+    ///
+    /// ```
+    /// use futures::executor::block_on;
+    /// use futures::stream::{self, StreamExt};
+    /// use std::thread;
+    ///
+    /// let stream = stream::iter(1..=3);
+    /// let shared1 = stream.shared(4);
+    /// let shared2 = shared1.clone();
+    /// let join_handle = thread::spawn(move || {
+    ///     assert_eq!(vec![1,2,3], block_on(shared2.collect::<Vec<_>>()));
+    /// });
+    /// assert_eq!(vec![1,2,3], block_on(shared1.collect::<Vec<_>>()));
+    /// join_handle.join().unwrap();
+    /// ```
+    ///
+    /// ```
+    /// # futures::executor::block_on(async {
+    /// use futures::stream::{self, StreamExt};
+    ///
+    /// let stream = stream::iter(vec![1,2,3]);
+    /// let mut shared1 = stream.shared(4);
+    ///
+    /// assert_eq!(Some(1), shared1.next().await);
+    ///
+    /// let mut shared2 = shared1.clone();
+    /// assert_eq!(Some(2), shared2.next().await);
+    /// assert_eq!(Some(3), shared2.next().await);
+    /// assert_eq!(vec![2,3], shared1.collect::<Vec<_>>().await);
+    /// assert_eq!(None, shared2.next().await);
+    /// # });
+    /// ```
+    #[cfg(feature = "std")]
+    fn shared(self, capacity: usize) -> Shared<Self>
+    where
+        Self: Sized,
+        Self::Item: Clone,
+    {
+        Shared::new(self, capacity)
+    }
+
     /// An adaptor for creating a buffered list of pending futures.
     ///
     /// If this stream's item can be converted into a future, then this adaptor