From f696097264a20b94bc91ec01d5e3ff6d713fbc9b Mon Sep 17 00:00:00 2001 From: Cong-Cong Date: Tue, 7 Jan 2025 16:50:04 +0800 Subject: [PATCH] perf: one shot ref for class instance --- .../rspack_binding_values/src/dependency.rs | 12 +- crates/rspack_binding_values/src/module.rs | 14 ++- crates/rspack_napi/src/js_values/mod.rs | 1 + .../src/js_values/one_shot_instance_ref.rs | 106 ++++++++++++++++++ crates/rspack_napi/src/lib.rs | 1 + 5 files changed, 122 insertions(+), 12 deletions(-) create mode 100644 crates/rspack_napi/src/js_values/one_shot_instance_ref.rs diff --git a/crates/rspack_binding_values/src/dependency.rs b/crates/rspack_binding_values/src/dependency.rs index 4381cf7b3464..5bbc84d0dc12 100644 --- a/crates/rspack_binding_values/src/dependency.rs +++ b/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>; +type DependencyInstanceRefs = HashMap>; type DependencyInstanceRefsByCompilationId = RefCell>; @@ -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) } } diff --git a/crates/rspack_binding_values/src/module.rs b/crates/rspack_binding_values/src/module.rs index 1325fbfe18c3..71f4bce98ea8 100644 --- a/crates/rspack_binding_values/src/module.rs +++ b/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>; +type ModuleInstanceRefs = IdentifierMap>; type ModuleInstanceRefsByCompilationId = RefCell>; @@ -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) } } diff --git a/crates/rspack_napi/src/js_values/mod.rs b/crates/rspack_napi/src/js_values/mod.rs index e63b3ddf4984..a1dfd46a1e53 100644 --- a/crates/rspack_napi/src/js_values/mod.rs +++ b/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; diff --git a/crates/rspack_napi/src/js_values/one_shot_instance_ref.rs b/crates/rspack_napi/src/js_values/one_shot_instance_ref.rs new file mode 100644 index 000000000000..b4f3db58e99a --- /dev/null +++ b/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>> = Default::default(); + + // cleanup references to be executed when the JS thread exits normally + static GLOBAL_CLEANUP_FLAG: Cell = 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 { + env: napi_env, + napi_ref: sys::napi_ref, + inner: *mut T, + ty: PhantomData, +} + +impl OneShotInstanceRef { + pub fn new(env: napi_env, val: T) -> Result { + 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 Drop for OneShotInstanceRef { + fn drop(&mut self) { + if !GLOBAL_CLEANUP_FLAG.get() { + unsafe { sys::napi_delete_reference(self.env, self.napi_ref) }; + } + } +} + +impl ToNapiValue for &OneShotInstanceRef { + unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result { + 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 ToNapiValue for &mut OneShotInstanceRef { + unsafe fn to_napi_value(env: sys::napi_env, val: Self) -> Result { + 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 Deref for OneShotInstanceRef { + type Target = T; + + fn deref(&self) -> &Self::Target { + unsafe { &*self.inner } + } +} + +impl DerefMut for OneShotInstanceRef { + fn deref_mut(&mut self) -> &mut T { + unsafe { &mut *self.inner } + } +} + +impl AsRef for OneShotInstanceRef { + fn as_ref(&self) -> &T { + unsafe { &*self.inner } + } +} diff --git a/crates/rspack_napi/src/lib.rs b/crates/rspack_napi/src/lib.rs index 0f3a5b53f6e8..4bedbab602f4 100644 --- a/crates/rspack_napi/src/lib.rs +++ b/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::*;