From 968d9e10b1d621eae69865c3b16ff403d97aaf76 Mon Sep 17 00:00:00 2001 From: LongYinan Date: Fri, 16 Dec 2022 20:07:22 +0800 Subject: [PATCH] chore(napi): fix tokio destroy logic --- .../src/bindgen_runtime/module_register.rs | 86 ++++++------------- crates/napi/src/tokio_runtime.rs | 32 +++++-- 2 files changed, 49 insertions(+), 69 deletions(-) diff --git a/crates/napi/src/bindgen_runtime/module_register.rs b/crates/napi/src/bindgen_runtime/module_register.rs index a1ca0676..429ebd67 100644 --- a/crates/napi/src/bindgen_runtime/module_register.rs +++ b/crates/napi/src/bindgen_runtime/module_register.rs @@ -1,5 +1,5 @@ use std::collections::{HashMap, HashSet}; -use std::ffi::{c_void, CStr}; +use std::ffi::CStr; use std::ptr; use std::sync::atomic::AtomicUsize; use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering}; @@ -21,17 +21,6 @@ struct PersistedPerInstanceVec { length: AtomicUsize, } -impl Drop for PersistedPerInstanceVec { - fn drop(&mut self) { - let length = self.length.load(Ordering::Relaxed); - if length == 0 { - return; - } - let inner = self.inner.load(Ordering::Relaxed); - unsafe { Vec::from_raw_parts(inner, length, length) }; - } -} - impl Default for PersistedPerInstanceVec { fn default() -> Self { let mut vec: Vec = Vec::with_capacity(1); @@ -86,12 +75,6 @@ unsafe impl Sync for PersistedPerInstanceVec {} pub(crate) struct PersistedPerInstanceHashMap(*mut HashMap); -impl Drop for PersistedPerInstanceHashMap { - fn drop(&mut self) { - unsafe { Box::from_raw(self.0) }; - } -} - impl PersistedPerInstanceHashMap { pub(crate) fn from_hashmap(hashmap: HashMap) -> Self { Self(Box::into_raw(Box::new(hashmap))) @@ -130,32 +113,11 @@ type RegisteredClassesMap = PersistedPerInstanceHashMap = Lazy::new(Default::default); static MODULE_CLASS_PROPERTIES: Lazy = Lazy::new(Default::default); -static REGISTERED: AtomicBool = AtomicBool::new(false); +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); -#[ctor::dtor] -fn destroy() { - { - let ptr = MODULE_REGISTER_CALLBACK.inner.load(Ordering::Relaxed); - let len = MODULE_REGISTER_CALLBACK.length.load(Ordering::Relaxed); - unsafe { Vec::from_raw_parts(ptr, len, len) }; - } - { - unsafe { Box::from_raw(MODULE_CLASS_PROPERTIES.0) }; - } - { - unsafe { Box::from_raw(FN_REGISTER_MAP.0) }; - } -} - -#[inline] -fn wait_first_thread_registered() { - while !REGISTERED.load(Ordering::SeqCst) { - std::hint::spin_loop(); - } -} - type RegisteredClasses = PersistedPerInstanceHashMap; @@ -164,9 +126,15 @@ type RegisteredClasses = static MODULE_EXPORTS: Lazy> = Lazy::new(Default::default); +#[inline] +fn wait_first_thread_registered() { + while !FIRST_MODULE_REGISTERED.load(Ordering::SeqCst) { + std::hint::spin_loop(); + } +} + #[doc(hidden)] pub fn get_class_constructor(js_name: &'static str) -> Option { - wait_first_thread_registered(); let current_id = std::thread::current().id(); REGISTERED_CLASSES.borrow_mut(|map| { map @@ -236,7 +204,6 @@ pub fn register_class( /// ``` /// pub fn get_js_function(env: &Env, raw_fn: ExportRegisterCallback) -> Result { - wait_first_thread_registered(); FN_REGISTER_MAP.borrow_mut(|inner| { inner .get(&raw_fn) @@ -290,7 +257,6 @@ pub fn get_js_function(env: &Env, raw_fn: ExportRegisterCallback) -> Result Result { - wait_first_thread_registered(); FN_REGISTER_MAP.borrow_mut(|inner| { inner .get(&raw_fn) @@ -317,6 +283,11 @@ unsafe extern "C" fn napi_register_module_v1( env: sys::napi_env, exports: sys::napi_value, ) -> sys::napi_value { + if IS_FIRST_MODULE.load(Ordering::SeqCst) { + IS_FIRST_MODULE.store(false, Ordering::SeqCst); + } else { + wait_first_thread_registered(); + } crate::__private::___CALL_FROM_FACTORY.get_or_default(); let mut exports_objects: HashSet = HashSet::default(); MODULE_REGISTER_CALLBACK.borrow_mut(|inner| { @@ -393,8 +364,7 @@ unsafe extern "C" fn napi_register_module_v1( }) }); - let mut registered_classes = - HashMap::with_capacity(MODULE_CLASS_PROPERTIES.borrow_mut(|inner| inner.len())); + let mut registered_classes = HashMap::new(); MODULE_CLASS_PROPERTIES.borrow_mut(|inner| { inner.iter().for_each(|(rust_name, js_mods)| { @@ -498,29 +468,21 @@ unsafe extern "C" fn napi_register_module_v1( }) }); - #[cfg(feature = "napi3")] + #[cfg(all(windows, feature = "napi4", feature = "tokio_rt"))] { + crate::tokio_runtime::RT_REFERENCE_COUNT.fetch_add(1, Ordering::SeqCst); unsafe { - sys::napi_add_env_cleanup_hook(env, Some(remove_registered_classes), env as *mut c_void) + sys::napi_add_env_cleanup_hook( + env, + Some(crate::tokio_runtime::drop_runtime), + ptr::null_mut(), + ) }; } - REGISTERED.store(true, Ordering::SeqCst); + FIRST_MODULE_REGISTERED.store(true, Ordering::SeqCst); exports } -unsafe extern "C" fn remove_registered_classes(env: *mut c_void) { - let env = env as sys::napi_env; - if let Some(registered_classes) = - REGISTERED_CLASSES.borrow_mut(|map| map.remove(&std::thread::current().id())) - { - registered_classes.borrow_mut(|map| { - map.iter().for_each(|(_, v)| { - unsafe { sys::napi_delete_reference(env, *v) }; - }) - }); - } -} - pub(crate) unsafe extern "C" fn noop( env: sys::napi_env, _info: sys::napi_callback_info, diff --git a/crates/napi/src/tokio_runtime.rs b/crates/napi/src/tokio_runtime.rs index a3e49103..ef923ac5 100644 --- a/crates/napi/src/tokio_runtime.rs +++ b/crates/napi/src/tokio_runtime.rs @@ -6,11 +6,29 @@ use tokio::runtime::Runtime; use crate::{check_status, sys, JsDeferred, JsUnknown, NapiValue, Result}; -pub(crate) static RT: Lazy = - Lazy::new(|| tokio::runtime::Runtime::new().expect("Create tokio runtime failed")); +pub(crate) static mut RT: Lazy> = Lazy::new(|| { + let runtime = tokio::runtime::Runtime::new().expect("Create tokio runtime failed"); + Some(runtime) +}); -pub fn runtime() -> &'static Runtime { - &RT +#[cfg(windows)] +pub(crate) static RT_REFERENCE_COUNT: std::sync::atomic::AtomicUsize = + std::sync::atomic::AtomicUsize::new(0); + +#[cfg(windows)] +pub(crate) unsafe extern "C" fn drop_runtime(arg: *mut std::ffi::c_void) { + use std::sync::atomic::Ordering; + + if RT_REFERENCE_COUNT.fetch_sub(1, Ordering::SeqCst) == 1 { + if let Some(rt) = Lazy::get_mut(unsafe { &mut RT }) { + rt.take(); + } + } + + unsafe { + let env: sys::napi_env = arg as *mut sys::napi_env__; + sys::napi_remove_env_cleanup_hook(env, Some(drop_runtime), arg); + } } /// Spawns a future onto the Tokio runtime. @@ -21,7 +39,7 @@ pub fn spawn(fut: F) -> tokio::task::JoinHandle where F: 'static + Send + Future, { - RT.spawn(fut) + unsafe { RT.as_ref() }.unwrap().spawn(fut) } /// Runs a future to completion @@ -31,7 +49,7 @@ pub fn block_on(fut: F) -> F::Output where F: 'static + Send + Future, { - RT.block_on(fut) + unsafe { RT.as_ref() }.unwrap().block_on(fut) } // This function's signature must be kept in sync with the one in lib.rs, otherwise napi @@ -41,7 +59,7 @@ where /// then call the provided closure. Otherwise it will just call the provided closure. #[inline] pub fn within_runtime_if_available T, T>(f: F) -> T { - let _rt_guard = RT.enter(); + let _rt_guard = unsafe { RT.as_ref() }.unwrap().enter(); f() }