From ff5c10302022140f2604b92676e5b7af322ab2bb Mon Sep 17 00:00:00 2001
From: LongYinan <lynweklm@gmail.com>
Date: Mon, 28 Dec 2020 14:02:28 +0800
Subject: [PATCH] chore: add thread safe function bench suite

---
 bench/Cargo.toml           |  2 +-
 bench/async.ts             | 15 +++++++++++++--
 bench/src/async_compute.rs | 33 ++++++++++++++++++++++++++++++---
 3 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/bench/Cargo.toml b/bench/Cargo.toml
index 245b1344..ddfa7d2e 100644
--- a/bench/Cargo.toml
+++ b/bench/Cargo.toml
@@ -8,7 +8,7 @@ version = "0.1.0"
 crate-type = ["cdylib"]
 
 [dependencies]
-napi = {path = "../napi"}
+napi = {path = "../napi", features = ["napi4"]}
 napi-derive = {path = "../napi-derive"}
 
 [target.'cfg(all(unix, not(target_env = "musl"), not(target_arch = "aarch64")))'.dependencies]
diff --git a/bench/async.ts b/bench/async.ts
index bf9179a1..e41ec044 100644
--- a/bench/async.ts
+++ b/bench/async.ts
@@ -1,15 +1,26 @@
 import b from 'benny'
 
-const { benchAsyncTask } = require('./index.node')
+const { benchAsyncTask, benchThreadsafeFunction } = require('./index.node')
 
 const buffer = Buffer.from('hello 🚀 rust!')
 
 export const benchAsync = () =>
   b.suite(
     'Async task',
-    b.add('napi-rs', async () => {
+    b.add('spawn task', async () => {
       await benchAsyncTask(buffer)
     }),
+    b.add('thread safe function', async () => {
+      await new Promise<number | undefined>((resolve, reject) => {
+        benchThreadsafeFunction(buffer, (err?: Error, value?: number) => {
+          if (err) {
+            reject(err)
+          } else {
+            resolve(value)
+          }
+        })
+      })
+    }),
     b.cycle(),
     b.complete(),
   )
diff --git a/bench/src/async_compute.rs b/bench/src/async_compute.rs
index 87300303..41d8c82f 100644
--- a/bench/src/async_compute.rs
+++ b/bench/src/async_compute.rs
@@ -1,5 +1,5 @@
-use napi::{CallContext, Env, JsBuffer, JsBufferValue, JsNumber, JsObject, Ref, Result, Task};
-
+use napi::threadsafe_function::*;
+use napi::*;
 
 struct BufferLength(Ref<JsBufferValue>);
 
@@ -8,7 +8,7 @@ impl Task for BufferLength {
   type JsValue = JsNumber;
 
   fn compute(&mut self) -> Result<Self::Output> {
-    Ok((&self.0).len())
+    Ok(self.0.len() + 1)
   }
 
   fn resolve(self, env: Env, output: Self::Output) -> Result<Self::JsValue> {
@@ -25,7 +25,34 @@ fn bench_async_task(ctx: CallContext) -> Result<JsObject> {
   Ok(async_promise.promise_object())
 }
 
+#[js_function(2)]
+fn bench_threadsafe_function(ctx: CallContext) -> Result<JsUndefined> {
+  let buffer_ref = ctx.get::<JsBuffer>(0)?.into_ref()?;
+  let callback = ctx.get::<JsFunction>(1)?;
+
+  let tsfn = ctx.env.create_threadsafe_function(
+    &callback,
+    0,
+    |ctx: ThreadSafeCallContext<(usize, Ref<JsBufferValue>)>| {
+      ctx
+        .env
+        .create_uint32(ctx.value.0 as u32)
+        .and_then(|v| ctx.value.1.unref(ctx.env).map(|_| vec![v]))
+    },
+  )?;
+
+  std::thread::spawn(move || {
+    tsfn.call(
+      Ok((buffer_ref.len() + 1, buffer_ref)),
+      ThreadsafeFunctionCallMode::NonBlocking,
+    );
+  });
+
+  ctx.env.get_undefined()
+}
+
 pub fn register_js(exports: &mut JsObject) -> Result<()> {
   exports.create_named_method("benchAsyncTask", bench_async_task)?;
+  exports.create_named_method("benchThreadsafeFunction", bench_threadsafe_function)?;
   Ok(())
 }