refactor: use Module.needBuild to trigger lazy modules compile (#11507)

Author: JSerFeng

Cargo.lock

@@ -2348,11 +2348,11 @@ export interface RawJsonParserOptions {
 }
 
 export interface RawLazyCompilationOption {
-  module: ((err: Error | null, arg: RawModuleArg) => RawModuleInfo)
+  currentActiveModules: ((err: Error | null, ) => Set<string>)
   test?: RawLazyCompilationTest
   entries: boolean
   imports: boolean
-  cacheable: boolean
+  client: string
 }
 
 export interface RawLibManifestPluginOptions {
@@ -2410,11 +2410,6 @@ export interface RawLimitChunkCountPluginOptions {
   maxChunks: number
 }
 
-export interface RawModuleArg {
-  module: string
-  path: string
-}
-
 export interface RawModuleFederationRuntimePluginOptions {
   entryRuntime?: string | undefined
 }

crates/node_binding/napi-binding.d.ts

@@ -720,11 +720,11 @@ impl<'a> BuiltinPlugin<'a> {
         let js_backend = JsBackend::from(&options);
         plugins.push(Box::new(
           rspack_plugin_lazy_compilation::plugin::LazyCompilationPlugin::new(
-            options.cacheable,
             js_backend,
             options.test.map(|test| test.into()),
             options.entries,
             options.imports,
+            options.client,
           ),
         ) as Box<dyn Plugin>)
       }

crates/rspack_binding_api/src/raw_options/raw_builtins/mod.rs

@@ -5,10 +5,11 @@ use napi::{
   bindgen_prelude::{FromNapiValue, ToNapiValue, ValidateNapiValue},
 };
 use napi_derive::napi;
+use rspack_collections::IdentifierSet;
 use rspack_core::{CompilationId, CompilerId, Module, ModuleIdentifier};
 use rspack_napi::threadsafe_function::ThreadsafeFunction;
 use rspack_plugin_lazy_compilation::{
-  backend::{Backend, ModuleInfo},
+  backend::Backend,
   plugin::{LazyCompilationTest, LazyCompilationTestCheck},
 };
 use rspack_regex::RspackRegex;
@@ -88,21 +89,15 @@ pub struct RawModuleInfo {
 
 #[napi(object, object_to_js = false)]
 pub struct RawLazyCompilationOption {
-  pub module: ThreadsafeFunction<RawModuleArg, RawModuleInfo>,
+  pub current_active_modules: ThreadsafeFunction<(), std::collections::HashSet<String>>,
   pub test: Option<RawLazyCompilationTest>,
   pub entries: bool,
   pub imports: bool,
-  pub cacheable: bool,
-}
-
-#[napi(object)]
-pub struct RawModuleArg {
-  pub module: String,
-  pub path: String,
+  pub client: String,
 }
 
 pub(crate) struct JsBackend {
-  module: ThreadsafeFunction<RawModuleArg, RawModuleInfo>,
+  current_active_modules: ThreadsafeFunction<(), std::collections::HashSet<String>>,
 }
 
 impl std::fmt::Debug for JsBackend {
@@ -114,31 +109,20 @@ impl std::fmt::Debug for JsBackend {
 impl From<&RawLazyCompilationOption> for JsBackend {
   fn from(value: &RawLazyCompilationOption) -> Self {
     Self {
-      module: value.module.clone(),
+      current_active_modules: value.current_active_modules.clone(),
     }
   }
 }
 
 #[async_trait::async_trait]
 impl Backend for JsBackend {
-  async fn module(
-    &mut self,
-    identifier: ModuleIdentifier,
-    path: String,
-  ) -> rspack_error::Result<ModuleInfo> {
-    let module_info = self
-      .module
-      .call_with_sync(RawModuleArg {
-        module: identifier.to_string(),
-        path,
-      })
+  async fn current_active_modules(&mut self) -> rspack_error::Result<IdentifierSet> {
+    let active_modules = self
+      .current_active_modules
+      .call_with_sync(())
       .await
       .expect("channel should have result");
 
-    Ok(ModuleInfo {
-      active: module_info.active,
-      client: module_info.client,
-      data: module_info.data,
-    })
+    Ok(active_modules.into_iter().map(Into::into).collect())
   }
 }

crates/rspack_binding_api/src/raw_options/raw_builtins/raw_lazy_compilation.rs

@@ -11,6 +11,7 @@ version.workspace       = true
 [dependencies]
 async-trait = { workspace = true }
 rustc-hash  = { workspace = true }
+serde_json  = { workspace = true }
 tokio       = { workspace = true }
 tracing     = { workspace = true }
 

crates/rspack_plugin_lazy_compilation/Cargo.toml

@@ -1,14 +1,7 @@
-use rspack_core::ModuleIdentifier;
+use rspack_collections::IdentifierSet;
 use rspack_error::Result;
 
-pub struct ModuleInfo {
-  pub active: bool,
-  pub data: String,
-  pub client: String,
-}
-
 #[async_trait::async_trait]
 pub trait Backend: std::fmt::Debug + Send + Sync {
-  async fn module(&mut self, original_module: ModuleIdentifier, path: String)
-  -> Result<ModuleInfo>;
+  async fn current_active_modules(&mut self) -> Result<IdentifierSet>;
 }

crates/rspack_plugin_lazy_compilation/src/backend.rs

@@ -7,8 +7,8 @@ use rspack_core::{
   BuildMeta, BuildResult, ChunkGraph, CodeGenerationData, CodeGenerationResult, Compilation,
   ConcatenationScope, Context, DependenciesBlock, DependencyId, DependencyRange, FactoryMeta,
   LibIdentOptions, Module, ModuleFactoryCreateData, ModuleGraph, ModuleIdentifier, ModuleLayer,
-  ModuleType, RuntimeGlobals, RuntimeSpec, SourceType, TemplateContext, impl_module_meta_info,
-  module_namespace_promise, module_update_hash,
+  ModuleType, RuntimeGlobals, RuntimeSpec, SourceType, TemplateContext, ValueCacheVersions,
+  impl_module_meta_info, module_namespace_promise, module_update_hash,
   rspack_sources::{BoxSource, RawStringSource},
 };
 use rspack_error::{Result, impl_empty_diagnosable_trait};
@@ -32,8 +32,6 @@ pub(crate) struct LazyCompilationProxyModule {
   build_info: BuildInfo,
   build_meta: BuildMeta,
   factory_meta: Option<FactoryMeta>,
-  cacheable: bool,
-
   readable_identifier: String,
   identifier: ModuleIdentifier,
   lib_ident: Option<String>,
@@ -46,11 +44,13 @@ pub(crate) struct LazyCompilationProxyModule {
   pub resource: String,
 
   pub active: bool,
-  pub data: String,
+
   // TODO:
   // The client field may be refreshed when rspack restart,
   // so this is not safe to cache at the moment
   pub client: String,
+
+  pub need_build: bool,
 }
 
 impl ModuleSourceMapConfig for LazyCompilationProxyModule {
@@ -70,9 +70,7 @@ impl LazyCompilationProxyModule {
     lib_ident: Option<String>,
     create_data: ModuleFactoryCreateData,
     resource: String,
-    cacheable: bool,
     active: bool,
-    data: String,
     client: String,
   ) -> Self {
     let readable_identifier = format!(
@@ -86,7 +84,6 @@ impl LazyCompilationProxyModule {
     Self {
       build_info: Default::default(),
       build_meta: Default::default(),
-      cacheable,
       create_data,
       readable_identifier,
       lib_ident,
@@ -97,8 +94,8 @@ impl LazyCompilationProxyModule {
       blocks: vec![],
       dependencies: vec![],
       active,
+      need_build: false,
       client,
-      data,
     }
   }
 }
@@ -138,6 +135,10 @@ impl Module for LazyCompilationProxyModule {
     self.lib_ident.as_ref().map(|s| Cow::Borrowed(s.as_str()))
   }
 
+  fn need_build(&self, _value_cache_versions: &ValueCacheVersions) -> bool {
+    self.need_build
+  }
+
   async fn build(
     &mut self,
     _build_context: BuildContext,
@@ -177,8 +178,6 @@ impl Module for LazyCompilationProxyModule {
       self.build_info.file_dependencies = files;
     }
 
-    self.build_info.cacheable = self.cacheable;
-
     Ok(BuildResult {
       dependencies,
       blocks,
@@ -208,10 +207,10 @@ impl Module for LazyCompilationProxyModule {
     let block = self.blocks.first();
 
     let client = format!(
-      "var client = __webpack_require__(\"{}\");\nvar data = \"{}\"",
+      "var client = __webpack_require__(\"{}\");\nvar data = {};",
       ChunkGraph::get_module_id(&compilation.module_ids_artifact, *client_module)
         .expect("should have module id"),
-      self.data
+      serde_json::to_string(&self.identifier).expect("should serialize identifier")
     );
 
     let keep_active = format!(
@@ -296,7 +295,7 @@ impl Module for LazyCompilationProxyModule {
     let mut hasher = RspackHash::from(&compilation.options.output);
     module_update_hash(self, &mut hasher, compilation, runtime);
     self.active.dyn_hash(&mut hasher);
-    self.data.dyn_hash(&mut hasher);
+    self.identifier.dyn_hash(&mut hasher);
     Ok(hasher.digest(&compilation.options.output.hash_digest))
   }
 }

crates/rspack_plugin_lazy_compilation/src/module.rs

@@ -3,15 +3,16 @@ use std::{
   sync::{Arc, LazyLock},
 };
 
+use rspack_collections::IdentifierSet;
 use rspack_core::{
   BoxModule, Compilation, CompilationId, CompilationParams, CompilerCompilation, CompilerId,
-  DependencyType, EntryDependency, LibIdentOptions, Module, ModuleFactory, ModuleFactoryCreateData,
-  NormalModuleCreateData, NormalModuleFactoryModule, Plugin,
+  CompilerMake, DependencyType, EntryDependency, LibIdentOptions, Module, ModuleFactory,
+  ModuleFactoryCreateData, NormalModuleCreateData, NormalModuleFactoryModule, Plugin,
 };
 use rspack_error::Result;
 use rspack_hook::{plugin, plugin_hook};
 use rspack_regex::RspackRegex;
-use tokio::sync::Mutex;
+use tokio::sync::{Mutex, RwLock};
 
 use crate::{
   backend::Backend, factory::LazyCompilationDependencyFactory, module::LazyCompilationProxyModule,
@@ -63,18 +64,26 @@ pub struct LazyCompilationPlugin<T: Backend, F: LazyCompilationTestCheck> {
   entries: bool, // enable for entries
   imports: bool, // enable for imports
   test: Option<LazyCompilationTest<F>>,
-  cacheable: bool,
+  client: String,
+  active_modules: RwLock<IdentifierSet>,
 }
 
 impl<T: Backend, F: LazyCompilationTestCheck> LazyCompilationPlugin<T, F> {
   pub fn new(
-    cacheable: bool,
     backend: T,
     test: Option<LazyCompilationTest<F>>,
     entries: bool,
     imports: bool,
+    client: String,
   ) -> Self {
-    Self::new_inner(Mutex::new(backend), entries, imports, test, cacheable)
+    Self::new_inner(
+      Mutex::new(backend),
+      entries,
+      imports,
+      test,
+      client,
+      Default::default(),
+    )
   }
 
   async fn check_test(
@@ -180,33 +189,53 @@ async fn normal_module_factory_module(
     return Ok(());
   }
 
-  let mut backend = self.backend.lock().await;
   let module_identifier = module.identifier();
 
   let lib_ident = module.lib_ident(LibIdentOptions {
     context: module_factory_create_data.options.context.as_str(),
   });
-  let info = backend
-    .module(
-      module_identifier,
-      create_data.resource_resolve_data.resource.clone(),
-    )
-    .await?;
 
   *module = Box::new(LazyCompilationProxyModule::new(
     module_identifier,
     lib_ident.map(|ident| ident.into_owned()),
     module_factory_create_data.clone(),
     create_data.resource_resolve_data.resource.clone(),
-    self.cacheable,
-    info.active,
-    info.data,
-    info.client,
+    self
+      .active_modules
+      .read()
+      .await
+      .contains(&format!("lazy-compilation-proxy|{module_identifier}").into()),
+    self.client.clone(),
   ));
 
   Ok(())
 }
 
+#[plugin_hook(CompilerMake for LazyCompilationPlugin<T: Backend, F: LazyCompilationTestCheck>)]
+async fn compiler_make(&self, compilation: &mut Compilation) -> Result<()> {
+  let active_modules = self.backend.lock().await.current_active_modules().await?;
+  let mut module_graph = compilation.get_module_graph_mut();
+  let mut errors = vec![];
+  for module_id in &active_modules {
+    let Some(active_module) = module_graph.module_by_identifier_mut(module_id) else {
+      errors.push(rspack_error::error!("cannot find module instance for id {module_id}").into());
+      continue;
+    };
+
+    let Some(active_module) = active_module.downcast_mut::<LazyCompilationProxyModule>() else {
+      errors.push(rspack_error::error!("cannot find module instance for id {module_id}").into());
+      continue;
+    };
+
+    active_module.need_build = true;
+  }
+
+  *self.active_modules.write().await = active_modules.into_iter().collect();
+
+  compilation.extend_diagnostics(errors);
+  Ok(())
+}
+
 #[async_trait::async_trait]
 impl<T: Backend + 'static, F: LazyCompilationTestCheck + 'static> Plugin
   for LazyCompilationPlugin<T, F>
@@ -218,6 +247,8 @@ impl<T: Backend + 'static, F: LazyCompilationTestCheck + 'static> Plugin
       .normal_module_factory_hooks
       .module
       .tap(normal_module_factory_module::new(self));
+
+    ctx.compiler_hooks.make.tap(compiler_make::new(self));
     Ok(())
   }
 }

crates/rspack_plugin_lazy_compilation/src/plugin.rs

@@ -10,7 +10,7 @@
 - Bundle: modules_module_js_lazy-compilation-proxy.chunk.CURRENT_HASH.js
 - Manifest: main.LAST_HASH.hot-update.json, size: 69
 - Update: main.LAST_HASH.hot-update.js, size: 582
-- Update: modules_demo_js_lazy-compilation-proxy.LAST_HASH.hot-update.js, size: 1412
+- Update: modules_demo_js_lazy-compilation-proxy.LAST_HASH.hot-update.js, size: 1436
 
 ## Manifest
 
@@ -79,7 +79,7 @@ self["webpackHotUpdate"]("modules_demo_js_lazy-compilation-proxy", {
   \*********************************************************************************************************************/
 (function (module, __unused_webpack_exports, __webpack_require__) {
 var client = __webpack_require__("../../../../../rspack/hot/lazy-compilation-web.js?http%3A%2F%2Flocalhost%3APORT%2Flazy-compilation-using-");
-var data = "<TEST_TOOLS_ROOT>/dist/helper/loaders/hot-update.js??ruleSet[1].rules[0].use[0]!<TEST_TOOLS_ROOT>/tests/hotCases/lazy-compilation/context/modules/demo.js"
+var data = "lazy-compilation-proxy|<TEST_TOOLS_ROOT>/dist/helper/loaders/hot-update.js??ruleSet[1].rules[0].use[0]!<TEST_TOOLS_ROOT>/tests/hotCases/lazy-compilation/context/modules/demo.js";
         module.exports = __webpack_require__.e(/*! import() */ "modules_demo_js").then(__webpack_require__.bind(__webpack_require__, /*! ./modules/demo.js */ "./modules/demo.js"));
         if (module.hot) {
           module.hot.accept();