Use a class to hold state for Rust GbtGenerator

backend/rust-gbt/src/lib.rs

@@ -1,9 +1,9 @@
 use napi::bindgen_prelude::*;
 use napi_derive::napi;
 
-use once_cell::sync::Lazy;
 use std::collections::HashMap;
-use std::sync::Mutex;
+use std::sync::{Arc, Mutex};
+use utils::U32HasherState;
 
 mod audit_transaction;
 mod gbt;
@@ -11,41 +11,48 @@ mod thread_transaction;
 mod utils;
 use thread_transaction::ThreadTransaction;
 
-static THREAD_TRANSACTIONS: Lazy<Mutex<HashMap<u32, ThreadTransaction>>> =
-    Lazy::new(|| Mutex::new(HashMap::new()));
+type ThreadTransactionsMap = HashMap<u32, ThreadTransaction, U32HasherState>;
 
-#[napi(ts_args_type = "mempoolBuffer: Uint8Array")]
-pub async fn make(mempool_buffer: Uint8Array) -> Result<GbtResult> {
-    let mut map = HashMap::new();
-    for tx in ThreadTransaction::batch_from_buffer(&mempool_buffer) {
-        map.insert(tx.uid, tx);
-    }
-
-    {
-        let mut global_map = THREAD_TRANSACTIONS
-            .lock()
-            .map_err(|_| napi::Error::from_reason("THREAD_TRANSACTIONS Mutex poisoned"))?;
-        *global_map = map;
-    }
-
-    run_in_thread().await
+#[napi]
+pub struct GbtGenerator {
+    thread_transactions: Arc<Mutex<ThreadTransactionsMap>>,
 }
 
-#[napi(ts_args_type = "newTxs: Uint8Array, removeTxs: Uint8Array")]
-pub async fn update(new_txs: Uint8Array, remove_txs: Uint8Array) -> Result<GbtResult> {
-    {
-        let mut map = THREAD_TRANSACTIONS
-            .lock()
-            .map_err(|_| napi::Error::from_reason("THREAD_TRANSACTIONS Mutex poisoned"))?;
-        for tx in ThreadTransaction::batch_from_buffer(&new_txs) {
-            map.insert(tx.uid, tx);
-        }
-        for txid in &utils::txids_from_buffer(&remove_txs) {
-            map.remove(txid);
+#[napi]
+impl GbtGenerator {
+    #[napi(constructor)]
+    #[allow(clippy::new_without_default)]
+    pub fn new() -> Self {
+        Self {
+            thread_transactions: Arc::new(Mutex::new(HashMap::with_capacity_and_hasher(
+                2048,
+                U32HasherState,
+            ))),
         }
     }
 
-    run_in_thread().await
+    #[napi]
+    pub async fn make(&self, mempool_buffer: Uint8Array) -> Result<GbtResult> {
+        run_task(Arc::clone(&self.thread_transactions), move |map| {
+            for tx in ThreadTransaction::batch_from_buffer(&mempool_buffer) {
+                map.insert(tx.uid, tx);
+            }
+        })
+        .await
+    }
+
+    #[napi]
+    pub async fn update(&self, new_txs: Uint8Array, remove_txs: Uint8Array) -> Result<GbtResult> {
+        run_task(Arc::clone(&self.thread_transactions), move |map| {
+            for tx in ThreadTransaction::batch_from_buffer(&new_txs) {
+                map.insert(tx.uid, tx);
+            }
+            for txid in &utils::txids_from_buffer(&remove_txs) {
+                map.remove(txid);
+            }
+        })
+        .await
+    }
 }
 
 /// The result from calling the gbt function.
@@ -61,11 +68,27 @@ pub struct GbtResult {
     pub rates: Vec<Vec<f64>>, // Tuples not supported. u32 fits inside f64
 }
 
-async fn run_in_thread() -> Result<GbtResult> {
+/// All on another thread, this runs an arbitrary task in between
+/// taking the lock and running gbt.
+///
+/// Rather than filling / updating the HashMap on the main thread,
+/// this allows for HashMap modifying tasks to be run before running and returning gbt results.
+///
+/// `thread_transactions` is a cloned Arc of the Mutex for the HashMap state.
+/// `callback` is a `'static + Send` `FnOnce` closure/function that takes a mutable reference
+/// to the HashMap as the only argument. (A move closure is recommended to meet the bounds)
+async fn run_task<F>(
+    thread_transactions: Arc<Mutex<ThreadTransactionsMap>>,
+    callback: F,
+) -> Result<GbtResult>
+where
+    F: FnOnce(&mut ThreadTransactionsMap) + Send + 'static,
+{
     let handle = napi::tokio::task::spawn_blocking(move || {
-        let mut map = THREAD_TRANSACTIONS
+        let mut map = thread_transactions
             .lock()
             .map_err(|_| napi::Error::from_reason("THREAD_TRANSACTIONS Mutex poisoned"))?;
+        callback(&mut map);
         gbt::gbt(&mut map).ok_or_else(|| napi::Error::from_reason("gbt failed"))
     });