fix: use Mutex<Arc<tokio::sync::OnceCell<JsLoaderRunner>>> to avoid loader runner deadlock

crates/node_binding/src/plugins/js_loader/mod.rs

@@ -2,23 +2,27 @@ mod context;
 mod resolver;
 mod scheduler;
 
-use std::{ffi::c_void, fmt::Debug, ptr};
+use std::{
+  ffi::c_void,
+  fmt::Debug,
+  ptr,
+  sync::{Arc, Mutex},
+};
 
 pub use context::{JsLoaderContext, JsLoaderItem};
 use napi::{
   bindgen_prelude::*,
   sys::{napi_call_threadsafe_function, napi_threadsafe_function},
   threadsafe_function::ThreadsafeFunction,
 };
-use once_cell::sync::OnceCell;
 use rspack_core::{
   ApplyContext, Compilation, CompilationParams, CompilerEmit, CompilerId, CompilerOptions,
   CompilerThisCompilation, Plugin, PluginContext,
 };
 use rspack_error::Result;
 use rspack_hook::{plugin, plugin_hook};
 use rustc_hash::FxHashSet;
-use tokio::sync::RwLock;
+use tokio::sync::{OnceCell, RwLock};
 
 use crate::{RspackResultToNapiResultExt, COMPILER_REFERENCES};
 
@@ -144,9 +148,11 @@ impl JsLoaderRunnerGetter {
 
 #[plugin]
 pub(crate) struct JsLoaderRspackPlugin {
-  compiler_id: OnceCell<CompilerId>,
+  compiler_id: once_cell::sync::OnceCell<CompilerId>,
   pub(crate) runner_getter: JsLoaderRunnerGetter,
-  pub(crate) runner: RwLock<Option<JsLoaderRunner>>,
+  /// This complex data structure is used to avoid deadlock when running loaders which contain `importModule`
+  /// See: https://github.com/web-infra-dev/rspack/pull/10632
+  pub(crate) runner: Mutex<Arc<tokio::sync::OnceCell<JsLoaderRunner>>>,
   pub(crate) loaders_without_pitch: RwLock<FxHashSet<String>>,
 }
 
@@ -155,7 +161,7 @@ impl JsLoaderRspackPlugin {
     Self::new_inner(
       Default::default(),
       runner_getter,
-      RwLock::new(None),
+      Mutex::default(),
       RwLock::new(FxHashSet::default()),
     )
   }
@@ -180,8 +186,7 @@ async fn this_compilation(
 
 #[plugin_hook(CompilerEmit for JsLoaderRspackPlugin)]
 async fn done(&self, _compilation: &mut Compilation) -> Result<()> {
-  let mut write_guard = self.runner.write().await;
-  *write_guard = None;
+  *self.runner.lock().expect("should get lock") = Arc::new(OnceCell::new());
   Ok(())
 }
 

crates/node_binding/src/plugins/js_loader/scheduler.rs

@@ -51,64 +51,32 @@ pub(crate) async fn loader_yield(
   &self,
   loader_context: &mut LoaderContext<RunnerContext>,
 ) -> Result<()> {
-  let read_guard = self.runner.read().await;
-  match &*read_guard {
-    Some(runner) => {
-      let new_cx = runner
-        .call_async(loader_context.try_into()?)
-        .await
-        .into_diagnostic()?
-        .await
-        .into_diagnostic()?;
-      drop(read_guard);
-
-      if loader_context.state() == LoaderState::Pitching {
-        let list = collect_loaders_without_pitch(loader_context, &new_cx);
-        if !list.is_empty() {
-          self.update_loaders_without_pitch(list).await;
-        }
-      }
-
-      merge_loader_context(loader_context, new_cx)?;
-    }
-    None => {
-      drop(read_guard);
-
-      {
-        let mut write_guard = self.runner.write().await;
-        #[allow(clippy::unwrap_used)]
-        let compiler_id = self.compiler_id.get().unwrap();
-        #[allow(clippy::unwrap_used)]
-        let runner = self
-          .runner_getter
-          .call(compiler_id)
-          .await
-          .into_diagnostic()?;
-        *write_guard = Some(runner);
-      };
-
-      let read_guard = self.runner.read().await;
+  let runner = self.runner.lock().expect("should get lock").clone();
+  let runner = runner
+    .get_or_try_init(|| async {
       #[allow(clippy::unwrap_used)]
-      let new_cx = read_guard
-        .as_ref()
-        .unwrap()
-        .call_async(loader_context.try_into()?)
-        .await
-        .into_diagnostic()?
-        .await
-        .into_diagnostic()?;
-      drop(read_guard);
+      let compiler_id = self.compiler_id.get().unwrap();
+      self.runner_getter.call(compiler_id).await
+    })
+    .await
+    .into_diagnostic()?;
+
+  let new_cx = runner
+    .call_async(loader_context.try_into()?)
+    .await
+    .into_diagnostic()?
+    .await
+    .into_diagnostic()?;
+
+  if loader_context.state() == LoaderState::Pitching {
+    let list = collect_loaders_without_pitch(loader_context, &new_cx);
+    if !list.is_empty() {
+      self.update_loaders_without_pitch(list).await;
+    }
+  }
 
-      if loader_context.state() == LoaderState::Pitching {
-        let list = collect_loaders_without_pitch(loader_context, &new_cx);
-        if !list.is_empty() {
-          self.update_loaders_without_pitch(list).await;
-        }
-      }
+  merge_loader_context(loader_context, new_cx)?;
 
-      merge_loader_context(loader_context, new_cx)?;
-    }
-  };
   Ok(())
 }