Implement rust gbt updateBlockTemplates

backend/rust-gbt/src/gbt.rs

@@ -34,13 +34,7 @@ impl Ord for TxPriority {
   }
 }
 
-pub fn gbt(mempool_array: Vec<ThreadTransaction>) -> (Vec<Vec<u32>>, Vec<(u32, f64)>, Vec<Vec<u32>>) {
-  let mut mempool: HashMap<u32,ThreadTransaction> = HashMap::new();
-  for transaction in mempool_array {
-    mempool.insert(transaction.uid, transaction);
-  }
-
-
+pub fn gbt(mempool: &mut HashMap<u32,ThreadTransaction>) -> (Vec<Vec<u32>>, Vec<(u32, f64)>, Vec<Vec<u32>>) {
   return make_block_templates(mempool);
 }
 
@@ -49,13 +43,13 @@ pub fn gbt(mempool_array: Vec<ThreadTransaction>) -> (Vec<Vec<u32>>, Vec<(u32, f
 * (see BlockAssembler in https://github.com/bitcoin/bitcoin/blob/master/src/node/miner.cpp)
 * Ported from https://github.com/mempool/mempool/blob/master/backend/src/api/tx-selection-worker.ts
 */
-fn make_block_templates(mempool: HashMap<u32,ThreadTransaction>) -> (Vec<Vec<u32>>, Vec<(u32, f64)>, Vec<Vec<u32>>) {
+fn make_block_templates(mempool: &mut HashMap<u32,ThreadTransaction>) -> (Vec<Vec<u32>>, Vec<(u32, f64)>, Vec<Vec<u32>>) {
   let mut audit_pool: HashMap<u32, AuditTransaction> = HashMap::new();
   let mut mempool_array: VecDeque<u32> = VecDeque::new();
   let mut cluster_array: Vec<Vec<u32>> = Vec::new();
 
   // Initialize working structs
-  for (uid, tx) in &mempool {
+  for (uid, tx) in mempool {
     let audit_tx = AuditTransaction {
       uid: tx.uid,
       fee: tx.fee,

backend/rust-gbt/src/lib.rs

@@ -1,11 +1,20 @@
 use neon::{prelude::*, types::buffer::TypedArray};
+use std::collections::HashMap;
+use std::ops::DerefMut;
+use std::sync::Mutex;
+use once_cell::sync::Lazy;
 
 mod gbt;
 mod thread_transaction;
 mod audit_transaction;
-use thread_transaction::{ThreadTransaction};
+mod utils;
+use thread_transaction::ThreadTransaction;
 
-fn go(mut cx: FunctionContext) -> JsResult<JsUndefined> {
+static THREAD_TRANSACTIONS: Lazy<Mutex<HashMap<u32, ThreadTransaction>>> = Lazy::new(|| {
+  Mutex::new(HashMap::new())
+});
+
+fn make(mut cx: FunctionContext) -> JsResult<JsUndefined> {
   let mempool_arg = cx.argument::<JsArrayBuffer>(0)?.root(&mut cx).into_inner(&mut cx);
   let callback = cx.argument::<JsFunction>(1)?.root(&mut cx);
   let channel = cx.channel();
@@ -13,8 +22,48 @@ fn go(mut cx: FunctionContext) -> JsResult<JsUndefined> {
   let buffer = mempool_arg.as_slice(&mut cx);
   let thread_transactions = ThreadTransaction::batch_from_buffer(buffer);
 
+  let mut map = THREAD_TRANSACTIONS.lock().unwrap();
+  map.clear();
+  for tx in thread_transactions {
+    map.insert(tx.uid, tx);
+  }
+  drop(map);
+
+  run_in_thread(channel, callback);
+
+  Ok(cx.undefined())
+}
+
+fn update(mut cx: FunctionContext) -> JsResult<JsUndefined> {
+  let new_txs_arg = cx.argument::<JsArrayBuffer>(0)?.root(&mut cx).into_inner(&mut cx);
+  let remove_txs_arg = cx.argument::<JsArrayBuffer>(1)?.root(&mut cx).into_inner(&mut cx);
+  let callback = cx.argument::<JsFunction>(2)?.root(&mut cx);
+  let channel = cx.channel();
+
+  let mut map = THREAD_TRANSACTIONS.lock().unwrap();
+  let new_tx_buffer = new_txs_arg.as_slice(&mut cx);
+  let thread_transactions = ThreadTransaction::batch_from_buffer(new_tx_buffer);
+  for tx in thread_transactions {
+    map.insert(tx.uid, tx);
+  }
+
+  let remove_tx_buffer = remove_txs_arg.as_slice(&mut cx);
+  let remove_ids = utils::txids_from_buffer(remove_tx_buffer);
+  for txid in &remove_ids {
+    map.remove(txid);
+  }
+  drop(map);
+
+  run_in_thread(channel, callback);
+
+  Ok(cx.undefined())
+}
+
+fn run_in_thread(channel: Channel, callback: Root<JsFunction>) {
   std::thread::spawn(move || {
-    let (blocks, rates, clusters) = gbt::gbt(thread_transactions);
+    let mut map = THREAD_TRANSACTIONS.lock().unwrap();
+    let (blocks, rates, clusters) = gbt::gbt(map.deref_mut());
+    drop(map);
 
     channel.send(move |mut cx| {
       let result = JsObject::new(&mut cx);
@@ -64,12 +113,11 @@ fn go(mut cx: FunctionContext) -> JsResult<JsUndefined> {
       Ok(())
     });
   });
-
-  Ok(cx.undefined())
 }
 
 #[neon::main]
 fn main(mut cx: ModuleContext) -> NeonResult<()> {
-  cx.export_function("go", go)?;
+  cx.export_function("make", make)?;
+  cx.export_function("update", update)?;
   Ok(())
 }

backend/rust-gbt/src/thread_transaction.rs

@@ -2,6 +2,7 @@
 extern crate bytes;
 use std::io::Cursor;
 use bytes::buf::Buf;
+
 pub struct ThreadTransaction {
   pub uid: u32,
   pub fee: u64,

backend/rust-gbt/src/utils.rs

@@ -0,0 +1,14 @@
+extern crate bytes;
+use std::io::Cursor;
+use bytes::buf::Buf;
+
+pub fn txids_from_buffer(buffer: &[u8]) -> Vec<u32> {
+  let mut txids: Vec<u32> = Vec::new();
+  let mut cursor = Cursor::new(buffer);
+  let size = cursor.get_u32();
+  for _ in 0..size {
+    txids.push(cursor.get_u32());
+  }
+
+  return txids;
+}