refactor(compilation): Add callback to factorizeTask and make process…

… dependencies happen after all dependencies get factorized (#4835)

refactor(compilation): process dependencies should happen after all dependencies get factorized

crates/rspack_core/src/compiler/compilation.rs

@@ -39,12 +39,12 @@ use crate::{
   CleanTask, CleanTaskResult, CodeGenerationResult, CodeGenerationResults, CompilationLogger,
   CompilationLogging, CompilerOptions, ContentHashArgs, ContextDependency, DependencyId,
   DependencyParents, DependencyType, Entry, EntryData, EntryOptions, Entrypoint, ErrorSpan,
-  FactorizeQueue, FactorizeTask, FactorizeTaskResult, Filename, Logger, Module, ModuleFactory,
-  ModuleGraph, ModuleGraphModule, ModuleIdentifier, ModuleProfile, NormalModuleSource, PathData,
-  ProcessAssetsArgs, ProcessDependenciesQueue, ProcessDependenciesResult, ProcessDependenciesTask,
-  RenderManifestArgs, Resolve, ResolverFactory, RuntimeGlobals, RuntimeModule,
-  RuntimeRequirementsInTreeArgs, RuntimeSpec, SharedPluginDriver, SourceType, Stats, TaskResult,
-  WorkerTask,
+  FactorizeQueue, FactorizeTask, FactorizeTaskResult, Filename, Logger, Module,
+  ModuleCreationCallback, ModuleFactory, ModuleGraph, ModuleGraphModule, ModuleIdentifier,
+  ModuleProfile, NormalModuleSource, PathData, ProcessAssetsArgs, ProcessDependenciesQueue,
+  ProcessDependenciesResult, ProcessDependenciesTask, RenderManifestArgs, Resolve, ResolverFactory,
+  RuntimeGlobals, RuntimeModule, RuntimeRequirementsInTreeArgs, RuntimeSpec, SharedPluginDriver,
+  SourceType, Stats, TaskResult, WorkerTask,
 };
 use crate::{tree_shaking::visitor::OptimizeAnalyzeResult, Context};
 
@@ -568,6 +568,7 @@ impl Compilation {
           parent_module
             .and_then(|m| m.as_normal_module())
             .and_then(|module| module.name_for_condition()),
+          None,
         );
       });
 
