fix: use correct Cell for EsmLibraryPlugin (#12067)

* fix: should clear oncell when rebuild

* fix: use atomicRefCell to avoid mutation panic in watch mode

* chore: explictly drop and write comments about the read lock

Author: JSerFeng

Cargo.lock

@@ -12,6 +12,7 @@ version.workspace       = true
 
 [dependencies]
 async-trait                = { workspace = true }
+atomic_refcell             = { workspace = true }
 rayon                      = { workspace = true }
 regex                      = { workspace = true }
 rspack_cacheable           = { workspace = true }

crates/rspack_plugin_esm_library/Cargo.toml

@@ -407,7 +407,8 @@ impl EsmLibraryPlugin {
       chunk_link.namespace_object_sources.insert(module, source);
     }
 
-    self.links.set(link).expect("should set chunk link");
+    let mut links = self.links.borrow_mut();
+    *links = link;
     Ok(())
   }
 

crates/rspack_plugin_esm_library/src/link.rs

@@ -3,6 +3,7 @@ use std::{
   sync::{Arc, LazyLock},
 };
 
+use atomic_refcell::AtomicRefCell;
 use regex::Regex;
 use rspack_collections::{IdentifierIndexMap, IdentifierSet, UkeyMap};
 use rspack_core::{
@@ -23,7 +24,7 @@ use rspack_plugin_javascript::{
   dependency::ImportDependencyTemplate, parser_and_generator::JavaScriptParserAndGenerator,
 };
 use sugar_path::SugarPath;
-use tokio::sync::{OnceCell, RwLock};
+use tokio::sync::RwLock;
 
 use crate::{
   chunk_link::ChunkLinkContext, dependency::dyn_import::DynamicImportDependencyTemplate,
@@ -39,9 +40,10 @@ pub struct EsmLibraryPlugin {
   pub(crate) preserve_modules: Option<PathBuf>,
   // module instance will hold this map till compile done, we can't mutate it,
   // normal concatenateModule just read the info from it
-  pub(crate) concatenated_modules_map_for_codegen: OnceCell<Arc<IdentifierIndexMap<ModuleInfo>>>,
+  pub(crate) concatenated_modules_map_for_codegen:
+    AtomicRefCell<Arc<IdentifierIndexMap<ModuleInfo>>>,
   pub(crate) concatenated_modules_map: RwLock<Arc<IdentifierIndexMap<ModuleInfo>>>,
-  pub(crate) links: OnceCell<UkeyMap<ChunkUkey, ChunkLinkContext>>,
+  pub(crate) links: AtomicRefCell<UkeyMap<ChunkUkey, ChunkLinkContext>>,
 }
 
 impl EsmLibraryPlugin {
@@ -247,10 +249,9 @@ async fn finish_modules(&self, compilation: &mut Compilation) -> Result<()> {
 
   // only used for scope
   // we mutably modify data in `self.concatenated_modules_map`
-  self
-    .concatenated_modules_map_for_codegen
-    .set(Arc::new(modules_map.clone()))
-    .expect("should set concatenated_modules_map_for_codegen");
+  let mut map = self.concatenated_modules_map_for_codegen.borrow_mut();
+  *map = Arc::new(modules_map.clone());
+  drop(map);
 
   *self.concatenated_modules_map.write().await = Arc::new(modules_map);
   // mark all entry exports as used
@@ -282,10 +283,7 @@ async fn concatenation_scope(
   _compilation: &Compilation,
   module: ModuleIdentifier,
 ) -> Result<Option<ConcatenationScope>> {
-  let modules_map = self
-    .concatenated_modules_map_for_codegen
-    .get()
-    .expect("should already initialized");
+  let modules_map = self.concatenated_modules_map_for_codegen.borrow();
 
   let Some(current_module) = modules_map.get(&module) else {
     return Ok(None);

crates/rspack_plugin_esm_library/src/plugin.rs

@@ -76,15 +76,18 @@ impl EsmLibraryPlugin {
     asset_info: &mut AssetInfo,
   ) -> Result<Option<RenderSource>> {
     let module_graph = compilation.get_module_graph();
-    let chunk_link_guard = self.links.get().expect("should have chunk link");
 
-    let chunk_link = chunk_link_guard.get(chunk_ukey).expect("should have chunk");
+    // In this phase we only read from the lock, no write happen in this phase, the
+    // next write happen only happen for next compile start
+    let chunk_link_guard = self.links.borrow();
+    let chunk_link = &chunk_link_guard[chunk_ukey];
 
     let mut chunk_init_fragments: Vec<Box<dyn InitFragment<ChunkRenderContext> + 'static>> =
       chunk_link.init_fragments.clone();
     let mut replace_auto_public_path = false;
     let mut replace_static_url = false;
 
+    // Same as above, we can only read here, the write happen only at the finish_modules phase
     let concatenated_modules_map = self.concatenated_modules_map.read().await;
 
     let chunk = get_chunk(compilation, *chunk_ukey);