From 23da8892ea15899e9d861f004e7bb827a7c8e5f4 Mon Sep 17 00:00:00 2001
From: Cong-Cong <dacongsama@live.com>
Date: Wed, 7 Aug 2024 18:15:39 +0800
Subject: [PATCH] fix: use add_env_cleanup_hook

---
 crates/node_binding/binding.d.ts              |  4 ++--
 crates/node_binding/src/lib.rs                | 12 +----------
 .../src/compilation/mod.rs                    | 20 +++++++++++++++++--
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/crates/node_binding/binding.d.ts b/crates/node_binding/binding.d.ts
index d7440ed818c0..014c2768ab10 100644
--- a/crates/node_binding/binding.d.ts
+++ b/crates/node_binding/binding.d.ts
@@ -64,7 +64,7 @@ export class JsCompilation {
   getAssets(): Readonly<JsAsset>[]
   getAsset(name: string): JsAsset | null
   getAssetSource(name: string): JsCompatSource | null
-  get modules(): Array<ModuleDTOSingleton>
+  get modules(): Readonly<ModuleDTO>[]
   getOptimizationBailout(): Array<JsStatsOptimizationBailout>
   getChunks(): Array<JsChunk>
   getNamedChunkKeys(): Array<string>
@@ -351,7 +351,7 @@ export interface JsBuildTimeExecutionOption {
 }
 
 export interface JsCacheGroupTestCtx {
-  module: ModuleDTOSingleton
+  module: ModuleDTO
 }
 
 export interface JsChunk {
diff --git a/crates/node_binding/src/lib.rs b/crates/node_binding/src/lib.rs
index f573251445ea..a0c69cc1daa3 100644
--- a/crates/node_binding/src/lib.rs
+++ b/crates/node_binding/src/lib.rs
@@ -44,25 +44,15 @@ impl Rspack {
     output_filesystem: ThreadsafeNodeFS,
     mut resolver_factory_reference: Reference<JsResolverFactory>,
   ) -> Result<Self> {
-    tracing::info!("raw_options: {:#?}", &options);
-
     let raw_env = env.raw();
+    // TODO: use napi_add_finalizer if N-API version >= 5
     env.add_env_cleanup_hook(raw_env, move |raw_env| {
-      // TODO: wait for a better hook to clean up
       COMPILATION_INSTANCE_REFS.with(|refs| {
         let mut refs = refs.borrow_mut();
         for (_, mut r) in refs.drain() {
           let _ = r.unref(raw_env);
         }
       });
-      MODULE_INSTANCE_REFS.with(|refs| {
-        let mut refs_by_compilation_id = refs.borrow_mut();
-        for (_, mut refs) in refs_by_compilation_id.drain() {
-          for (_, mut r) in refs.drain() {
-            let _ = r.unref(raw_env);
-          }
-        }
-      });
     })?;
 
     let mut plugins = Vec::new();
diff --git a/crates/rspack_binding_values/src/compilation/mod.rs b/crates/rspack_binding_values/src/compilation/mod.rs
index 8499f3c3d988..842060866b4d 100644
--- a/crates/rspack_binding_values/src/compilation/mod.rs
+++ b/crates/rspack_binding_values/src/compilation/mod.rs
@@ -29,6 +29,7 @@ use crate::utils::callbackify;
 use crate::JsStatsOptimizationBailout;
 use crate::LocalJsFilename;
 use crate::ModuleDTOWrapper;
+use crate::MODULE_INSTANCE_REFS;
 use crate::{
   chunk::JsChunk, CompatSource, JsAsset, JsAssetInfo, JsChunkGroup, JsCompatSource, JsPathData,
   JsStats, ToJsCompatSource,
@@ -132,7 +133,7 @@ impl JsCompilation {
       .transpose()
   }
 
-  #[napi(getter, ts_return_type = "ModuleDTO")]
+  #[napi(getter, ts_return_type = "Readonly<ModuleDTO>[]")]
   pub fn modules(&'static self) -> Vec<ModuleDTOWrapper> {
     self
       .0
@@ -581,6 +582,8 @@ impl JsCompilationWrapper {
 
 impl ToNapiValue for JsCompilationWrapper {
   unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
+    let mut env_wrapper = Env::from_raw(env);
+
     COMPILATION_INSTANCE_REFS.with(|refs| {
       let mut refs = refs.borrow_mut();
       match refs.entry(val.0.id()) {
@@ -589,7 +592,20 @@ impl ToNapiValue for JsCompilationWrapper {
           ToNapiValue::to_napi_value(env, r)
         }
         std::collections::hash_map::Entry::Vacant(entry) => {
-          let instance = JsCompilation(val.0).into_instance(Env::from_raw(env))?;
+          // TODO: use napi_add_finalizer if N-API version >= 5
+          let _ = env_wrapper.add_env_cleanup_hook(env, move |env| {
+            MODULE_INSTANCE_REFS.with(|refs| {
+              let mut refs_by_compilation_id = refs.borrow_mut();
+              let refs = refs_by_compilation_id.remove(&val.0.id());
+              if let Some(mut refs) = refs {
+                for (_, mut r) in refs.drain() {
+                  let _ = r.unref(env);
+                }
+              }
+            });
+          });
+
+          let instance = JsCompilation(val.0).into_instance(env_wrapper)?;
           let napi_value = ToNapiValue::to_napi_value(env, instance)?;
           let r = Ref::new(env, napi_value, 1)?;
           let r = entry.insert(r);