Init get/set limit methods

Author: yoshuawuyts

src/from_stream.rs

@@ -17,6 +17,7 @@ pin_project! {
     pub struct FromStream<S> {
         #[pin]
         stream: S,
+        limit: Option<usize>,
     }
 }
 
@@ -26,6 +27,7 @@ where
     S: Send + Sync,
 {
     FromStream {
+        limit: None,
         stream: stream.into_stream(),
     }
 }
@@ -40,4 +42,13 @@ where
         let this = self.project();
         this.stream.poll_next(cx)
     }
+
+    fn set_limit(mut self, limit: impl Into<Option<usize>>) -> Self {
+        self.limit = limit.into();
+        self
+    }
+
+    fn limit(&self) -> Option<usize> {
+        self.limit
+    }
 }

src/par_stream/for_each.rs

@@ -19,12 +19,14 @@ pin_project_lite::pin_project! {
         exhausted: Arc<AtomicBool>,
         // Count how many tasks are executing.
         ref_count: Arc<AtomicU64>,
+        // Max concurrency limit.
+        limit: Option<usize>,
     }
 }
 
 impl ForEach {
     /// Create a new instance of `ForEach`.
-    pub fn new<S, F, Fut>(mut input: S, mut f: F) -> Self
+    pub fn new<S, F, Fut>(mut stream: S, mut f: F) -> Self
     where
         S: ParallelStream,
         F: FnMut(S::Item) -> Fut + Send + Sync + Copy + 'static,
@@ -39,10 +41,11 @@ impl ForEach {
             receiver,
             exhausted: exhausted.clone(),
             ref_count: ref_count.clone(),
+            limit: stream.limit(),
         };
 
         task::spawn(async move {
-            while let Some(item) = input.next().await {
+            while let Some(item) = stream.next().await {
                 let sender = sender.clone();
                 let exhausted = exhausted.clone();
                 let ref_count = ref_count.clone();

src/par_stream/map.rs

@@ -14,6 +14,7 @@ pin_project_lite::pin_project! {
     pub struct Map<T> {
         #[pin]
         receiver: Receiver<T>,
+        limit: Option<usize>,
     }
 }
 
@@ -26,6 +27,7 @@ impl<T: Send + 'static> Map<T> {
         Fut: Future<Output = T> + Send,
     {
         let (sender, receiver) = sync::channel(1);
+        let limit = stream.limit();
         task::spawn(async move {
             while let Some(item) = stream.next().await {
                 let sender = sender.clone();
@@ -35,7 +37,7 @@ impl<T: Send + 'static> Map<T> {
                 });
             }
         });
-        Map { receiver }
+        Map { receiver, limit }
     }
 }
 
@@ -46,6 +48,15 @@ impl<T: Send + 'static> ParallelStream for Map<T> {
         let this = self.project();
         this.receiver.poll_next(cx)
     }
+
+    fn set_limit(mut self, limit: impl Into<Option<usize>>) -> Self {
+        self.limit = limit.into();
+        self
+    }
+
+    fn limit(&self) -> Option<usize> {
+        self.limit
+    }
 }
 
 #[async_std::test]

src/par_stream/mod.rs

@@ -21,6 +21,12 @@ pub trait ParallelStream: Sized + Send + Sync + Unpin + 'static {
     /// Attempts to receive the next item from the stream.
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>;
 
+    /// Set a max concurrency limit
+    fn set_limit(self, limit: impl Into<Option<usize>>) -> Self;
+
+    /// Get the max concurrency limit
+    fn limit(&self) -> Option<usize>;
+
     /// Applies `f` to each item of this stream in parallel, producing a new
     /// stream with the results.
     fn map<F, T, Fut>(self, f: F) -> Map<T>

src/par_stream/take.rs

@@ -17,14 +17,19 @@ pin_project! {
     #[derive(Clone, Debug)]
     pub struct Take<S> {
         #[pin]
-        pub(crate) stream: S,
-        pub(crate) remaining: usize,
+        stream: S,
+        remaining: usize,
+        limit: Option<usize>,
     }
 }
 
-impl<S> Take<S> {
+impl<S: ParallelStream> Take<S> {
     pub(super) fn new(stream: S, remaining: usize) -> Self {
-        Self { stream, remaining }
+        Self {
+            limit: stream.limit(),
+            remaining,
+            stream,
+        }
     }
 }
 
@@ -44,4 +49,13 @@ impl<S: ParallelStream> ParallelStream for Take<S> {
             Poll::Ready(next)
         }
     }
+
+    fn set_limit(mut self, limit: impl Into<Option<usize>>) -> Self {
+        self.limit = limit.into();
+        self
+    }
+
+    fn limit(&self) -> Option<usize> {
+        self.limit
+    }
 }

src/vec.rs

@@ -17,6 +17,7 @@ pin_project_lite::pin_project! {
     pub struct IntoParStream<T> {
         #[pin]
         stream: FromStream<FromIter<vec::IntoIter<T>>>,
+        limit: Option<usize>,
     }
 }
 
@@ -26,6 +27,15 @@ impl<T: Send + Sync + 'static> ParallelStream for IntoParStream<T> {
         let this = self.project();
         this.stream.poll_next(cx)
     }
+
+    fn set_limit(mut self, limit: impl Into<Option<usize>>) -> Self {
+        self.limit = limit.into();
+        self
+    }
+
+    fn limit(&self) -> Option<usize> {
+        self.limit
+    }
 }
 
 impl<T: Send + Sync + 'static> IntoParallelStream for Vec<T> {
@@ -36,6 +46,7 @@ impl<T: Send + Sync + 'static> IntoParallelStream for Vec<T> {
     fn into_par_stream(self) -> Self::IntoParStream {
         IntoParStream {
             stream: from_stream(from_iter(self)),
+            limit: None,
         }
     }
 }