Use N-API ThreadsafeFunction

backend/rust-gbt/index.d.ts

@@ -3,8 +3,16 @@
 
 /* auto-generated by NAPI-RS */
 
-export function make(mempoolBuffer: Uint8Array, callback: (arg0: GbtResult) => void): void
-export function update(newTxs: Uint8Array, removeTxs: Uint8Array, callback: (arg0: GbtResult) => void): void
+export function make(mempoolBuffer: Uint8Array, callback: (result: GbtResult) => void): void
+export function update(newTxs: Uint8Array, removeTxs: Uint8Array, callback: (result: GbtResult) => void): void
+/**
+ * The result from calling the gbt function.
+ *
+ * This tuple contains the following:
+ *   blocks: A 2D Vector of transaction IDs (u32), the inner Vecs each represent a block.
+ * clusters: A 2D Vector of transaction IDs representing clusters of dependent mempool transactions
+ *    rates: A Vector of tuples containing transaction IDs (u32) and effective fee per vsize (f64)
+ */
 export class GbtResult {
   blocks: Array<Array<number>>
   clusters: Array<Array<number>>

backend/rust-gbt/src/lib.rs

@@ -1,4 +1,7 @@
-use napi::bindgen_prelude::*;
+use napi::{
+    bindgen_prelude::*,
+    threadsafe_function::{ErrorStrategy, ThreadsafeFunction, ThreadsafeFunctionCallMode},
+};
 use napi_derive::napi;
 
 use once_cell::sync::Lazy;
@@ -11,14 +14,14 @@ mod thread_transaction;
 mod utils;
 use thread_transaction::ThreadTransaction;
 
+/// Used for ThreadsafeFunction's queue size parameter
+const UNBOUNDED_QUEUE: usize = 0;
+
 static THREAD_TRANSACTIONS: Lazy<Mutex<HashMap<u32, ThreadTransaction>>> =
     Lazy::new(|| Mutex::new(HashMap::new()));
 
-#[napi]
-pub fn make<F>(mempool_buffer: Uint8Array, callback: F) -> Result<()>
-where
-    F: Fn(GbtResult) -> Result<()>,
-{
+#[napi(ts_args_type = "mempoolBuffer: Uint8Array, callback: (result: GbtResult) => void")]
+pub fn make(mempool_buffer: Uint8Array, callback: JsFunction) -> Result<()> {
     let mut map = HashMap::new();
     for tx in ThreadTransaction::batch_from_buffer(&mempool_buffer) {
         map.insert(tx.uid, tx);
@@ -30,11 +33,10 @@ where
     run_in_thread(callback)
 }
 
-#[napi]
-pub fn update<F>(new_txs: Uint8Array, remove_txs: Uint8Array, callback: F) -> Result<()>
-where
-    F: Fn(GbtResult) -> Result<()>,
-{
+#[napi(
+    ts_args_type = "newTxs: Uint8Array, removeTxs: Uint8Array, callback: (result: GbtResult) => void"
+)]
+pub fn update(new_txs: Uint8Array, remove_txs: Uint8Array, callback: JsFunction) -> Result<()> {
     let mut map = THREAD_TRANSACTIONS.lock().unwrap();
     for tx in ThreadTransaction::batch_from_buffer(&new_txs) {
         map.insert(tx.uid, tx);
@@ -60,20 +62,24 @@ pub struct GbtResult {
     pub rates: Vec<Vec<f64>>, // Tuples not supported. u32 fits inside f64
 }
 
-fn run_in_thread<F>(callback: F) -> Result<()>
-where
-    F: Fn(GbtResult) -> Result<()>,
-{
-    let handle = std::thread::spawn(|| -> Result<GbtResult> {
+fn run_in_thread(callback: JsFunction) -> Result<()> {
+    let thread_safe_callback: ThreadsafeFunction<GbtResult, ErrorStrategy::Fatal> =
+        callback.create_threadsafe_function(UNBOUNDED_QUEUE, |ctx| Ok(vec![ctx.value]))?;
+
+    let handle = std::thread::spawn(move || {
         let mut map = THREAD_TRANSACTIONS
             .lock()
             .map_err(|_| napi::Error::from_reason("THREAD_TRANSACTIONS Mutex poisoned"))?;
-        gbt::gbt(&mut map).ok_or_else(|| napi::Error::from_reason("gbt failed"))
+        let result = gbt::gbt(&mut map).ok_or_else(|| napi::Error::from_reason("gbt failed"))?;
+
+        // Note: A call mode of Blocking does not mean it will block, but rather it tells
+        // the N-API what to do in the event of a full queue.
+        // The queue will never be full, so Blocking is fine.
+        thread_safe_callback.call(result, ThreadsafeFunctionCallMode::Blocking);
+        Ok(())
     });
 
-    callback(
-        handle
-            .join()
-            .map_err(|_| napi::Error::from_reason("thread panicked"))??,
-    )
+    handle
+        .join()
+        .map_err(|_| napi::Error::from_reason("thread panicked"))?
 }