1use std::{
2 any::Any,
3 cell::RefCell,
4 future::Future,
5 pin::Pin,
6 sync::{
7 LazyLock, RwLock,
8 atomic::{AtomicUsize, Ordering},
9 },
10};
11
12use napi::{
13 CleanupEnvHook, Env, Error, JsValue, Result, Status,
14 bindgen_prelude::{PromiseRaw, ToNapiValue},
15};
16
17static RUNTIME: LazyLock<RwLock<Option<tokio::runtime::Runtime>>> =
18 LazyLock::new(|| RwLock::new(None));
19static ACTIVE_ENVS: AtomicUsize = AtomicUsize::new(0);
20
21type BoxedUnitFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
22
23thread_local! {
24 static RUNTIME_CLEANUP_HOOK: RefCell<Option<CleanupEnvHook<()>>> = Default::default();
25}
26
27pub fn within_runtime_if_available<F: FnOnce() -> T, T>(f: F) -> T {
28 f()
29}
30
31pub fn ensure_runtime(env: &Env) -> Result<()> {
32 start_runtime();
33 register_env_cleanup(env)
34}
35
36pub fn spawn<F>(future: F)
37where
38 F: Future + Send + 'static,
39 F::Output: Send + 'static,
40{
41 std::mem::drop(spawn_inner(future));
42}
43
44pub fn block_on<F: Future>(future: F) -> F::Output {
45 with_runtime(|runtime| runtime.block_on(future))
46}
47
48pub fn promise_from_future<'env, T, F>(env: &'env Env, future: F) -> Result<PromiseRaw<'env, T>>
49where
50 T: 'static + Send + ToNapiValue,
51 F: 'static + Send + Future<Output = Result<T>>,
52{
53 ensure_runtime(env)?;
54
55 let (deferred, promise) = env.create_deferred()?;
56 let promise = PromiseRaw::new(env.raw(), promise.raw());
57 let deferred_for_panic = deferred.clone();
58
59 let task: BoxedUnitFuture = Box::pin(async move {
60 match future.await {
61 Ok(value) => deferred.resolve(|_| Ok(value)),
62 Err(error) => deferred.reject(error),
63 }
64 });
65 let handle = spawn_inner(task);
66
67 let monitor: BoxedUnitFuture = Box::pin(async move {
68 if let Err(error) = handle.await {
69 deferred_for_panic.reject(join_error_to_napi_error(error));
70 }
71 });
72 spawn_inner(monitor);
73
74 Ok(promise)
75}
76
77pub fn panic_to_napi_error(payload: Box<dyn Any + Send + 'static>) -> Error {
78 Error::new(Status::GenericFailure, panic_message(payload))
79}
80
81fn join_error_to_napi_error(error: tokio::task::JoinError) -> Error {
82 if error.is_panic() {
83 panic_to_napi_error(error.into_panic())
84 } else {
85 Error::new(Status::GenericFailure, "Async task was cancelled")
86 }
87}
88
89fn panic_message(payload: Box<dyn Any + Send + 'static>) -> String {
90 if let Some(message) = payload.downcast_ref::<&str>() {
91 (*message).to_string()
92 } else if let Some(message) = payload.downcast_ref::<String>() {
93 message.clone()
94 } else {
95 "Panic in async function".to_string()
96 }
97}
98
99fn register_env_cleanup(env: &Env) -> Result<()> {
100 RUNTIME_CLEANUP_HOOK.with(|cleanup_hook| {
101 let mut cleanup_hook = cleanup_hook.borrow_mut();
102 if cleanup_hook.is_some() {
103 return Ok(());
104 }
105
106 ACTIVE_ENVS.fetch_add(1, Ordering::SeqCst);
107 match env.add_env_cleanup_hook((), |_| {
108 RUNTIME_CLEANUP_HOOK.with_borrow_mut(|cleanup_hook| *cleanup_hook = None);
109 if ACTIVE_ENVS.fetch_sub(1, Ordering::SeqCst) == 1 {
110 shutdown_runtime();
111 }
112 }) {
113 Ok(hook) => {
114 *cleanup_hook = Some(hook);
115 Ok(())
116 }
117 Err(error) => {
118 ACTIVE_ENVS.fetch_sub(1, Ordering::SeqCst);
119 Err(error)
120 }
121 }
122 })
123}
124
125fn spawn_inner<F>(future: F) -> tokio::task::JoinHandle<F::Output>
126where
127 F: Future + Send + 'static,
128 F::Output: Send + 'static,
129{
130 with_runtime(|runtime| runtime.spawn(future))
131}
132
133fn with_runtime<R>(f: impl FnOnce(&tokio::runtime::Handle) -> R) -> R {
134 start_runtime();
135 let handle = RUNTIME
138 .read()
139 .expect("Read tokio runtime failed")
140 .as_ref()
141 .expect("Access tokio runtime failed after initialization")
142 .handle()
143 .clone();
144 f(&handle)
145}
146
147fn start_runtime() {
148 let mut runtime = RUNTIME.write().expect("Write tokio runtime failed");
149 if runtime.is_none() {
150 *runtime = Some(create_runtime());
151 }
152}
153
154fn shutdown_runtime() {
155 if let Some(runtime) = RUNTIME.write().expect("Write tokio runtime failed").take() {
156 runtime.shutdown_background();
157 }
158}
159
160fn create_runtime() -> tokio::runtime::Runtime {
161 let mut builder = tokio::runtime::Builder::new_multi_thread();
162 builder
163 .max_blocking_threads(blocking_threads())
164 .thread_name_fn(|| {
165 static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
166 let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
167 format!("tokio-{id}")
168 })
169 .enable_all()
170 .build()
171 .expect("Create tokio runtime failed")
172}
173
174fn blocking_threads() -> usize {
175 const ENV_BLOCKING_THREADS: &str = "RSPACK_BLOCKING_THREADS";
176
177 std::env::var(ENV_BLOCKING_THREADS)
178 .ok()
179 .and_then(|v| v.parse::<usize>().ok())
180 .unwrap_or(default_blocking_threads())
181}
182
183fn default_blocking_threads() -> usize {
184 #[cfg(target_family = "wasm")]
185 {
186 1
187 }
188
189 #[cfg(not(target_family = "wasm"))]
190 {
191 4
193 }
194}