feat(compilation): add queue handlers representing task_queue

web-infra-dev · Jan 12, 2024 · 435726f · 435726f
1 parent 21b6f6f
commit 435726f
Show file tree

Hide file tree

Showing 7 changed files with 257 additions and 28 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/rspack_core/src/compiler/compilation.rs b/crates/rspack_core/src/compiler/compilation.rs
@@ -30,7 +30,7 @@ use super::{
 use crate::{
   build_chunk_graph::build_chunk_graph,
   cache::{use_code_splitting_cache, Cache, CodeSplittingCache},
-  get_chunk_from_ukey, get_mut_chunk_from_ukey, is_source_equal,
+  create_queue_handle, get_chunk_from_ukey, get_mut_chunk_from_ukey, is_source_equal,
   tree_shaking::{optimizer, visitor::SymbolRef, BailoutFlag, OptimizeDependencyResult},
   AddQueue, AddTask, AddTaskResult, AdditionalChunkRuntimeRequirementsArgs,
   AdditionalModuleRequirementsArgs, AsyncDependenciesBlock, BoxDependency, BoxModule, BuildQueue,
@@ -40,11 +40,12 @@ use crate::{
   CompilationLogging, CompilerOptions, ContentHashArgs, ContextDependency, DependencyId,
   DependencyParents, DependencyType, Entry, EntryData, EntryOptions, Entrypoint, ErrorSpan,
   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,
+  ModuleCreationCallback, ModuleFactory, ModuleFactoryResult, ModuleGraph, ModuleGraphModule,
+  ModuleIdentifier, ModuleProfile, NormalModuleSource, PathData, ProcessAssetsArgs,
+  ProcessDependenciesQueue, ProcessDependenciesResult, ProcessDependenciesTask, QueueHandler,
+  RenderManifestArgs, Resolve, ResolverFactory, RuntimeGlobals, RuntimeModule,
+  RuntimeRequirementsInTreeArgs, RuntimeSpec, SharedPluginDriver, SourceType, Stats, TaskResult,
+  WorkerTask,
 };
 use crate::{tree_shaking::visitor::OptimizeAnalyzeResult, Context};
 
@@ -109,6 +110,8 @@ pub struct Compilation {
   pub build_dependencies: IndexSet<PathBuf, BuildHasherDefault<FxHasher>>,
   pub side_effects_free_modules: IdentifierSet,
   pub module_item_map: IdentifierMap<Vec<ModuleItem>>,
+
+  pub queue_handle: Option<QueueHandler>,
 }
 
 impl Compilation {
@@ -170,6 +173,8 @@ impl Compilation {
       side_effects_free_modules: IdentifierSet::default(),
       module_item_map: IdentifierMap::default(),
       include_module_ids: IdentifierSet::default(),
+
+      queue_handle: None,
     }
   }
 
@@ -529,14 +534,19 @@ impl Compilation {
     let mut active_task_count = 0usize;
     let is_expected_shutdown = Arc::new(AtomicBool::new(false));
     let (result_tx, mut result_rx) = tokio::sync::mpsc::unbounded_channel::<Result<TaskResult>>();
+
     let mut factorize_queue = FactorizeQueue::new();
     let mut add_queue = AddQueue::new();
     let mut build_queue = BuildQueue::new();
     let mut process_dependencies_queue = ProcessDependenciesQueue::new();
+
     let mut make_failed_dependencies: HashSet<BuildDependency> = HashSet::default();
     let mut make_failed_module: HashSet<ModuleIdentifier> = HashSet::default();
     let mut errored = None;
 
+    let (handler, mut processor) = create_queue_handle();
+    self.queue_handle = Some(handler);
+
     deps_builder
       .revoke_modules(&mut self.module_graph)
       .into_iter()
@@ -587,6 +597,14 @@ impl Compilation {
 
     tokio::task::block_in_place(|| loop {
       let start = factorize_time.start();
+      processor.try_process(
+        self,
+        &mut factorize_queue,
+        &mut add_queue,
+        &mut build_queue,
+        &mut process_dependencies_queue,
+      );
+
       while let Some(task) = factorize_queue.get_task() {
         tokio::spawn({
           let result_tx = result_tx.clone();
@@ -729,6 +747,43 @@ impl Compilation {
 
       match result_rx.try_recv() {
         Ok(item) => {
+          if let Ok(item) = &item {
+            match item {
+              TaskResult::Factorize(result) => {
+                if let Some(ModuleFactoryResult {
+                  module: Some(module),
+                  ..
+                }) = &result.factory_result
+                {
+                  processor.complete_task(
+                    crate::WaitTask::Factorize(result.dependency),
+                    module.identifier(),
+                    self,
+                  )
+                }
+              }
+              TaskResult::Add(result) => {
+                let module = match result.as_ref() {
+                  AddTaskResult::ModuleReused { module } => module.identifier(),
+                  AddTaskResult::ModuleAdded { module, .. } => module.identifier(),
+                };
+
+                processor.complete_task(crate::WaitTask::Add(module), module, self)
+              }
+              TaskResult::Build(result) => {
+                let id = result.module.identifier();
+                processor.complete_task(crate::WaitTask::Build(id), id, self);
+              }
+              TaskResult::ProcessDependencies(result) => {
+                processor.complete_task(
+                  crate::WaitTask::ProcessDependencies(result.module_identifier),
+                  result.module_identifier,
+                  self,
+                );
+              }
+            }
+          }
+
           match item {
             Ok(TaskResult::Factorize(box task_result)) => {
               let FactorizeTaskResult {
@@ -743,6 +798,7 @@ impl Compilation {
                 missing_dependencies,
                 diagnostics,
                 callback,
+                ..
               } = task_result;
               if !diagnostics.is_empty() {
                 if let Some(id) = original_module_identifier {