use deno_core::futures::channel::mpsc;
use deno_core::V8CrossThreadTaskSpawner;
use deno_runtime::deno_napi::*;
use once_cell::sync::Lazy;
use std::mem::forget;
use std::ptr::NonNull;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
#[repr(transparent)]
pub struct SendPtr<T>(pub *const T);
unsafe impl<T> Send for SendPtr<T> {}
unsafe impl<T> Sync for SendPtr<T> {}
static TS_FN_ID_COUNTER: Lazy<AtomicUsize> = Lazy::new(|| AtomicUsize::new(0));
pub struct TsFn {
pub id: usize,
pub env: *mut Env,
pub maybe_func: Option<v8::Global<v8::Function>>,
pub maybe_call_js_cb: Option<napi_threadsafe_function_call_js>,
pub context: *mut c_void,
pub thread_counter: usize,
pub ref_counter: Arc<AtomicUsize>,
finalizer: Option<napi_finalize>,
finalizer_data: *mut c_void,
sender: V8CrossThreadTaskSpawner,
tsfn_sender: mpsc::UnboundedSender<ThreadSafeFunctionStatus>,
}
impl Drop for TsFn {
fn drop(&mut self) {
let env = unsafe { self.env.as_mut().unwrap() };
env.remove_threadsafe_function_ref_counter(self.id);
if let Some(finalizer) = self.finalizer {
unsafe {
(finalizer)(self.env as _, self.finalizer_data, ptr::null_mut());
}
}
}
}
impl TsFn {
pub fn acquire(&mut self) -> napi_status {
self.thread_counter += 1;
napi_ok
}
pub fn release(mut self) -> napi_status {
self.thread_counter -= 1;
if self.thread_counter == 0 {
if self
.tsfn_sender
.unbounded_send(ThreadSafeFunctionStatus::Dead)
.is_err()
{
return napi_generic_failure;
}
drop(self);
} else {
forget(self);
}
napi_ok
}
pub fn ref_(&mut self) -> napi_status {
self
.ref_counter
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
napi_ok
}
pub fn unref(&mut self) -> napi_status {
let _ = self.ref_counter.fetch_update(
std::sync::atomic::Ordering::SeqCst,
std::sync::atomic::Ordering::SeqCst,
|x| {
if x == 0 {
None
} else {
Some(x - 1)
}
},
);
napi_ok
}
pub fn call(&self, data: *mut c_void, is_blocking: bool) {
let js_func = self.maybe_func.clone();
let env = SendPtr(self.env);
let context = SendPtr(self.context);
let data = SendPtr(data);
#[inline(always)]
fn spawn(
sender: &V8CrossThreadTaskSpawner,
is_blocking: bool,
f: impl FnOnce(&mut v8::HandleScope) + Send + 'static,
) {
if is_blocking {
sender.spawn_blocking(f);
} else {
sender.spawn(f);
}
}
if let Some(call_js_cb) = self.maybe_call_js_cb {
if let Some(func) = js_func {
let func = SendPtr(func.into_raw().as_ptr());
#[inline(always)]
fn call(
scope: &mut v8::HandleScope,
call_js_cb: napi_threadsafe_function_call_js,
func: SendPtr<v8::Function>,
env: SendPtr<Env>,
context: SendPtr<c_void>,
data: SendPtr<c_void>,
) {
let func: v8::Global<v8::Function> = unsafe {
v8::Global::<v8::Function>::from_raw(
scope,
NonNull::new_unchecked(func.0 as _),
)
};
let func: v8::Local<v8::Value> =
func.open(scope).to_object(scope).unwrap().into();
unsafe {
call_js_cb(env.0 as _, func.into(), context.0 as _, data.0 as _)
}
}
spawn(&self.sender, is_blocking, move |scope| {
call(scope, call_js_cb, func, env, context, data);
});
} else {
#[inline(always)]
fn call(
call_js_cb: napi_threadsafe_function_call_js,
env: SendPtr<Env>,
context: SendPtr<c_void>,
data: SendPtr<c_void>,
) {
unsafe {
call_js_cb(
env.0 as _,
std::mem::zeroed(),
context.0 as _,
data.0 as _,
)
}
}
spawn(&self.sender, is_blocking, move |_| {
call(call_js_cb, env, context, data);
});
}
} else {
spawn(&self.sender, is_blocking, |_| {
});
};
}
}
#[napi_sym::napi_sym]
fn napi_create_threadsafe_function(
env: *mut Env,
func: napi_value,
_async_resource: napi_value,
_async_resource_name: napi_value,
_max_queue_size: usize,
initial_thread_count: usize,
thread_finalize_data: *mut c_void,
thread_finalize_cb: Option<napi_finalize>,
context: *mut c_void,
maybe_call_js_cb: Option<napi_threadsafe_function_call_js>,
result: *mut napi_threadsafe_function,
) -> napi_status {
let Some(env_ref) = env.as_mut() else {
return napi_generic_failure;
};
if initial_thread_count == 0 {
return napi_invalid_arg;
}
let mut maybe_func = None;
if let Some(value) = *func {
let Ok(func) = v8::Local::<v8::Function>::try_from(value) else {
return napi_function_expected;
};
maybe_func = Some(v8::Global::new(&mut env_ref.scope(), func));
}
let id = TS_FN_ID_COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let tsfn = TsFn {
id,
maybe_func,
maybe_call_js_cb,
context,
thread_counter: initial_thread_count,
sender: env_ref.async_work_sender.clone(),
finalizer: thread_finalize_cb,
finalizer_data: thread_finalize_data,
tsfn_sender: env_ref.threadsafe_function_sender.clone(),
ref_counter: Arc::new(AtomicUsize::new(1)),
env,
};
env_ref
.add_threadsafe_function_ref_counter(tsfn.id, tsfn.ref_counter.clone());
if env_ref
.threadsafe_function_sender
.unbounded_send(ThreadSafeFunctionStatus::Alive)
.is_err()
{
return napi_generic_failure;
}
*result = transmute::<Box<TsFn>, _>(Box::new(tsfn));
napi_ok
}
#[napi_sym::napi_sym]
fn napi_acquire_threadsafe_function(
tsfn: napi_threadsafe_function,
_mode: napi_threadsafe_function_release_mode,
) -> napi_status {
let tsfn: &mut TsFn = &mut *(tsfn as *mut TsFn);
tsfn.acquire()
}
#[napi_sym::napi_sym]
fn napi_unref_threadsafe_function(
_env: &mut Env,
tsfn: napi_threadsafe_function,
) -> napi_status {
let tsfn: &mut TsFn = &mut *(tsfn as *mut TsFn);
tsfn.unref()
}
#[napi_sym::napi_sym]
pub fn napi_get_threadsafe_function_context(
func: napi_threadsafe_function,
result: *mut *const c_void,
) -> napi_status {
let tsfn: &TsFn = &*(func as *const TsFn);
*result = tsfn.context;
napi_ok
}
#[napi_sym::napi_sym]
fn napi_call_threadsafe_function(
func: napi_threadsafe_function,
data: *mut c_void,
is_blocking: napi_threadsafe_function_call_mode,
) -> napi_status {
let tsfn: &TsFn = &*(func as *const TsFn);
tsfn.call(data, is_blocking != 0);
napi_ok
}
#[napi_sym::napi_sym]
fn napi_ref_threadsafe_function(
_env: &mut Env,
func: napi_threadsafe_function,
) -> napi_status {
let tsfn: &mut TsFn = &mut *(func as *mut TsFn);
tsfn.ref_()
}
#[napi_sym::napi_sym]
fn napi_release_threadsafe_function(
tsfn: napi_threadsafe_function,
_mode: napi_threadsafe_function_release_mode,
) -> napi_status {
let tsfn: Box<TsFn> = Box::from_raw(tsfn as *mut TsFn);
tsfn.release()
}