fix pin violation in much of concurrentstream

Author: conradludgate

Cargo.toml

@@ -38,6 +38,7 @@ futures-lite = "1.12.0"
 pin-project = "1.0.8"
 slab = { version = "0.4.8", optional = true }
 smallvec = { version = "1.11.0", optional = true }
+futures-buffered = "0.2.6"
 
 [dev-dependencies]
 async-io = "2.3.2"

src/concurrent_stream/for_each.rs

@@ -1,5 +1,5 @@
 use super::{Consumer, ConsumerState};
-use crate::future::FutureGroup;
+use futures_buffered::FuturesUnordered;
 use futures_lite::StreamExt;
 use pin_project::pin_project;
 
@@ -22,7 +22,7 @@ where
     // NOTE: we can remove the `Arc` here if we're willing to make this struct self-referential
     count: Arc<AtomicUsize>,
     #[pin]
-    group: FutureGroup<ForEachFut<F, FutT, T, FutB>>,
+    group: FuturesUnordered<ForEachFut<F, FutT, T, FutB>>,
     limit: usize,
     f: F,
     _phantom: PhantomData<(T, FutB)>,
@@ -44,7 +44,7 @@ where
             f,
             _phantom: PhantomData,
             count: Arc::new(AtomicUsize::new(0)),
-            group: FutureGroup::new(),
+            group: FuturesUnordered::new(),
         }
     }
 }
@@ -69,7 +69,7 @@ where
         // Space was available! - insert the item for posterity
         this.count.fetch_add(1, Ordering::Relaxed);
         let fut = ForEachFut::new(this.f.clone(), future, this.count.clone());
-        this.group.as_mut().insert_pinned(fut);
+        this.group.as_mut().push(fut);
 
         ConsumerState::Continue
     }

src/concurrent_stream/from_concurrent_stream.rs

@@ -1,5 +1,5 @@
 use super::{ConcurrentStream, Consumer, ConsumerState, IntoConcurrentStream};
-use crate::future::FutureGroup;
+use futures_buffered::FuturesUnordered;
 #[cfg(all(feature = "alloc", not(feature = "std")))]
 use alloc::vec::Vec;
 use core::future::Future;
@@ -32,14 +32,14 @@ impl<T> FromConcurrentStream<T> for Vec<T> {
 #[pin_project]
 pub(crate) struct VecConsumer<'a, Fut: Future> {
     #[pin]
-    group: FutureGroup<Fut>,
+    group: FuturesUnordered<Fut>,
     output: &'a mut Vec<Fut::Output>,
 }
 
 impl<'a, Fut: Future> VecConsumer<'a, Fut> {
     pub(crate) fn new(output: &'a mut Vec<Fut::Output>) -> Self {
         Self {
-            group: FutureGroup::new(),
+            group: FuturesUnordered::new(),
             output,
         }
     }
@@ -54,7 +54,7 @@ where
     async fn send(self: Pin<&mut Self>, future: Fut) -> super::ConsumerState {
         let mut this = self.project();
         // unbounded concurrency, so we just goooo
-        this.group.as_mut().insert_pinned(future);
+        this.group.as_mut().push(future);
         ConsumerState::Continue
     }
 

src/concurrent_stream/try_for_each.rs

@@ -1,5 +1,5 @@
 use crate::concurrent_stream::ConsumerState;
-use crate::future::FutureGroup;
+use futures_buffered::FuturesUnordered;
 use crate::private::Try;
 use futures_lite::StreamExt;
 use pin_project::pin_project;
@@ -26,7 +26,7 @@ where
     count: Arc<AtomicUsize>,
     // TODO: remove the `Pin<Box>` from this signature by requiring this struct is pinned
     #[pin]
-    group: FutureGroup<TryForEachFut<F, FutT, T, FutB, B>>,
+    group: FuturesUnordered<TryForEachFut<F, FutT, T, FutB, B>>,
     limit: usize,
     residual: Option<B::Residual>,
     f: F,
@@ -50,7 +50,7 @@ where
             f,
             residual: None,
             count: Arc::new(AtomicUsize::new(0)),
-            group: FutureGroup::new(),
+            group: FuturesUnordered::new(),
             _phantom: PhantomData,
         }
     }
@@ -93,7 +93,7 @@ where
         // Space was available! - insert the item for posterity
         this.count.fetch_add(1, Ordering::Relaxed);
         let fut = TryForEachFut::new(this.f.clone(), future, this.count.clone());
-        this.group.as_mut().insert_pinned(fut);
+        this.group.as_mut().push(fut);
         ConsumerState::Continue
     }
 

src/future/future_group.rs

@@ -274,37 +274,6 @@ impl<F: Future> FutureGroup<F> {
         Key(index)
     }
 
-    /// Insert a value into a pinned `FutureGroup`
-    ///
-    /// This method is private because it serves as an implementation detail for
-    /// `ConcurrentStream`. We should never expose this publicly, as the entire
-    /// point of this crate is that we abstract the futures poll machinery away
-    /// from end-users.
-    pub(crate) fn insert_pinned(self: Pin<&mut Self>, future: F) -> Key
-    where
-        F: Future,
-    {
-        let mut this = self.project();
-        // SAFETY: inserting a value into the futures slab does not ever move
-        // any of the existing values.
-        let index = unsafe { this.futures.as_mut().get_unchecked_mut() }.insert(future);
-        this.keys.insert(index);
-        let key = Key(index);
-
-        // If our slab allocated more space we need to
-        // update our tracking structures along with it.
-        let max_len = this.futures.as_ref().capacity().max(index);
-        this.wakers.resize(max_len);
-        this.states.resize(max_len);
-
-        // Set the corresponding state
-        this.states[index].set_pending();
-        let mut readiness = this.wakers.readiness();
-        readiness.set_ready(index);
-
-        key
-    }
-
     /// Create a stream which also yields the key of each item.
     ///
     /// # Example