normalize tokio runtime

crates/napi/src/env.rs

@@ -5,11 +5,6 @@ use std::mem;
 use std::os::raw::{c_char, c_void};
 use std::ptr;
 
-#[cfg(all(feature = "tokio_rt", feature = "napi4"))]
-use once_cell::sync::Lazy;
-#[cfg(all(feature = "tokio_rt", feature = "napi4"))]
-use tokio::{runtime::Handle, sync::mpsc};
-
 use crate::{
   async_work::{self, AsyncWorkPromise},
   check_status,
@@ -25,8 +20,6 @@ use crate::async_cleanup_hook::AsyncCleanupHook;
 use crate::cleanup_env::{CleanupEnvHook, CleanupEnvHookData};
 #[cfg(all(feature = "serde-json"))]
 use crate::js_values::{De, Ser};
-#[cfg(all(feature = "tokio_rt", feature = "napi4"))]
-use crate::promise;
 #[cfg(feature = "napi4")]
 use crate::threadsafe_function::{ThreadSafeCallContext, ThreadsafeFunction};
 #[cfg(all(feature = "serde-json"))]
@@ -38,35 +31,6 @@ use std::future::Future;
 
 pub type Callback = extern "C" fn(sys::napi_env, sys::napi_callback_info) -> sys::napi_value;
 
-#[cfg(all(feature = "tokio_rt", feature = "napi4"))]
-static RT: Lazy<(Handle, mpsc::Sender<()>)> = Lazy::new(|| {
-  let rt = tokio::runtime::Runtime::new();
-  let (tx, mut rx) = mpsc::channel::<()>(1);
-  rt.map(|rt| {
-    let h = rt.handle();
-    let handle = h.clone();
-    handle.spawn(async move {
-      if rx.recv().await.is_some() {
-        rt.shutdown_background();
-      }
-    });
-
-    (handle, tx)
-  })
-  .expect("Create tokio runtime failed")
-});
-
-#[doc(hidden)]
-#[cfg(all(feature = "tokio_rt", feature = "napi4"))]
-#[inline(never)]
-pub extern "C" fn shutdown_tokio_rt(_arg: *mut c_void) {
-  let sender = &RT.1;
-  sender
-    .clone()
-    .try_send(())
-    .expect("Shutdown tokio runtime failed");
-}
-
 #[derive(Clone, Copy)]
 /// `Env` is used to represent a context that the underlying N-API implementation can use to persist VM-specific state.
 ///
@@ -1062,19 +1026,13 @@ impl Env {
     fut: F,
     resolver: R,
   ) -> Result<JsObject> {
-    let handle = &RT.0;
+    use crate::tokio_runtime;
 
-    let mut raw_promise = ptr::null_mut();
+    let promise = tokio_runtime::execute_tokio_future(self.0, fut, |env, val| unsafe {
-    let mut raw_deferred = ptr::null_mut();
+      resolver(&mut Env::from_raw(env), val).map(|v| v.raw())
-    check_status!(unsafe {
-      sys::napi_create_promise(self.0, &mut raw_deferred, &mut raw_promise)
     })?;
 
-    let raw_env = self.0;
+    Ok(unsafe { JsObject::from_raw_unchecked(self.0, promise) })
-    let future_promise = promise::FuturePromise::create(raw_env, raw_deferred, resolver)?;
-    let future_to_resolve = promise::resolve_from_future(future_promise.start()?, fut);
-    handle.spawn(future_to_resolve);
-    Ok(unsafe { JsObject::from_raw_unchecked(self.0, raw_promise) })
   }
 
   /// This API does not observe leap seconds; they are ignored, as ECMAScript aligns with POSIX time specification.

crates/napi/src/lib.rs

@@ -93,6 +93,8 @@ mod js_values;
 mod promise;
 mod status;
 mod task;
+#[cfg(all(feature = "tokio_rt", feature = "napi4"))]
+mod tokio_runtime;
 mod value_type;
 #[cfg(feature = "napi3")]
 pub use cleanup_env::CleanupEnvHook;
@@ -113,6 +115,8 @@ pub use error::*;
 pub use js_values::*;
 pub use status::Status;
 pub use task::Task;
+#[cfg(all(feature = "tokio_rt", feature = "napi4"))]
+pub use tokio_runtime::shutdown_tokio_rt;
 pub use value_type::*;
 pub use version::NodeVersion;
 #[cfg(feature = "serde-json")]

crates/napi/src/promise.rs

@@ -3,22 +3,27 @@ use std::marker::PhantomData;
 use std::os::raw::{c_char, c_void};
 use std::ptr;
 
-use crate::{check_status, sys, Env, JsError, NapiRaw, Result};
+use crate::{check_status, sys, JsError, Result};
 
-pub struct FuturePromise<T, V: NapiRaw, F: FnOnce(&mut Env, T) -> Result<V>> {
+pub struct FuturePromise<Data, Resolver: FnOnce(sys::napi_env, Data) -> Result<sys::napi_value>> {
   deferred: sys::napi_deferred,
   env: sys::napi_env,
   tsfn: sys::napi_threadsafe_function,
   async_resource_name: sys::napi_value,
-  resolver: F,
+  resolver: Resolver,
-  _data: PhantomData<T>,
+  _data: PhantomData<Data>,
-  _value: PhantomData<V>,
+  _value: PhantomData<sys::napi_value>,
 }
 
-unsafe impl<T, V: NapiRaw, F: FnOnce(&mut Env, T) -> Result<V>> Send for FuturePromise<T, V, F> {}
+unsafe impl<T, F: FnOnce(sys::napi_env, T) -> Result<sys::napi_value>> Send
+  for FuturePromise<T, F>
+{
+}
 
-impl<T, V: NapiRaw, F: FnOnce(&mut Env, T) -> Result<V>> FuturePromise<T, V, F> {
+impl<Data, Resolver: FnOnce(sys::napi_env, Data) -> Result<sys::napi_value>>
-  pub fn create(env: sys::napi_env, raw_deferred: sys::napi_deferred, resolver: F) -> Result<Self> {
+  FuturePromise<Data, Resolver>
+{
+  pub fn new(env: sys::napi_env, dererred: sys::napi_deferred, resolver: Resolver) -> Result<Self> {
     let mut async_resource_name = ptr::null_mut();
     let s = "napi_resolve_promise_from_future";
     check_status!(unsafe {
@@ -31,7 +36,7 @@ impl<T, V: NapiRaw, F: FnOnce(&mut Env, T) -> Result<V>> FuturePromise<T, V, F>
     })?;
 
     Ok(FuturePromise {
-      deferred: raw_deferred,
+      deferred: dererred,
       resolver,
       env,
       tsfn: ptr::null_mut(),
@@ -56,8 +61,8 @@ impl<T, V: NapiRaw, F: FnOnce(&mut Env, T) -> Result<V>> FuturePromise<T, V, F>
         1,
         ptr::null_mut(),
         None,
-        self_ref as *mut FuturePromise<T, V, F> as *mut c_void,
+        self_ref as *mut FuturePromise<Data, Resolver> as *mut c_void,
-        Some(call_js_cb::<T, V, F>),
+        Some(call_js_cb::<Data, Resolver>),
         &mut tsfn_value,
       )
     })?;
@@ -70,15 +75,15 @@ pub(crate) struct TSFNValue(sys::napi_threadsafe_function);
 
 unsafe impl Send for TSFNValue {}
 
-pub(crate) async fn resolve_from_future<T: Send, F: Future<Output = Result<T>>>(
+pub(crate) async fn resolve_from_future<Data: Send, Fut: Future<Output = Result<Data>>>(
   tsfn_value: TSFNValue,
-  fut: F,
+  fut: Fut,
 ) {
   let val = fut.await;
   check_status!(unsafe {
     sys::napi_call_threadsafe_function(
       tsfn_value.0,
-      Box::into_raw(Box::from(val)) as *mut T as *mut c_void,
+      Box::into_raw(Box::from(val)) as *mut Data as *mut c_void,
       sys::napi_threadsafe_function_call_mode::napi_tsfn_nonblocking,
     )
   })
@@ -92,26 +97,27 @@ pub(crate) async fn resolve_from_future<T: Send, F: Future<Output = Result<T>>>(
   .expect("Failed to release thread safe function");
 }
 
-unsafe extern "C" fn call_js_cb<T, V: NapiRaw, F: FnOnce(&mut Env, T) -> Result<V>>(
+unsafe extern "C" fn call_js_cb<
-  raw_env: sys::napi_env,
+  Data,
+  Resolver: FnOnce(sys::napi_env, Data) -> Result<sys::napi_value>,
+>(
+  env: sys::napi_env,
   _js_callback: sys::napi_value,
   context: *mut c_void,
   data: *mut c_void,
 ) {
-  let mut env = Env::from_raw(raw_env);
+  let future_promise = Box::from_raw(context as *mut FuturePromise<Data, Resolver>);
-  let future_promise = Box::from_raw(context as *mut FuturePromise<T, V, F>);
+  let value = Box::from_raw(data as *mut Result<Data>);
-  let value = Box::from_raw(data as *mut Result<T>);
   let resolver = future_promise.resolver;
   let deferred = future_promise.deferred;
-  let js_value_to_resolve = value.and_then(move |v| (resolver)(&mut env, v));
+  let js_value_to_resolve = value.and_then(move |v| (resolver)(env, v));
   match js_value_to_resolve {
     Ok(v) => {
-      let status = sys::napi_resolve_deferred(raw_env, deferred, v.raw());
+      let status = sys::napi_resolve_deferred(env, deferred, v);
       debug_assert!(status == sys::Status::napi_ok, "Resolve promise failed");
     }
     Err(e) => {
-      let status =
+      let status = sys::napi_reject_deferred(env, deferred, JsError::from(e).into_value(env));
-        sys::napi_reject_deferred(raw_env, deferred, JsError::from(e).into_value(raw_env));
       debug_assert!(status == sys::Status::napi_ok, "Reject promise failed");
     }
   };

crates/napi/src/tokio_runtime.rs

@@ -0,0 +1,61 @@
+use std::{ffi::c_void, future::Future, ptr};
+
+use crate::{check_status, promise, sys, Result};
+use once_cell::sync::Lazy;
+use tokio::{runtime::Handle, sync::mpsc};
+
+static RT: Lazy<(Handle, mpsc::Sender<()>)> = Lazy::new(|| {
+  let runtime = tokio::runtime::Runtime::new();
+  let (sender, mut receiver) = mpsc::channel::<()>(1);
+  runtime
+    .map(|rt| {
+      let h = rt.handle();
+      let handle = h.clone();
+      handle.spawn(async move {
+        if receiver.recv().await.is_some() {
+          rt.shutdown_background();
+        }
+      });
+
+      (handle, sender)
+    })
+    .expect("Create tokio runtime failed")
+});
+
+#[doc(hidden)]
+#[inline(never)]
+pub extern "C" fn shutdown_tokio_rt(_arg: *mut c_void) {
+  let sender = &RT.1;
+  sender
+    .clone()
+    .try_send(())
+    .expect("Shutdown tokio runtime failed");
+}
+
+pub fn spawn<F>(fut: F)
+where
+  F: 'static + Send + Future<Output = ()>,
+{
+  RT.0.spawn(fut);
+}
+
+pub fn execute_tokio_future<
+  Data: 'static + Send,
+  Fut: 'static + Send + Future<Output = Result<Data>>,
+  Resolver: 'static + Send + Sync + FnOnce(sys::napi_env, Data) -> Result<sys::napi_value>,
+>(
+  env: sys::napi_env,
+  fut: Fut,
+  resolver: Resolver,
+) -> Result<sys::napi_value> {
+  let mut promise = ptr::null_mut();
+  let mut deferred = ptr::null_mut();
+
+  check_status!(unsafe { sys::napi_create_promise(env, &mut deferred, &mut promise) })?;
+
+  let future_promise = promise::FuturePromise::new(env, deferred, resolver)?;
+  let future_to_resolve = promise::resolve_from_future(future_promise.start()?, fut);
+  spawn(future_to_resolve);
+
+  Ok(promise)
+}