perf: add iterator consumer to improve performance

crates/rspack_core/src/utils/iterator_consumer/future.rs

@@ -0,0 +1,97 @@
+use std::future::Future;
+use std::sync::Arc;
+
+use tokio::sync::mpsc::unbounded_channel;
+
+/// Tools for consume iterator which return future.
+#[async_trait::async_trait]
+pub trait FutureConsumer {
+  type Item;
+  /// Use to immediately consume the data produced by the future in the iterator
+  /// without waiting for all the data to be processed.
+  /// The closures runs in the current thread.
+  async fn fut_consume<F>(self, func: impl Fn(Self::Item) -> F + Send)
+  where
+    F: Future + Send;
+}
+
+#[async_trait::async_trait]
+impl<I, Fut> FutureConsumer for I
+where
+  I: Iterator<Item = Fut> + Send,
+  Fut: Future + Send + 'static,
+  Fut::Output: Send + 'static,
+{
+  type Item = Fut::Output;
+  async fn fut_consume<F>(self, func: impl Fn(Self::Item) -> F + Send)
+  where
+    F: Future + Send,
+  {
+    let mut rx = {
+      // Create the channel in the closure to ensure all sender are dropped when iterator completes
+      // This ensures that the receiver does not get stuck in an infinite loop.
+      let (tx, rx) = unbounded_channel::<Self::Item>();
+      let tx = Arc::new(tx);
+      self.for_each(|fut| {
+        let tx = tx.clone();
+        tokio::spawn(async move {
+          let data = fut.await;
+          tx.send(data).expect("should send success");
+        });
+      });
+      rx
+    };
+
+    while let Some(data) = rx.recv().await {
+      func(data).await;
+    }
+  }
+}
+
+#[cfg(test)]
+mod test {
+  use std::time::SystemTime;
+
+  use rspack_futures::FuturesResults;
+  use tokio::time::{sleep, Duration};
+
+  use super::FutureConsumer;
+
+  #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+  async fn available() {
+    (0..10)
+      .map(|item| async move { item * 2 })
+      .fut_consume(|item| async move { assert_eq!(item % 2, 0) })
+      .await;
+  }
+
+  #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+  async fn time_check() {
+    let start = SystemTime::now();
+    vec![100, 200]
+      .into_iter()
+      .map(|item| async move {
+        sleep(Duration::from_millis(item)).await;
+        item
+      })
+      .fut_consume(|_| async move {
+        sleep(Duration::from_millis(20)).await;
+      })
+      .await;
+    let time1 = SystemTime::now().duration_since(start).unwrap();
+
+    let start = SystemTime::now();
+    let data = vec![100, 200]
+      .into_iter()
+      .map(|item| async move {
+        sleep(Duration::from_millis(item)).await;
+        item
+      })
+      .collect::<FuturesResults<_>>();
+    for _ in data.into_inner() {
+      sleep(Duration::from_millis(20)).await;
+    }
+    let time2 = SystemTime::now().duration_since(start).unwrap();
+    assert!(time1 < time2);
+  }
+}

crates/rspack_core/src/utils/iterator_consumer/mod.rs

@@ -0,0 +1,7 @@
+mod future;
+mod rayon;
+mod rayon_fut;
+
+pub use future::FutureConsumer;
+pub use rayon::RayonConsumer;
+pub use rayon_fut::RayonFutureConsumer;

crates/rspack_core/src/utils/iterator_consumer/rayon.rs

@@ -0,0 +1,78 @@
+use std::sync::mpsc::channel;
+
+use rayon::iter::ParallelIterator;
+
+/// Tools for consume rayon iterator.
+pub trait RayonConsumer {
+  type Item: Send;
+  /// Use to immediately consume the data produced by the rayon iterator
+  /// without waiting for all the data to be processed.
+  /// The closures runs in the current thread.
+  fn consume(self, func: impl Fn(Self::Item));
+}
+
+impl<P, I> RayonConsumer for P
+where
+  P: ParallelIterator<Item = I>,
+  I: Send,
+{
+  type Item = I;
+  fn consume(self, func: impl Fn(Self::Item)) {
+    let (tx, rx) = channel::<Self::Item>();
+    std::thread::scope(|s| {
+      // move rx to s.spawn, otherwise rx.into_iter will never stop
+      s.spawn(move || {
+        self.for_each(|item| tx.send(item).expect("should send success"));
+      });
+      while let Ok(data) = rx.recv() {
+        func(data);
+      }
+    });
+  }
+}
+
+#[cfg(test)]
+mod test {
+  use std::time::{Duration, SystemTime};
+
+  use rayon::prelude::*;
+
+  use super::RayonConsumer;
+
+  #[test]
+  fn available() {
+    (0..10)
+      .into_par_iter()
+      .map(|item| item * 2)
+      .consume(|item| assert_eq!(item % 2, 0));
+  }
+
+  #[test]
+  fn time_check() {
+    let start = SystemTime::now();
+    vec![100, 200]
+      .into_par_iter()
+      .map(|item| {
+        std::thread::sleep(Duration::from_millis(item));
+        item
+      })
+      .consume(|_| {
+        std::thread::sleep(Duration::from_millis(20));
+      });
+    let time1 = SystemTime::now().duration_since(start).unwrap();
+
+    let start = SystemTime::now();
+    let data: Vec<_> = vec![100, 200]
+      .into_par_iter()
+      .map(|item| {
+        std::thread::sleep(Duration::from_millis(item));
+        item
+      })
+      .collect();
+    for _ in data {
+      std::thread::sleep(Duration::from_millis(20));
+    }
+    let time2 = SystemTime::now().duration_since(start).unwrap();
+    assert!(time1 < time2);
+  }
+}

