init limit functionality

Updates the traits to support limiting

Author: yoshuawuyts

src/from_stream.rs

@@ -43,12 +43,12 @@ where
         this.stream.poll_next(cx)
     }
 
-    fn set_limit(mut self, limit: impl Into<Option<usize>>) -> Self {
+    fn limit(mut self, limit: impl Into<Option<usize>>) -> Self {
         self.limit = limit.into();
         self
     }
 
-    fn limit(&self) -> Option<usize> {
+    fn get_limit(&self) -> Option<usize> {
         self.limit
     }
 }

src/lib.rs

@@ -56,3 +56,5 @@ pub use par_stream::{ForEach, Map, NextFuture, ParallelStream, Take};
 
 pub mod prelude;
 pub mod vec;
+
+pub(crate) mod utils;

src/par_stream/for_each.rs

@@ -19,8 +19,6 @@ pin_project_lite::pin_project! {
         exhausted: Arc<AtomicBool>,
         // Count how many tasks are executing.
         ref_count: Arc<AtomicU64>,
-        // Max concurrency limit.
-        limit: Option<usize>,
     }
 }
 
@@ -35,13 +33,13 @@ impl ForEach {
         let exhausted = Arc::new(AtomicBool::new(false));
         let ref_count = Arc::new(AtomicU64::new(0));
         let (sender, receiver): (Sender<()>, Receiver<()>) = sync::channel(1);
+        let _limit = stream.get_limit();
 
         // Initialize the return type here to prevent borrowing issues.
         let this = Self {
             receiver,
             exhausted: exhausted.clone(),
             ref_count: ref_count.clone(),
-            limit: stream.limit(),
         };
 
         task::spawn(async move {

src/par_stream/map.rs

@@ -27,7 +27,7 @@ impl<T: Send + 'static> Map<T> {
         Fut: Future<Output = T> + Send,
     {
         let (sender, receiver) = sync::channel(1);
-        let limit = stream.limit();
+        let limit = stream.get_limit();
         task::spawn(async move {
             while let Some(item) = stream.next().await {
                 let sender = sender.clone();
@@ -49,12 +49,12 @@ impl<T: Send + 'static> ParallelStream for Map<T> {
         this.receiver.poll_next(cx)
     }
 
-    fn set_limit(mut self, limit: impl Into<Option<usize>>) -> Self {
+    fn limit(mut self, limit: impl Into<Option<usize>>) -> Self {
         self.limit = limit.into();
         self
     }
 
-    fn limit(&self) -> Option<usize> {
+    fn get_limit(&self) -> Option<usize> {
         self.limit
     }
 }

src/par_stream/mod.rs

@@ -22,10 +22,10 @@ pub trait ParallelStream: Sized + Send + Sync + Unpin + 'static {
     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;
+    fn limit(self, limit: impl Into<Option<usize>>) -> Self;
 
     /// Get the max concurrency limit
-    fn limit(&self) -> Option<usize>;
+    fn get_limit(&self) -> Option<usize>;
 
     /// Applies `f` to each item of this stream in parallel, producing a new
     /// stream with the results.

src/par_stream/take.rs

@@ -26,7 +26,7 @@ pin_project! {
 impl<S: ParallelStream> Take<S> {
     pub(super) fn new(stream: S, remaining: usize) -> Self {
         Self {
-            limit: stream.limit(),
+            limit: stream.get_limit(),
             remaining,
             stream,
         }
@@ -50,12 +50,12 @@ impl<S: ParallelStream> ParallelStream for Take<S> {
         }
     }
 
-    fn set_limit(mut self, limit: impl Into<Option<usize>>) -> Self {
+    fn limit(mut self, limit: impl Into<Option<usize>>) -> Self {
         self.limit = limit.into();
         self
     }
 
-    fn limit(&self) -> Option<usize> {
+    fn get_limit(&self) -> Option<usize> {
         self.limit
     }
 }

src/utils.rs

@@ -0,0 +1,55 @@
+// use core::pin::Pin;
+// use core::task::{Context, Poll};
+
+// use std::sync::atomic::{AtomicUsize, Ordering};
+// use std::sync::Arc;
+
+// use async_std::stream::Stream;
+
+// /// A stream that has a max concurrency of N.
+// pub(crate) struct LimitStream {
+//     limit: Option<usize>,
+//     ref_count: Arc<AtomicUsize>,
+// }
+
+// impl LimitStream {
+//     /// Create a new instance of LimitStream.
+//     pub(crate) fn new(limit: Option<usize>) -> Self {
+//         Self {
+//             limit,
+//             ref_count: Arc::new(AtomicUsize::new(0)),
+//         }
+//     }
+// }
+
+// #[derive(Debug)]
+// pub(crate) struct Guard {
+//     limit: Option<usize>,
+//     ref_count: Arc<AtomicUsize>,
+// }
+
+// impl Guard {
+//     fn new(limit: Option<usize>, ref_count: Arc<AtomicUsize>) -> Self {
+//         Self { limit, ref_count }
+//     }
+// }
+
+// impl Drop for Guard {
+//     fn drop(&mut self) {
+//         if self.limit.is_some() {
+//             self.ref_count.fetch_sub(1, Ordering::SeqCst);
+//         }
+//     }
+// }
+
+// impl Stream for LimitStream {
+//     type Item = Guard;
+
+//     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+//         if self.limit.is_none() {
+//             let guard = Guard::new(self.limit, self.ref_count.clone());
+//             return Poll::Ready(Some(guard));
+//         }
+//         todo!();
+//     }
+// }

src/vec.rs

@@ -28,12 +28,12 @@ impl<T: Send + Sync + 'static> ParallelStream for IntoParStream<T> {
         this.stream.poll_next(cx)
     }
 
-    fn set_limit(mut self, limit: impl Into<Option<usize>>) -> Self {
+    fn limit(mut self, limit: impl Into<Option<usize>>) -> Self {
         self.limit = limit.into();
         self
     }
 
-    fn limit(&self) -> Option<usize> {
+    fn get_limit(&self) -> Option<usize> {
         self.limit
     }
 }