From e3a403b441da344cac0797fb638cdd039abd4a82 Mon Sep 17 00:00:00 2001
From: LongYinan <lynweklm@gmail.com>
Date: Mon, 15 Apr 2024 23:58:28 +0800
Subject: [PATCH] feat(napi): allow user defined tokio runtime (#2040)

---
 crates/napi/src/tokio_runtime.rs | 30 ++++++++++++++++++++++++++++--
 examples/napi/src/lib.rs         | 15 +++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/crates/napi/src/tokio_runtime.rs b/crates/napi/src/tokio_runtime.rs
index 3074ba70..b612ddb4 100644
--- a/crates/napi/src/tokio_runtime.rs
+++ b/crates/napi/src/tokio_runtime.rs
@@ -1,6 +1,6 @@
 use std::{future::Future, marker::PhantomData, sync::RwLock};
 
-use once_cell::sync::Lazy;
+use once_cell::sync::{Lazy, OnceCell};
 use tokio::runtime::Runtime;
 
 use crate::{sys, JsDeferred, JsUnknown, NapiValue, Result};
@@ -21,7 +21,33 @@ fn create_runtime() -> Option<Runtime> {
   }
 }
 
-pub(crate) static RT: Lazy<RwLock<Option<Runtime>>> = Lazy::new(|| RwLock::new(create_runtime()));
+pub(crate) static RT: Lazy<RwLock<Option<Runtime>>> = Lazy::new(|| {
+  if let Some(user_defined_rt) = unsafe { USER_DEFINED_RT.take() } {
+    RwLock::new(user_defined_rt)
+  } else {
+    RwLock::new(create_runtime())
+  }
+});
+
+static mut USER_DEFINED_RT: OnceCell<Option<Runtime>> = OnceCell::new();
+
+/// Create a custom Tokio runtime used by the NAPI-RS.
+/// You can control the tokio runtime configuration by yourself.
+/// ### Example
+/// ```no_run
+/// use tokio::runtime::Builder;
+/// use napi::create_custom_tokio_runtime;
+///
+/// #[napi::module_init]
+/// fn init() {
+///    let rt = Builder::new_multi_thread().enable_all().thread_stack_size(32 * 1024 * 1024).build().unwrap();
+///    create_custom_tokio_runtime(rt);
+/// }
+pub fn create_custom_tokio_runtime(rt: Runtime) {
+  unsafe {
+    USER_DEFINED_RT.get_or_init(move || Some(rt));
+  }
+}
 
 #[cfg(not(any(target_os = "macos", target_family = "wasm")))]
 static RT_REFERENCE_COUNT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
diff --git a/examples/napi/src/lib.rs b/examples/napi/src/lib.rs
index 4d8f0133..2cae0595 100644
--- a/examples/napi/src/lib.rs
+++ b/examples/napi/src/lib.rs
@@ -5,6 +5,8 @@
 #![allow(clippy::new_without_default)]
 #![allow(deprecated)]
 
+use napi::bindgen_prelude::create_custom_tokio_runtime;
+
 #[macro_use]
 extern crate napi_derive;
 #[macro_use]
@@ -14,6 +16,19 @@ extern crate serde_derive;
 #[global_allocator]
 static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;
 
+#[cfg(not(target_family = "wasm"))]
+#[napi::module_init]
+fn init() {
+  let rt = tokio::runtime::Builder::new_multi_thread()
+    .enable_all()
+    .on_thread_start(|| {
+      println!("tokio thread started");
+    })
+    .build()
+    .unwrap();
+  create_custom_tokio_runtime(rt);
+}
+
 #[napi]
 /// This is a const
 pub const DEFAULT_COST: u32 = 12;