fix: rsdoctor multi compiler data scope (#11310)

LingyuCoder · web-flow · commit 1e3a0c1f4170 · 2025-08-06T21:38:54.000+08:00
diff --git a/crates/rspack_plugin_rsdoctor/src/chunk_graph.rs b/crates/rspack_plugin_rsdoctor/src/chunk_graph.rs
@@ -7,7 +7,6 @@ use rspack_core::{
   Chunk, ChunkByUkey, ChunkGraph, ChunkGroupByUkey, ChunkGroupUkey, ChunkUkey, CompilationAsset,
   ModuleGraph,
 };
-use rspack_util::fx_hash::FxDashMap;
 use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
 
 use crate::{
@@ -179,7 +178,7 @@ pub fn collect_assets(
 
 pub fn collect_chunk_modules(
   chunk_by_ukey: &ChunkByUkey,
-  module_ukeys: &FxDashMap<Identifier, RsdoctorChunkUkey>,
+  module_ukeys: &HashMap<Identifier, RsdoctorChunkUkey>,
   chunk_graph: &ChunkGraph,
   module_graph: &ModuleGraph,
 ) -> Vec<RsdoctorChunkModules> {
@@ -203,7 +202,7 @@ pub fn collect_chunk_modules(
               concatenated_module
                 .get_modules()
                 .iter()
-                .filter_map(|m| module_ukeys.get(&m.id).map(|u| *u)),
+                .filter_map(|m| module_ukeys.get(&m.id).copied()),
             );
           }
           res
@@ -238,7 +237,7 @@ pub fn collect_chunk_assets(
 pub fn collect_entrypoint_assets(
   entrypoints: &IndexMap<String, ChunkGroupUkey>,
   rsd_assets: &HashMap<String, RsdoctorAsset>,
-  entrypoint_ukey_map: &FxDashMap<ChunkGroupUkey, RsdoctorEntrypointUkey>,
+  entrypoint_ukey_map: &HashMap<ChunkGroupUkey, RsdoctorEntrypointUkey>,
   chunk_group_by_ukey: &ChunkGroupByUkey,
   chunk_by_ukey: &ChunkByUkey,
 ) -> Vec<RsdoctorEntrypointAssets> {
diff --git a/crates/rspack_plugin_rsdoctor/src/module_graph.rs b/crates/rspack_plugin_rsdoctor/src/module_graph.rs
@@ -7,7 +7,6 @@ use rspack_core::{
   ModuleIdsArtifact, rspack_sources::MapOptions,
 };
 use rspack_paths::Utf8PathBuf;
-use rspack_util::fx_hash::FxDashMap;
 use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
 
 use crate::{
@@ -121,7 +120,7 @@ pub fn collect_concatenated_modules(
 
 pub fn collect_module_original_sources(
   modules: &IdentifierMap<&BoxModule>,
-  module_ukeys: &FxDashMap<Identifier, ModuleUkey>,
+  module_ukeys: &HashMap<Identifier, ModuleUkey>,
   module_graph: &ModuleGraph,
   compilation: &Compilation,
 ) -> Vec<RsdoctorModuleOriginalSource> {
@@ -167,7 +166,7 @@ pub fn collect_module_original_sources(
 
 pub fn collect_module_dependencies(
   modules: &IdentifierMap<&BoxModule>,
-  module_ukeys: &FxDashMap<Identifier, ModuleUkey>,
+  module_ukeys: &HashMap<Identifier, ModuleUkey>,
   module_graph: &ModuleGraph,
 ) -> HashMap<Identifier, HashMap<Identifier, (DependencyId, RsdoctorDependency)>> {
   let dependency_ukey_counter = Arc::new(AtomicI32::new(0));
@@ -220,7 +219,7 @@ pub fn collect_module_dependencies(
 
 pub fn collect_module_ids(
   modules: &IdentifierMap<&BoxModule>,
-  module_ukeys: &FxDashMap<Identifier, ModuleUkey>,
+  module_ukeys: &HashMap<Identifier, ModuleUkey>,
   module_ids: &ModuleIdsArtifact,
 ) -> Vec<RsdoctorModuleId> {
   modules
diff --git a/crates/rspack_plugin_rsdoctor/src/plugin.rs b/crates/rspack_plugin_rsdoctor/src/plugin.rs
@@ -50,10 +50,11 @@ type ArcRsdoctorPluginHooks = Arc<AtomicRefCell<RsdoctorPluginHooks>>;
 static COMPILATION_HOOKS_MAP: LazyLock<FxDashMap<CompilationId, ArcRsdoctorPluginHooks>> =
   LazyLock::new(Default::default);
 
-static MODULE_UKEY_MAP: LazyLock<FxDashMap<Identifier, ModuleUkey>> =
-  LazyLock::new(FxDashMap::default);
-static ENTRYPOINT_UKEY_MAP: LazyLock<FxDashMap<ChunkGroupUkey, EntrypointUkey>> =
+static MODULE_UKEY_MAP: LazyLock<FxDashMap<CompilationId, HashMap<Identifier, ModuleUkey>>> =
   LazyLock::new(FxDashMap::default);
+static ENTRYPOINT_UKEY_MAP: LazyLock<
+  FxDashMap<CompilationId, HashMap<ChunkGroupUkey, EntrypointUkey>>,
+> = LazyLock::new(FxDashMap::default);
 
 #[derive(Debug, Hash, PartialEq, Eq)]
 pub enum RsdoctorPluginModuleGraphFeature {
@@ -178,10 +179,11 @@ impl RsdoctorPlugin {
 #[plugin_hook(CompilerCompilation for RsdoctorPlugin)]
 async fn compilation(
   &self,
-  _compilation: &mut Compilation,
+  compilation: &mut Compilation,
   _params: &mut CompilationParams,
 ) -> Result<()> {
-  MODULE_UKEY_MAP.clear();
+  MODULE_UKEY_MAP.insert(compilation.id(), HashMap::default());
+  ENTRYPOINT_UKEY_MAP.insert(compilation.id(), HashMap::default());
   Ok(())
 }
 
@@ -221,8 +223,13 @@ async fn optimize_chunks(&self, compilation: &mut Compilation) -> Result<Option<
     chunk_group_by_ukey,
   ));
 
-  for (entrypoint_ukey, entrypoint) in rsd_entrypoints.iter() {
-    ENTRYPOINT_UKEY_MAP.insert(*entrypoint_ukey, entrypoint.ukey);
+  {
+    let mut entrypoint_ukey_map = ENTRYPOINT_UKEY_MAP
+      .get_mut(&compilation.id())
+      .expect("should have entrypoint ukey map");
+    for (entrypoint_ukey, entrypoint) in rsd_entrypoints.iter() {
+      entrypoint_ukey_map.insert(*entrypoint_ukey, entrypoint.ukey);
+    }
   }
 
   tokio::spawn(async move {
@@ -267,18 +274,26 @@ async fn optimize_chunk_modules(&self, compilation: &mut Compilation) -> Result<
     &compilation.options.context,
   ));
 
-  for (module_id, module) in rsd_modules.iter() {
-    MODULE_UKEY_MAP.insert(*module_id, module.ukey);
+  {
+    let mut module_ukey_map = MODULE_UKEY_MAP
+      .get_mut(&compilation.id())
+      .expect("should have module ukey map");
+    for (module_id, module) in rsd_modules.iter() {
+      module_ukey_map.insert(*module_id, module.ukey);
+    }
   }
 
+  let module_ukey_map = MODULE_UKEY_MAP
+    .get(&compilation.id())
+    .expect("should have module ukey map");
   // 2. collect concatenate children
   let (child_map, parent_map) = collect_concatenated_modules(&modules);
   for (module_id, children) in child_map {
     if let Some(rsd_module) = rsd_modules.get_mut(&module_id) {
       rsd_module.modules.extend(
         children
           .iter()
-          .filter_map(|i| MODULE_UKEY_MAP.get(i).map(|u| *u))
+          .filter_map(|i| module_ukey_map.get(i).copied())
           .collect::<HashSet<_>>(),
       );
     }
@@ -290,14 +305,14 @@ async fn optimize_chunk_modules(&self, compilation: &mut Compilation) -> Result<
       rsd_module.belong_modules.extend(
         parents
           .iter()
-          .filter_map(|i| MODULE_UKEY_MAP.get(i).map(|u| *u))
+          .filter_map(|i| module_ukey_map.get(i).copied())
           .collect::<HashSet<_>>(),
       );
     }
   }
 
   // 4. collect module dependencies
-  let dependency_infos = collect_module_dependencies(&modules, &MODULE_UKEY_MAP, &module_graph);
+  let dependency_infos = collect_module_dependencies(&modules, &module_ukey_map, &module_graph);
   for (origin_module_id, dependencies) in dependency_infos {
     for (dep_module_id, (dep_id, dependency)) in dependencies {
       if let Some(rsd_module) = rsd_modules.get_mut(&dep_module_id) {
@@ -338,7 +353,7 @@ async fn optimize_chunk_modules(&self, compilation: &mut Compilation) -> Result<
 
   // 6. collect chunk modules
   let chunk_modules =
-    collect_chunk_modules(chunk_by_ukey, &MODULE_UKEY_MAP, chunk_graph, &module_graph);
+    collect_chunk_modules(chunk_by_ukey, &module_ukey_map, chunk_graph, &module_graph);
 
   tokio::spawn(async move {
     match hooks
@@ -368,8 +383,13 @@ async fn module_ids(&self, compilation: &mut Compilation) -> Result<()> {
   let hooks = RsdoctorPlugin::get_compilation_hooks(compilation.id());
   let module_graph = compilation.get_module_graph();
   let modules = module_graph.modules();
-  let rsd_module_ids =
-    collect_module_ids(&modules, &MODULE_UKEY_MAP, &compilation.module_ids_artifact);
+  let rsd_module_ids = collect_module_ids(
+    &modules,
+    &MODULE_UKEY_MAP
+      .get(&compilation.id())
+      .expect("should have module ukey map"),
+    &compilation.module_ids_artifact,
+  );
 
   tokio::spawn(async move {
     match hooks
@@ -397,8 +417,14 @@ async fn after_code_generation(&self, compilation: &mut Compilation) -> Result<(
   let hooks = RsdoctorPlugin::get_compilation_hooks(compilation.id());
   let module_graph = compilation.get_module_graph();
   let modules = module_graph.modules();
-  let rsd_module_original_sources =
-    collect_module_original_sources(&modules, &MODULE_UKEY_MAP, &module_graph, compilation);
+  let rsd_module_original_sources = collect_module_original_sources(
+    &modules,
+    &MODULE_UKEY_MAP
+      .get(&compilation.id())
+      .expect("should have module ukey map"),
+    &module_graph,
+    compilation,
+  );
 
   tokio::spawn(async move {
     match hooks
@@ -431,7 +457,9 @@ async fn after_process_assets(&self, compilation: &mut Compilation) -> Result<()
   let rsd_entrypoint_assets = collect_entrypoint_assets(
     &compilation.entrypoints,
     &rsd_assets,
-    &ENTRYPOINT_UKEY_MAP,
+    &ENTRYPOINT_UKEY_MAP
+      .get(&compilation.id())
+      .expect("should have entrypoint ukey map"),
     chunk_group_by_ukey,
     chunk_by_ukey,
   );
