fix(napi): re-create async runtime (#2519)

Author: Brooooooklyn

crates/napi/src/bindgen_runtime/module_register.rs

@@ -480,7 +480,7 @@ pub unsafe extern "C" fn napi_register_module_v1(
     create_custom_gc(env, current_thread_id);
     #[cfg(feature = "tokio_rt")]
     {
-      crate::tokio_runtime::ensure_runtime();
+      crate::tokio_runtime::start_async_runtime();
     }
   }
   FIRST_MODULE_REGISTERED.store(true, Ordering::SeqCst);
@@ -570,7 +570,7 @@ unsafe extern "C" fn thread_cleanup(
   if MODULE_COUNT.fetch_sub(1, Ordering::Relaxed) == 1 {
     #[cfg(all(feature = "tokio_rt", feature = "napi4"))]
     {
-      crate::tokio_runtime::shutdown_tokio_runtime();
+      crate::tokio_runtime::shutdown_async_runtime();
     }
     crate::bindgen_runtime::REFERENCE_MAP.borrow_mut(|m| m.clear());
     #[allow(clippy::needless_return)]

crates/napi/src/tokio_runtime.rs

@@ -10,6 +10,12 @@ use crate::{JsDeferred, JsUnknown, NapiValue};
 
 #[cfg(not(feature = "noop"))]
 fn create_runtime() -> Runtime {
+  if let Some(user_defined_rt) = USER_DEFINED_RT
+    .get()
+    .and_then(|rt| rt.write().ok().and_then(|mut rt| rt.take()))
+  {
+    return user_defined_rt;
+  }
   #[cfg(any(
     all(target_family = "wasm", tokio_unstable),
     not(target_family = "wasm")
@@ -57,19 +63,25 @@ pub fn create_custom_tokio_runtime(rt: Runtime) {
 pub fn create_custom_tokio_runtime(_: Runtime) {}
 
 #[cfg(not(feature = "noop"))]
-/// Ensure that the Tokio runtime is initialized.
-/// In Node.js the Tokio runtime will be dropped when Node env exits.
+/// Start the async runtime (Currently is tokio).
+///
+/// In Node.js native targets the async runtime will be dropped when Node env exits.
 /// But in Electron renderer process, the Node env will exits and recreate when the window reloads.
-/// So we need to ensure that the Tokio runtime is initialized when the Node env is created.
-pub(crate) fn ensure_runtime() {
-  let mut rt = RT.write().unwrap();
-  if rt.is_none() {
-    *rt = Some(create_runtime());
+/// So we need to ensure that the async runtime is initialized when the Node env is created.
+///
+/// In wasm targets, the async runtime will not been shutdown automatically due to the limitation of the wasm runtime.
+/// So, you need to call `shutdown_async_runtime` function to manually shutdown the async runtime.
+/// In some scenarios, you may want to start the async runtime again like in tests.
+pub fn start_async_runtime() {
+  if let Ok(mut rt) = RT.write() {
+    if rt.is_none() {
+      *rt = Some(create_runtime());
+    }
   }
 }
 
 #[cfg(not(feature = "noop"))]
-pub fn shutdown_tokio_runtime() {
+pub fn shutdown_async_runtime() {
   if let Some(rt) = RT.write().ok().and_then(|mut rt| rt.take()) {
     rt.shutdown_background();
   }

examples/napi/src/lib.rs

@@ -44,7 +44,7 @@ pub const TYPE_SKIPPED_CONST: u32 = 12;
 pub fn shutdown_runtime() {
   #[cfg(all(target_family = "wasm", tokio_unstable))]
   {
-    napi::bindgen_prelude::shutdown_tokio_runtime();
+    napi::bindgen_prelude::shutdown_async_runtime();
   }
 }