chore(rolldown): use bounded channel to avoid dynamic memory allocation

Brooooooklyn · hyf0 · commit e60dc50c4b13 · 2024-04-16T15:02:52.000+08:00
diff --git a/crates/rolldown/src/module_loader/module_loader.rs b/crates/rolldown/src/module_loader/module_loader.rs
@@ -44,7 +44,7 @@ impl IntermediateNormalModules {
 pub struct ModuleLoader {
   input_options: SharedOptions,
   shared_context: Arc<TaskContext>,
-  rx: tokio::sync::mpsc::UnboundedReceiver<Msg>,
+  rx: tokio::sync::mpsc::Receiver<Msg>,
   visited: FxHashMap<Arc<str>, ModuleId>,
   runtime_id: NormalModuleId,
   remaining: u32,
@@ -72,7 +72,9 @@ impl ModuleLoader {
     fs: OsFileSystem,
     resolver: SharedResolver,
   ) -> Self {
-    let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Msg>();
+    // 1024 should be enough for most cases
+    // over 1024 pending tasks are insane
+    let (tx, rx) = tokio::sync::mpsc::channel::<Msg>(1024);
 
     let tx_to_runtime_module = tx.clone();
 
diff --git a/crates/rolldown/src/module_loader/normal_module_task.rs b/crates/rolldown/src/module_loader/normal_module_task.rs
@@ -46,11 +46,11 @@ impl NormalModuleTask {
     match self.run_inner().await {
       Ok(()) => {
         if !self.errors.is_empty() {
-          self.ctx.tx.send(Msg::BuildErrors(self.errors)).expect("Send should not fail");
+          self.ctx.tx.send(Msg::BuildErrors(self.errors)).await.expect("Send should not fail");
         }
       }
       Err(err) => {
-        self.ctx.tx.send(Msg::Panics(err)).expect("Send should not fail");
+        self.ctx.tx.send(Msg::Panics(err)).await.expect("Send should not fail");
       }
     }
   }
@@ -124,6 +124,7 @@ impl NormalModuleTask {
         raw_import_records: import_records,
         ast,
       }))
+      .await
       .expect("Send should not fail");
     tracing::trace!("end process {:?}", self.resolved_path);
     Ok(())
diff --git a/crates/rolldown/src/module_loader/runtime_normal_module_task.rs b/crates/rolldown/src/module_loader/runtime_normal_module_task.rs
@@ -15,7 +15,7 @@ use crate::{
   types::{ast_symbols::AstSymbols, normal_module_builder::NormalModuleBuilder},
 };
 pub struct RuntimeNormalModuleTask {
-  tx: tokio::sync::mpsc::UnboundedSender<Msg>,
+  tx: tokio::sync::mpsc::Sender<Msg>,
   module_id: NormalModuleId,
   warnings: Vec<BuildError>,
 }
@@ -29,7 +29,7 @@ pub struct RuntimeNormalModuleTaskResult {
 }
 
 impl RuntimeNormalModuleTask {
-  pub fn new(id: NormalModuleId, tx: tokio::sync::mpsc::UnboundedSender<Msg>) -> Self {
+  pub fn new(id: NormalModuleId, tx: tokio::sync::mpsc::Sender<Msg>) -> Self {
     Self { module_id: id, tx, warnings: Vec::default() }
   }
 
@@ -77,13 +77,15 @@ impl RuntimeNormalModuleTask {
     builder.is_user_defined_entry = Some(false);
     builder.is_virtual = true;
 
-    if let Err(_err) = self.tx.send(Msg::RuntimeNormalModuleDone(RuntimeNormalModuleTaskResult {
-      warnings: self.warnings,
-      ast_symbol: symbol,
-      builder,
-      runtime,
-      ast,
-    })) {
+    if let Err(_err) =
+      self.tx.try_send(Msg::RuntimeNormalModuleDone(RuntimeNormalModuleTaskResult {
+        warnings: self.warnings,
+        ast_symbol: symbol,
+        builder,
+        runtime,
+        ast,
+      }))
+    {
       // hyf0: If main thread is dead, we should handle errors of main thread. So we just ignore the error here.
     };
   }
diff --git a/crates/rolldown/src/module_loader/task_context.rs b/crates/rolldown/src/module_loader/task_context.rs
@@ -8,7 +8,7 @@ use super::Msg;
 /// Used to store common data shared between all tasks.
 pub struct TaskContext {
   pub input_options: SharedOptions,
-  pub tx: tokio::sync::mpsc::UnboundedSender<Msg>,
+  pub tx: tokio::sync::mpsc::Sender<Msg>,
   pub resolver: SharedResolver,
   pub fs: OsFileSystem,
   pub plugin_driver: SharedPluginDriver,

Original file line number	Diff line number	Diff line change
`@@ -46,11 +46,11 @@ impl NormalModuleTask {`
`46`	`46`	`match self.run_inner().await {`
`47`	`47`	`Ok(()) => {`
`48`	`48`	`if !self.errors.is_empty() {`
`49`		`- self.ctx.tx.send(Msg::BuildErrors(self.errors)).expect("Send should not fail");`
	`49`	`+ self.ctx.tx.send(Msg::BuildErrors(self.errors)).await.expect("Send should not fail");`
`50`	`50`	`}`
`51`	`51`	`}`
`52`	`52`	`Err(err) => {`
`53`		`- self.ctx.tx.send(Msg::Panics(err)).expect("Send should not fail");`
	`53`	`+ self.ctx.tx.send(Msg::Panics(err)).await.expect("Send should not fail");`
`54`	`54`	`}`
`55`	`55`	`}`
`56`	`56`	`}`
`@@ -124,6 +124,7 @@ impl NormalModuleTask {`
`124`	`124`	`raw_import_records: import_records,`
`125`	`125`	`ast,`
`126`	`126`	`}))`
	`127`	`+ .await`
`127`	`128`	`.expect("Send should not fail");`
`128`	`129`	`tracing::trace!("end process {:?}", self.resolved_path);`
`129`	`130`	`Ok(())`