@@ -681,7 +682,10 @@ impl Compilation {
           .module_by_identifier(original_module_identifier)
           .expect("Module expected");
 
+        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
+        let mut remaining = sorted_dependencies.len();
         for dependencies in sorted_dependencies.into_values() {
+          let tx = tx.clone();
           self.handle_module_creation(
             &mut factorize_queue,
             Some(module.identifier()),
@@ -693,16 +697,33 @@ impl Compilation {
             module
               .as_normal_module()
               .and_then(|module| module.name_for_condition()),
+            Some(Box::new(move |_| {
+              tx.send(())
+                .expect("Failed to send callback to process_dependencies");
+            })),
           );
         }
+        drop(tx);
 
-        result_tx
-          .send(Ok(TaskResult::ProcessDependencies(Box::new(
+        let tx = result_tx.clone();
+
+        tokio::spawn(async move {
+          loop {
+            if remaining == 0 {
+              break;
+            }
+
+            rx.recv().await;
+            remaining -= 1;
+          }
+
+          tx.send(Ok(TaskResult::ProcessDependencies(Box::new(
             ProcessDependenciesResult {
               module_identifier: task.original_module_identifier,
             },
           ))))
           .expect("Failed to send process dependencies result");
+        });
       }
       process_deps_time.end(start);
 
@@ -721,6 +742,7 @@ impl Compilation {
                 context_dependencies,
                 missing_dependencies,
                 diagnostics,
+                callback,
               } = task_result;
               if !diagnostics.is_empty() {
                 if let Some(id) = original_module_identifier {
@@ -780,6 +802,7 @@ impl Compilation {
                     dependencies,
                     is_entry,
                     current_profile,
+                    callback,
                   });
                   tracing::trace!("Module created: {}", &module_identifier);
                 } else {
@@ -1119,6 +1142,7 @@ impl Compilation {
     resolve_options: Option<Box<Resolve>>,
     lazy_visit_modules: std::collections::HashSet<String>,
     issuer: Option<Box<str>>,
+    callback: Option<ModuleCreationCallback>,
   ) {
     let current_profile = self.options.profile.then(Box::<ModuleProfile>::default);
     let dependency = dependencies[0].get_dependency(&self.module_graph).clone();
@@ -1149,6 +1173,7 @@ impl Compilation {
       plugin_driver: self.plugin_driver.clone(),
       cache: self.cache.clone(),
       current_profile,
+      callback,
     });
   }
 

crates/rspack_core/src/compiler/queue.rs

@@ -11,7 +11,7 @@ use crate::{
   ModuleGraph, ModuleGraphModule, ModuleIdentifier, ModuleProfile, Resolve, ResolverFactory,
   SharedPluginDriver, WorkerQueue,
 };
-use crate::{DependencyId, ExportInfo, ExportsInfo, UsageState};
+use crate::{BoxModule, DependencyId, ExportInfo, ExportsInfo, UsageState};
 
 #[derive(Debug)]
 pub enum TaskResult {
@@ -43,6 +43,7 @@ pub struct FactorizeTask {
   pub plugin_driver: SharedPluginDriver,
   pub cache: Arc<Cache>,
   pub current_profile: Option<Box<ModuleProfile>>,
+  pub callback: Option<ModuleCreationCallback>,
 }
 
 /// a struct temporarily used creating ExportsInfo
@@ -52,7 +53,7 @@ pub struct ExportsInfoRelated {
   pub other_exports_info: ExportInfo,
   pub side_effects_info: ExportInfo,
 }
-#[derive(Debug)]
+
 pub struct FactorizeTaskResult {
   pub original_module_identifier: Option<ModuleIdentifier>,
   /// Result will be available if [crate::ModuleFactory::create] returns `Ok`.
@@ -66,6 +67,28 @@ pub struct FactorizeTaskResult {
   pub context_dependencies: HashSet<PathBuf>,
   pub missing_dependencies: HashSet<PathBuf>,
   pub diagnostics: Vec<Diagnostic>,
+  pub callback: Option<ModuleCreationCallback>,
+}
+
+impl std::fmt::Debug for FactorizeTaskResult {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    f.debug_struct("FactorizeTaskResult")
+      .field(
+        "original_module_identifier",
+        &self.original_module_identifier,
+      )
+      .field("factory_result", &self.factory_result)
+      .field("dependencies", &self.dependencies)
+      .field("diagnostics", &self.diagnostics)
+      .field("is_entry", &self.is_entry)
+      .field("current_profile", &self.current_profile)
+      .field("exports_info_related", &self.exports_info_related)
+      .field("file_dependencies", &self.file_dependencies)
+      .field("context_dependencies", &self.context_dependencies)
+      .field("missing_dependencies", &self.missing_dependencies)
+      .field("diagnostics", &self.diagnostics)
+      .finish()
+  }
 }
 
 impl FactorizeTaskResult {
@@ -134,6 +157,7 @@ impl WorkerTask for FactorizeTask {
       context_dependencies: Default::default(),
       missing_dependencies: Default::default(),
       diagnostics: Default::default(),
+      callback: self.callback,
     };
 
     // Error and result are not mutually exclusive in webpack module factorization.
@@ -204,6 +228,7 @@ pub struct AddTask {
   pub dependencies: Vec<DependencyId>,
   pub is_entry: bool,
   pub current_profile: Option<Box<ModuleProfile>>,
+  pub callback: Option<ModuleCreationCallback>,
 }
 
 #[derive(Debug)]
@@ -256,6 +281,10 @@ impl AddTask {
         module_identifier,
       )?;
 
+      if let Some(callback) = self.callback {
+        callback(&self.module);
+      }
+
       return Ok(TaskResult::Add(Box::new(AddTaskResult::ModuleReused {
         module: self.module,
       })));
@@ -282,6 +311,10 @@ impl AddTask {
       current_profile.mark_integration_end();
     }
 
+    if let Some(callback) = self.callback {
+      callback(&self.module);
+    }
+
     Ok(TaskResult::Add(Box::new(AddTaskResult::ModuleAdded {
       module: self.module,
       current_profile: self.current_profile,
@@ -463,3 +496,5 @@ impl CleanTask {
 }
 
 pub type CleanQueue = WorkerQueue<CleanTask>;
+
+pub type ModuleCreationCallback = Box<dyn FnOnce(&BoxModule) + Send>;