perf: one shot ref for class instance

crates/rspack_binding_values/src/dependency.rs

@@ -3,7 +3,7 @@ use std::{cell::RefCell, ptr::NonNull};
 use napi::{bindgen_prelude::ToNapiValue, Either, Env, JsString};
 use napi_derive::napi;
 use rspack_core::{Compilation, CompilationId, Dependency, DependencyId};
-use rspack_napi::OneShotRef;
+use rspack_napi::OneShotInstanceRef;
 use rspack_plugin_javascript::dependency::{
   CommonJsExportRequireDependency, ESMExportImportedSpecifierDependency,
   ESMImportSpecifierDependency,
@@ -140,7 +140,7 @@ impl JsDependency {
   }
 }
 
-type DependencyInstanceRefs = HashMap<DependencyId, OneShotRef<JsDependency>>;
+type DependencyInstanceRefs = HashMap<DependencyId, OneShotInstanceRef<JsDependency>>;
 
 type DependencyInstanceRefsByCompilationId =
   RefCell<HashMap<CompilationId, DependencyInstanceRefs>>;
@@ -199,9 +199,9 @@ impl ToNapiValue for JsDependencyWrapper {
       };
 
       match refs.entry(val.dependency_id) {
-        std::collections::hash_map::Entry::Occupied(occupied_entry) => {
-          let r = occupied_entry.get();
-          let instance = r.from_napi_mut_ref()?;
+        std::collections::hash_map::Entry::Occupied(mut occupied_entry) => {
+          let r = occupied_entry.get_mut();
+          let instance = &mut **r;
           instance.compilation = val.compilation;
           instance.dependency = val.dependency;
 
@@ -213,7 +213,7 @@ impl ToNapiValue for JsDependencyWrapper {
             dependency_id: val.dependency_id,
             dependency: val.dependency,
           };
-          let r = vacant_entry.insert(OneShotRef::new(env, js_dependency)?);
+          let r = vacant_entry.insert(OneShotInstanceRef::new(env, js_dependency)?);
           ToNapiValue::to_napi_value(env, r)
         }
       }

crates/rspack_binding_values/src/module.rs

@@ -8,7 +8,9 @@ use rspack_core::{
   ExportsArgument, LibIdentOptions, Module, ModuleArgument, ModuleIdentifier, RuntimeModuleStage,
   SourceType,
 };
-use rspack_napi::{napi::bindgen_prelude::*, threadsafe_function::ThreadsafeFunction, OneShotRef};
+use rspack_napi::{
+  napi::bindgen_prelude::*, threadsafe_function::ThreadsafeFunction, OneShotInstanceRef,
+};
 use rspack_plugin_runtime::RuntimeModuleFromJs;
 use rspack_util::source_map::SourceMapKind;
 use rustc_hash::FxHashMap as HashMap;
@@ -345,7 +347,7 @@ impl JsModule {
   }
 }
 
-type ModuleInstanceRefs = IdentifierMap<OneShotRef<JsModule>>;
+type ModuleInstanceRefs = IdentifierMap<OneShotInstanceRef<JsModule>>;
 
 type ModuleInstanceRefsByCompilationId = RefCell<HashMap<CompilationId, ModuleInstanceRefs>>;
 
