fix(napi): add back custom gc for Send Buffer (#1399)

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) },

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) };
       }

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<RegisteredClassesMap> = Lazy::new(Default::default);
 static FN_REGISTER_MAP: Lazy<FnRegisterMap> = Lazy::new(Default::default);
+#[cfg(feature = "napi4")]
+pub(crate) static CUSTOM_GC_TSFN: AtomicPtr<sys::napi_threadsafe_function__> =
+  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<ThreadId> =
+  once_cell::sync::OnceCell::new();
 
 type RegisteredClasses =
   PersistedPerInstanceHashMap</* export name */ String, /* constructor */ sys::napi_ref>;
@@ -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"
+  );
+}