diff --git a/crates/napi/src/bindgen_runtime/js_values/arraybuffer.rs b/crates/napi/src/bindgen_runtime/js_values/arraybuffer.rs index 663360a35a..c7363b4c84 100644 --- a/crates/napi/src/bindgen_runtime/js_values/arraybuffer.rs +++ b/crates/napi/src/bindgen_runtime/js_values/arraybuffer.rs @@ -7,6 +7,8 @@ use std::sync::{ Arc, }; +#[cfg(feature = "napi4")] +use crate::bindgen_prelude::{CUSTOM_GC_TSFN, CUSTOM_GC_TSFN_CLOSED, MAIN_THREAD_ID}; pub use crate::js_values::TypedArrayType; use crate::{check_status, sys, Error, Result, Status}; @@ -64,6 +66,31 @@ macro_rules! impl_typed_array { fn drop(&mut self) { if Arc::strong_count(&self.drop_in_vm) == 1 { if let Some((ref_, env)) = self.raw { + #[cfg(feature = "napi4")] + { + if CUSTOM_GC_TSFN_CLOSED.load(std::sync::atomic::Ordering::SeqCst) { + return; + } + if !MAIN_THREAD_ID + .get() + .map(|id| &std::thread::current().id() == id) + .unwrap_or(false) + { + let status = unsafe { + sys::napi_call_threadsafe_function( + CUSTOM_GC_TSFN.load(std::sync::atomic::Ordering::SeqCst), + ref_ as *mut c_void, + 1, + ) + }; + assert!( + status == sys::Status::napi_ok, + "Call custom GC in ArrayBuffer::drop failed {:?}", + Status::from(status) + ); + return; + } + } crate::check_status_or_throw!( env, unsafe { sys::napi_reference_unref(env, ref_, &mut 0) }, diff --git a/crates/napi/src/bindgen_runtime/js_values/buffer.rs b/crates/napi/src/bindgen_runtime/js_values/buffer.rs index c3f986772d..212fe99cbb 100644 --- a/crates/napi/src/bindgen_runtime/js_values/buffer.rs +++ b/crates/napi/src/bindgen_runtime/js_values/buffer.rs @@ -9,6 +9,8 @@ use std::sync::Arc; #[cfg(all(debug_assertions, not(windows)))] use std::sync::Mutex; +#[cfg(feature = "napi4")] +use crate::bindgen_prelude::{CUSTOM_GC_TSFN, CUSTOM_GC_TSFN_CLOSED, MAIN_THREAD_ID}; use crate::{bindgen_prelude::*, check_status, sys, Result, ValueType}; #[cfg(all(debug_assertions, not(windows)))] @@ -43,6 +45,30 @@ impl Drop for Buffer { unsafe { sys::napi_delete_reference(env, ref_) }, "Failed to delete Buffer reference in drop" ); + #[cfg(feature = "napi4")] + { + if CUSTOM_GC_TSFN_CLOSED.load(std::sync::atomic::Ordering::SeqCst) { + return; + } + if !MAIN_THREAD_ID + .get() + .map(|id| &std::thread::current().id() == id) + .unwrap_or(false) + { + let status = unsafe { + sys::napi_call_threadsafe_function( + CUSTOM_GC_TSFN.load(std::sync::atomic::Ordering::SeqCst), + ref_ as *mut c_void, + 1, + ) + }; + assert!( + status == sys::Status::napi_ok, + "Call custom GC in ArrayBuffer::drop failed {:?}", + Status::from(status) + ); + } + } } else { unsafe { Vec::from_raw_parts(self.inner.as_ptr(), self.len, self.capacity) }; } diff --git a/crates/napi/src/bindgen_runtime/module_register.rs b/crates/napi/src/bindgen_runtime/module_register.rs index 429ebd67d4..306bd23994 100644 --- a/crates/napi/src/bindgen_runtime/module_register.rs +++ b/crates/napi/src/bindgen_runtime/module_register.rs @@ -117,6 +117,16 @@ static IS_FIRST_MODULE: AtomicBool = AtomicBool::new(true); static FIRST_MODULE_REGISTERED: AtomicBool = AtomicBool::new(false); static REGISTERED_CLASSES: Lazy = Lazy::new(Default::default); static FN_REGISTER_MAP: Lazy = Lazy::new(Default::default); +#[cfg(feature = "napi4")] +pub(crate) static CUSTOM_GC_TSFN: AtomicPtr = + AtomicPtr::new(ptr::null_mut()); +#[cfg(feature = "napi4")] +// CustomGC ThreadsafeFunction may be deleted during the process exit. +// And there may still some Buffer alive after that. +pub(crate) static CUSTOM_GC_TSFN_CLOSED: AtomicBool = AtomicBool::new(false); +#[cfg(feature = "napi4")] +pub(crate) static MAIN_THREAD_ID: once_cell::sync::OnceCell = + once_cell::sync::OnceCell::new(); type RegisteredClasses = PersistedPerInstanceHashMap; @@ -479,6 +489,8 @@ unsafe extern "C" fn napi_register_module_v1( ) }; } + #[cfg(feature = "napi4")] + create_custom_gc(env); FIRST_MODULE_REGISTERED.store(true, Ordering::SeqCst); exports } @@ -500,3 +512,101 @@ pub(crate) unsafe extern "C" fn noop( } ptr::null_mut() } + +#[cfg(feature = "napi4")] +fn create_custom_gc(env: sys::napi_env) { + use std::os::raw::c_char; + + let mut custom_gc_fn = ptr::null_mut(); + check_status_or_throw!( + env, + unsafe { + sys::napi_create_function( + env, + "custom_gc".as_ptr() as *const c_char, + 9, + Some(empty), + ptr::null_mut(), + &mut custom_gc_fn, + ) + }, + "Create Custom GC Function in napi_register_module_v1 failed" + ); + let mut async_resource_name = ptr::null_mut(); + check_status_or_throw!( + env, + unsafe { + sys::napi_create_string_utf8( + env, + "CustomGC".as_ptr() as *const c_char, + 8, + &mut async_resource_name, + ) + }, + "Create async resource string in napi_register_module_v1 napi_register_module_v1" + ); + let mut custom_gc_tsfn = ptr::null_mut(); + check_status_or_throw!( + env, + unsafe { + sys::napi_create_threadsafe_function( + env, + custom_gc_fn, + ptr::null_mut(), + async_resource_name, + 0, + 1, + ptr::null_mut(), + Some(custom_gc_finalize), + ptr::null_mut(), + Some(custom_gc), + &mut custom_gc_tsfn, + ) + }, + "Create Custom GC ThreadsafeFunction in napi_register_module_v1 failed" + ); + check_status_or_throw!( + env, + unsafe { sys::napi_unref_threadsafe_function(env, custom_gc_tsfn) }, + "Unref Custom GC ThreadsafeFunction in napi_register_module_v1 failed" + ); + CUSTOM_GC_TSFN.store(custom_gc_tsfn, Ordering::SeqCst); + MAIN_THREAD_ID.get_or_init(|| std::thread::current().id()); +} + +#[cfg(feature = "napi4")] +#[allow(unused)] +unsafe extern "C" fn empty(env: sys::napi_env, info: sys::napi_callback_info) -> sys::napi_value { + ptr::null_mut() +} + +#[cfg(feature = "napi4")] +#[allow(unused)] +unsafe extern "C" fn custom_gc_finalize( + env: sys::napi_env, + finalize_data: *mut std::ffi::c_void, + finalize_hint: *mut std::ffi::c_void, +) { + CUSTOM_GC_TSFN_CLOSED.store(true, Ordering::SeqCst); +} + +#[cfg(feature = "napi4")] +// recycle the ArrayBuffer/Buffer Reference if the ArrayBuffer/Buffer is not dropped on the main thread +extern "C" fn custom_gc( + env: sys::napi_env, + _js_callback: sys::napi_value, + _context: *mut std::ffi::c_void, + data: *mut std::ffi::c_void, +) { + let mut ref_count = 0; + check_status_or_throw!( + env, + unsafe { sys::napi_reference_unref(env, data as sys::napi_ref, &mut ref_count) }, + "Failed to unref Buffer reference in Custom GC" + ); + check_status_or_throw!( + env, + unsafe { sys::napi_delete_reference(env, data as sys::napi_ref) }, + "Failed to delete Buffer reference in Custom GC" + ); +}