@@ -418,9 +420,9 @@ impl ToNapiValue for JsModuleWrapper {
       };
 
       match refs.entry(module.identifier()) {
-        std::collections::hash_map::Entry::Occupied(entry) => {
-          let r = entry.get();
-          let instance = r.from_napi_mut_ref()?;
+        std::collections::hash_map::Entry::Occupied(mut entry) => {
+          let r = entry.get_mut();
+          let instance = &mut **r;
           instance.compilation = val.compilation;
           instance.module = val.module;
           ToNapiValue::to_napi_value(env, r)
@@ -432,7 +434,7 @@ impl ToNapiValue for JsModuleWrapper {
             compilation_id: val.compilation_id,
             compilation: val.compilation,
           };
-          let r = entry.insert(OneShotRef::new(env, js_module)?);
+          let r = entry.insert(OneShotInstanceRef::new(env, js_module)?);
           ToNapiValue::to_napi_value(env, r)
         }
       }

crates/rspack_napi/src/js_values/mod.rs

@@ -1,3 +1,4 @@
 pub mod js_value_ref;
+pub mod one_shot_instance_ref;
 pub mod one_shot_value_ref;
 pub mod value_ref;

crates/rspack_napi/src/js_values/one_shot_instance_ref.rs

@@ -0,0 +1,106 @@
+#![allow(clippy::not_unsafe_ptr_arg_deref)]
+
+use std::cell::{Cell, RefCell};
+use std::marker::PhantomData;
+use std::ops::{Deref, DerefMut};
+use std::ptr;
+
+use napi::bindgen_prelude::{check_status, JavaScriptClassExt, ToNapiValue};
+use napi::sys::{self, napi_env};
+use napi::{CleanupEnvHook, Env, Result};
+
+thread_local! {
+  static CLEANUP_ENV_HOOK: RefCell<Option<CleanupEnvHook<()>>> = Default::default();
+
+  // cleanup references to be executed when the JS thread exits normally
+  static GLOBAL_CLEANUP_FLAG: Cell<bool> = const { Cell::new(false) };
+}
+
+// A RAII (Resource Acquisition Is Initialization) style wrapper around `Ref` that ensures the
+// reference is unreferenced when it goes out of scope. This struct maintains a single reference
+// count and automatically cleans up when it is dropped.
+pub struct OneShotInstanceRef<T: 'static> {
+  env: napi_env,
+  napi_ref: sys::napi_ref,
+  inner: *mut T,
+  ty: PhantomData<T>,
+}
+
+impl<T: JavaScriptClassExt + 'static> OneShotInstanceRef<T> {
+  pub fn new(env: napi_env, val: T) -> Result<Self> {
+    let env_wrapper = Env::from_raw(env);
+    let mut instance = val.into_instance(&env_wrapper)?;
+
+    let mut napi_ref = ptr::null_mut();
+    check_status!(unsafe { sys::napi_create_reference(env, instance.value, 1, &mut napi_ref) })?;
+
+    CLEANUP_ENV_HOOK.with(|ref_cell| {
+      if ref_cell.borrow().is_none() {
+        let result = env_wrapper.add_env_cleanup_hook((), move |_| {
+          CLEANUP_ENV_HOOK.with_borrow_mut(|cleanup_env_hook| *cleanup_env_hook = None);
+          GLOBAL_CLEANUP_FLAG.set(true);
+        });
+        if let Ok(cleanup_env_hook) = result {
+          *ref_cell.borrow_mut() = Some(cleanup_env_hook);
+        }
+      }
+    });
+
+    Ok(Self {
+      env,
+      napi_ref,
+      inner: &mut *instance,
+      ty: PhantomData,
+    })
+  }
+}
+
+impl<T> Drop for OneShotInstanceRef<T> {
+  fn drop(&mut self) {
+    if !GLOBAL_CLEANUP_FLAG.get() {
+      unsafe { sys::napi_delete_reference(self.env, self.napi_ref) };
+    }
+  }
+}
+
+impl<T: ToNapiValue> ToNapiValue for &OneShotInstanceRef<T> {
+  unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
+    let mut result = ptr::null_mut();
+    check_status!(
+      unsafe { sys::napi_get_reference_value(env, val.napi_ref, &mut result) },
+      "Failed to get reference value"
+    )?;
+    Ok(result)
+  }
+}
+
+impl<T: ToNapiValue> ToNapiValue for &mut OneShotInstanceRef<T> {
+  unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result<sys::napi_value> {
+    let mut result = ptr::null_mut();
+    check_status!(
+      unsafe { sys::napi_get_reference_value(env, val.napi_ref, &mut result) },
+      "Failed to get reference value"
+    )?;
+    Ok(result)
+  }
+}
+
+impl<T> Deref for OneShotInstanceRef<T> {
+  type Target = T;
+
+  fn deref(&self) -> &Self::Target {
+    unsafe { &*self.inner }
+  }
+}
+
+impl<T> DerefMut for OneShotInstanceRef<T> {
+  fn deref_mut(&mut self) -> &mut T {
+    unsafe { &mut *self.inner }
+  }
+}
+
+impl<T> AsRef<T> for OneShotInstanceRef<T> {
+  fn as_ref(&self) -> &T {
+    unsafe { &*self.inner }
+  }
+}

crates/rspack_napi/src/lib.rs

@@ -24,5 +24,6 @@ pub mod napi {
   pub use napi::*;
 }
 
+pub use js_values::one_shot_instance_ref::*;
 pub use js_values::one_shot_value_ref::*;
 pub use js_values::value_ref::*;