crates/rspack_core/src/utils/iterator_consumer/rayon_fut.rs

@@ -0,0 +1,102 @@
+use std::future::Future;
+use std::sync::Arc;
+
+use rayon::iter::ParallelIterator;
+use tokio::sync::mpsc::unbounded_channel;
+
+use super::RayonConsumer;
+
+/// Tools for consume rayon iterator which return feature.
+#[async_trait::async_trait]
+pub trait RayonFutureConsumer {
+  type Item;
+  /// Use to immediately consume the data produced by the future in the rayon iterator
+  /// without waiting for all the data to be processed.
+  /// The closures runs in the current thread.
+  async fn fut_consume<F>(self, func: impl Fn(Self::Item) -> F + Send)
+  where
+    F: Future + Send;
+}
+
+#[async_trait::async_trait]
+impl<I, Fut> RayonFutureConsumer for I
+where
+  I: ParallelIterator<Item = Fut> + Send,
+  Fut: Future + Send + 'static,
+  Fut::Output: Send + 'static,
+{
+  type Item = Fut::Output;
+  async fn fut_consume<F>(self, func: impl Fn(Self::Item) -> F + Send)
+  where
+    F: Future + Send,
+  {
+    let mut rx = {
+      // Create the channel in the closure to ensure all sender are dropped when iterator completes
+      // This ensures that the receiver does not get stuck in an infinite loop.
+      let (tx, rx) = unbounded_channel::<Self::Item>();
+      let tx = Arc::new(tx);
+      self.consume(|fut| {
+        let tx = tx.clone();
+        tokio::spawn(async move {
+          let data = fut.await;
+          tx.send(data).expect("should send success");
+        });
+      });
+      rx
+    };
+
+    while let Some(data) = rx.recv().await {
+      func(data).await;
+    }
+  }
+}
+
+#[cfg(test)]
+mod test {
+  use std::time::SystemTime;
+
+  use rayon::prelude::*;
+  use rspack_futures::FuturesResults;
+  use tokio::time::{sleep, Duration};
+
+  use super::RayonFutureConsumer;
+
+  #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+  async fn available() {
+    (0..10)
+      .into_par_iter()
+      .map(|item| async move { item * 2 })
+      .fut_consume(|item| async move { assert_eq!(item % 2, 0) })
+      .await;
+  }
+
+  #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+  async fn time_check() {
+    let start = SystemTime::now();
+    vec![100, 200]
+      .into_par_iter()
+      .map(|item| async move {
+        sleep(Duration::from_millis(item)).await;
+        item
+      })
+      .fut_consume(|_| async move {
+        sleep(Duration::from_millis(20)).await;
+      })
+      .await;
+    let time1 = SystemTime::now().duration_since(start).unwrap();
+
+    let start = SystemTime::now();
+    let data = vec![100, 200]
+      .into_iter()
+      .map(|item| async move {
+        sleep(Duration::from_millis(item)).await;
+        item
+      })
+      .collect::<FuturesResults<_>>();
+    for _ in data.into_inner() {
+      sleep(Duration::from_millis(20)).await;
+    }
+    let time2 = SystemTime::now().duration_since(start).unwrap();
+    assert!(time1 < time2);
+  }
+}

crates/rspack_core/src/utils/mod.rs

@@ -19,6 +19,7 @@ mod fs_trim;
 pub use fs_trim::*;
 mod hash;
 mod identifier;
+mod iterator_consumer;
 mod module_rules;
 mod property_access;
 mod property_name;
@@ -39,6 +40,7 @@ pub use self::file_counter::FileCounter;
 pub use self::find_graph_roots::*;
 pub use self::hash::*;
 pub use self::identifier::*;
+pub use self::iterator_consumer::{FutureConsumer, RayonConsumer, RayonFutureConsumer};
 pub use self::module_rules::*;
 pub use self::property_access::*;
 pub use self::property_name::*;