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 { } } -pub(crate) static RT: Lazy>> = Lazy::new(|| RwLock::new(create_runtime())); +pub(crate) static RT: Lazy>> = 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> = 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;