fix: introduce AsyncTask, make rebuildModule safer (#9244)

fix: introduce AsyncTask to make rebuildModule safer and correct the order of invoking module_executor hook_after_make

Author: JSerFeng

Diff for: crates/rspack_binding_values/src/compilation/mod.rs

@@ -590,7 +590,6 @@ impl JsCompilation {
         )
         .await
         .map_err(|e| Error::new(napi::Status::GenericFailure, format!("{e}")))?;
-
       modules
         .iter_mut()
         .for_each(|module| module.attach(compilation));

Diff for: crates/rspack_core/src/compiler/compilation.rs

@@ -1201,13 +1201,6 @@ impl Compilation {
 
     self.in_finish_make.store(false, Ordering::Release);
 
-    // sync assets to compilation from module_executor
-    if let Some(module_executor) = &mut self.module_executor {
-      let mut module_executor = std::mem::take(module_executor);
-      module_executor.hook_after_finish_modules(self).await;
-      self.module_executor = Some(module_executor);
-    }
-
     // take built_modules
     if let Some(mutations) = self.incremental.mutations_write() {
       mutations.extend(
@@ -1239,6 +1232,14 @@ impl Compilation {
       .finish_modules
       .call(self)
       .await?;
+
+    // sync assets to compilation from module_executor
+    if let Some(module_executor) = &mut self.module_executor {
+      let mut module_executor = std::mem::take(module_executor);
+      module_executor.hook_after_finish_modules(self).await;
+      self.module_executor = Some(module_executor);
+    }
+
     logger.time_end(start);
     // Collect dependencies diagnostics at here to make sure:
     // 1. after finish_modules: has provide exports info

Diff for: packages/rspack-test-tools/tests/configCases/compilation/rebuild-module-multi-times/a.js

@@ -0,0 +1 @@
+export const a = 1;

Diff for: packages/rspack-test-tools/tests/configCases/compilation/rebuild-module-multi-times/index.js

@@ -0,0 +1,7 @@
+import { a as a1 } from "./a?1";
+import { a as a2 } from "./a?2";
+import { a as a3 } from "./a?3";
+
+it("compilation rebuild module should works", () => {
+	expect(a1 + a2 + a3).toBe(15)
+});

Diff for: packages/rspack-test-tools/tests/configCases/compilation/rebuild-module-multi-times/loader.js

@@ -0,0 +1,5 @@
+let times = 0;
+module.exports = function loader(content) {
+	times++;
+	return content.replace("1", times);
+};

Diff for: packages/rspack-test-tools/tests/configCases/compilation/rebuild-module-multi-times/rspack.config.js

@@ -0,0 +1,49 @@
+const pluginName = "plugin";
+
+class Plugin {
+	apply(compiler) {
+		let initial = true;
+		compiler.hooks.compilation.tap(pluginName, compilation => {
+			compilation.hooks.finishModules.tapPromise(pluginName, async modules => {
+				modules = [...modules];
+				if (initial) {
+					initial = false;
+
+					const results = await Promise.all(
+						modules.map(m => {
+							return new Promise((resolve, reject) => {
+								compilation.rebuildModule(m, (err, module) => {
+									if (err) {
+										reject(err);
+									} else {
+										resolve(module);
+									}
+								});
+							});
+						})
+					);
+
+					// should compile success
+					expect(results.length).toBe(4);
+				}
+			});
+		});
+	}
+}
+
+/**@type {import("@rspack/core").Configuration}*/
+module.exports = {
+	module: {
+		rules: [
+			{
+				test: /a\.js$/,
+				use: [
+					{
+						loader: "./loader"
+					}
+				]
+			}
+		]
+	},
+	plugins: [new Plugin()]
+};

Diff for: packages/rspack/etc/core.api.md

@@ -838,7 +838,7 @@ export class Compilation {
     // (undocumented)
     static PROCESS_ASSETS_STAGE_SUMMARIZE: number;
     // (undocumented)
-    rebuildModule(m: Module, f: (err: Error, m: Module) => void): void;
+    rebuildModule(m: Module, f: (err: Error | null, m: Module | null) => void): void;
     // (undocumented)
     renameAsset(filename: string, newFilename: string): void;
     // (undocumented)

Diff for: packages/rspack/src/Compilation.ts

@@ -12,7 +12,6 @@ import {
 	type ExternalObject,
 	type JsCompatSourceOwned,
 	type JsCompilation,
-	type JsModule,
 	type JsPathData,
 	JsRspackSeverity,
 	type JsRuntimeModule
@@ -53,7 +52,7 @@ import { LogType, Logger } from "./logging/Logger";
 import { StatsFactory } from "./stats/StatsFactory";
 import { StatsPrinter } from "./stats/StatsPrinter";
 import { type AssetInfo, JsAssetInfo } from "./util/AssetInfo";
-import MergeCaller from "./util/MergeCaller";
+import { AsyncTask } from "./util/AsyncTask";
 import { createReadonlyMap } from "./util/createReadonlyMap";
 import { createFakeCompilationDependencies } from "./util/fake";
 import type { InputFileSystem } from "./util/fs";
@@ -1075,27 +1074,31 @@ BREAKING CHANGE: Asset processing hooks in Compilation has been merged into a si
 		);
 	}
 
-	#rebuildModuleCaller = ((compilation: Compilation) =>
-		new MergeCaller(
-			(args: Array<[string, (err: Error, m: Module) => void]>) => {
-				compilation.#inner.rebuildModule(
-					args.map(item => item[0]),
-					(err: Error, modules: JsModule[]) => {
-						for (const [id, callback] of args) {
-							const m = modules.find(item => item.moduleIdentifier === id);
-							if (m) {
-								callback(err, Module.__from_binding(m));
-							} else {
-								callback(err || new Error("module no found"), null as any);
-							}
-						}
+	#rebuildModuleTask = new AsyncTask<string, Module>(
+		(moduleIdentifiers, doneWork) => {
+			this.#inner.rebuildModule(
+				moduleIdentifiers,
+				(err: Error | null, modules: binding.JsModule[]) => {
+					/*
+					 * 	TODO:
+					 *	batch all call parameters, once a module is failed, we cannot know which module
+					 * 	is failed to rebuild, we have to make all modules failed, this should be improved
+					 *	in the future
+					 */
+					if (err) {
+						doneWork(new Array(moduleIdentifiers.length).fill([err, null]));
+					} else {
+						doneWork(
+							modules.map(jsModule => [null, Module.__from_binding(jsModule)])
+						);
 					}
-				);
-			}
-		))(this);
+				}
+			);
+		}
+	);
 
-	rebuildModule(m: Module, f: (err: Error, m: Module) => void) {
-		this.#rebuildModuleCaller.push([m.identifier(), f]);
+	rebuildModule(m: Module, f: (err: Error | null, m: Module | null) => void) {
+		this.#rebuildModuleTask.exec(m.identifier(), f);
 	}
 
 	addRuntimeModule(chunk: Chunk, runtimeModule: RuntimeModule) {
@@ -1248,7 +1251,7 @@ class AddIncludeDispatcher {
 		this.#cbs = [];
 		this.#inner(args, (wholeErr, results) => {
 			if (this.#args.length !== 0) {
-				queueMicrotask(this.#execute);
+				queueMicrotask(this.#execute.bind(this));
 			}
 
 			if (wholeErr) {
@@ -1281,7 +1284,7 @@ class AddIncludeDispatcher {
 		callback: (err?: null | WebpackError, module?: Module) => void
 	) {
 		if (this.#args.length === 0) {
-			queueMicrotask(this.#execute);
+			queueMicrotask(this.#execute.bind(this));
 		}
 
 		this.#args.push([context, dependency, options as any]);

Diff for: packages/rspack/src/util/AsyncTask.ts

@@ -0,0 +1,52 @@
+type TaskCallback<Ret> = (err: Error | null, ret: Ret | null) => void;
+
+export class AsyncTask<Param, Ret> {
+	#isRunning = false;
+	#params: Param[] = [];
+	#callbacks: TaskCallback<Ret>[] = [];
+
+	#task: (
+		param: Param[],
+		callback: (results: [Error | null, Ret | null][]) => void
+	) => void;
+
+	constructor(
+		task: (
+			param: Param[],
+			callback: (results: [Error | null, Ret | null][]) => void
+		) => void
+	) {
+		this.#task = task;
+	}
+
+	#exec_internal() {
+		const params = this.#params;
+		const callbacks = this.#callbacks;
+		this.#params = [];
+		this.#callbacks = [];
+
+		this.#task(params, results => {
+			this.#isRunning = false;
+			if (this.#params.length) {
+				this.#isRunning = true;
+				queueMicrotask(() => this.#exec_internal());
+			}
+
+			for (let i = 0; i < results.length; i++) {
+				const [err, result] = results[i];
+				const callback = callbacks[i];
+				callback(err, result);
+			}
+		});
+	}
+
+	exec(param: Param, callback: TaskCallback<Ret>) {
+		if (!this.#isRunning) {
+			queueMicrotask(() => this.#exec_internal());
+			this.#isRunning = true;
+		}
+
+		this.#params.push(param);
+		this.#callbacks.push(callback);
+	}
+}

Diff for: website/docs/en/api/javascript-api/compilation.mdx

@@ -479,7 +479,7 @@ Triggers a re-build of the module.
 ```ts
 rebuildModule(
   m: Module, // module to be rebuilt
-  f: (err: Error, m: Module) => void //  function to be invoked when the module finishes rebuilding
+  f: (err: Error | null, m: Module | null) => void //  function to be invoked when the module finishes rebuilding
 ): void;
 ```
 

Diff for: website/docs/zh/api/javascript-api/compilation.mdx

@@ -477,7 +477,7 @@ export default {
 重新构建指定模块。
 
 ```ts
-rebuildModule(m: Module, f: (err: Error, m: Module) => void): void;
+rebuildModule(m: Module, f: (err: Error | null, m: Module | null) => void): void;
 ```
 
 以下示例将重新构建结尾为 `a.js` 模块：