Implement rust gbt updateBlockTemplates

backend/rust-gbt/Cargo.lock

@@ -2,6 +2,18 @@
 # It is not intended for manual editing.
 version = 3
 
+[[package]]
+name = "autocfg"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
+
+[[package]]
+name = "bytes"
+version = "1.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "89b2fd2a0dcf38d7971e2194b6b6eebab45ae01067456a7fd93d5547a61b70be"
+
 [[package]]
 name = "cfg-if"
 version = "1.0.0"
@@ -12,7 +24,26 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 name = "gbt"
 version = "0.1.0"
 dependencies = [
+ "bytes",
  "neon",
+ "once_cell",
+ "priority-queue",
+]
+
+[[package]]
+name = "hashbrown"
+version = "0.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
+
+[[package]]
+name = "indexmap"
+version = "1.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
+dependencies = [
+ "autocfg",
+ "hashbrown",
 ]
 
 [[package]]
@@ -66,6 +97,22 @@ dependencies = [
  "smallvec",
 ]
 
+[[package]]
+name = "once_cell"
+version = "1.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d"
+
+[[package]]
+name = "priority-queue"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fff39edfcaec0d64e8d0da38564fad195d2d51b680940295fcc307366e101e61"
+dependencies = [
+ "autocfg",
+ "indexmap",
+]
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.60"

backend/rust-gbt/Cargo.toml

@@ -14,6 +14,7 @@ crate-type = ["cdylib"]
 [dependencies]
 priority-queue = "1.3.2"
 bytes = "1.4.0"
+once_cell = "1.18.0"
 
 [dependencies.neon]
 version = "0.10"

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>>) {
+pub fn gbt(mempool: &mut HashMap<u32,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);
-  }
-
-
   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;
+}

backend/src/api/mempool-blocks.ts

@@ -11,6 +11,7 @@ class MempoolBlocks {
   private mempoolBlocks: MempoolBlockWithTransactions[] = [];
   private mempoolBlockDeltas: MempoolBlockDelta[] = [];
   private txSelectionWorker: Worker | null = null;
+  private rustInitialized: boolean = false;
 
   private nextUid: number = 1;
   private uidMap: Map<number, string> = new Map(); // map short numerical uids to full txids
@@ -284,7 +285,7 @@ class MempoolBlocks {
     const start = Date.now();
 
     for (const tx of Object.values(added)) {
-      this.setUid(tx);
+      this.setUid(tx, true);
     }
     const removedUids = removed.map(tx => this.getUid(tx)).filter(uid => uid != null) as number[];
     // prepare a stripped down version of the mempool with only the minimum necessary data
@@ -337,23 +338,49 @@ class MempoolBlocks {
 
     // serialize relevant mempool data into an ArrayBuffer
     // to reduce the overhead of passing this data to the rust thread
-    const mempoolBuffer = this.mempoolToArrayBuffer(newMempool);
+    const mempoolBuffer = this.mempoolToArrayBuffer(Object.values(newMempool), newMempool);
 
     // run the block construction algorithm in a separate thread, and wait for a result
     try {
-      const { blocks, rates, clusters } = this.convertNeonResultTxids(await new Promise((resolve) => { neonAddon.go(mempoolBuffer, resolve); }));
+      const { blocks, rates, clusters } = this.convertNeonResultTxids(await new Promise((resolve) => { neonAddon.make(mempoolBuffer, resolve); }));
+      this.rustInitialized = true;
       const processed = this.processBlockTemplates(newMempool, blocks, rates, clusters, saveResults);
       logger.debug(`RUST makeBlockTemplates completed in ${(Date.now() - start)/1000} seconds`);
       return processed;
     } catch (e) {
+      this.rustInitialized = false;
       logger.err('RUST makeBlockTemplates failed. ' + (e instanceof Error ? e.message : e));
     }
     return this.mempoolBlocks;
   }
 
   public async $rustUpdateBlockTemplates(newMempool: { [txid: string]: MempoolTransactionExtended }, added: MempoolTransactionExtended[], removed: MempoolTransactionExtended[], saveResults: boolean = false): Promise<void> {
-    await this.$rustMakeBlockTemplates(newMempool, saveResults);
+    if (!this.rustInitialized) {
-    return;
+      // need to reset the worker
+      await this.$rustMakeBlockTemplates(newMempool, saveResults);
+      return;
+    }
+
+    const start = Date.now();
+
+    for (const tx of Object.values(added)) {
+      this.setUid(tx, true);
+    }
+    const removedUids = removed.map(tx => this.getUid(tx)).filter(uid => uid != null) as number[];
+    // serialize relevant mempool data into an ArrayBuffer
+    // to reduce the overhead of passing this data to the rust thread
+    const addedBuffer = this.mempoolToArrayBuffer(added, newMempool);
+    const removedBuffer = this.uidsToArrayBuffer(removedUids);
+
+    // run the block construction algorithm in a separate thread, and wait for a result
+    try {
+      const { blocks, rates, clusters } = this.convertNeonResultTxids(await new Promise((resolve) => { neonAddon.update(addedBuffer, removedBuffer, resolve); }));
+      this.processBlockTemplates(newMempool, blocks, rates, clusters, saveResults);
+      logger.debug(`RUST updateBlockTemplates completed in ${(Date.now() - start)/1000} seconds`);
+    } catch (e) {
+      this.rustInitialized = false;
+      logger.err('RUST updateBlockTemplates failed. ' + (e instanceof Error ? e.message : e));
+    }
   }
 
   private processBlockTemplates(mempool, blocks: string[][], rates: { [root: string]: number }, clusters: { [root: string]: string[] }, saveResults): MempoolBlockWithTransactions[] {
@@ -487,12 +514,16 @@ class MempoolBlocks {
     this.nextUid = 1;
   }
 
-  private setUid(tx: MempoolTransactionExtended): number {
+  private setUid(tx: MempoolTransactionExtended, skipSet = false): number {
-    const uid = this.nextUid;
+    if (tx.uid == null || !skipSet) {
-    this.nextUid++;
+      const uid = this.nextUid;
-    this.uidMap.set(uid, tx.txid);
+      this.nextUid++;
-    tx.uid = uid;
+      this.uidMap.set(uid, tx.txid);
-    return uid;
+      tx.uid = uid;
+      return uid;
+    } else {
+      return tx.uid;
+    }
   }
 
   private getUid(tx: MempoolTransactionExtended): number | void {
@@ -565,14 +596,14 @@ class MempoolBlocks {
     return { blocks: convertedBlocks, rates: convertedRates, clusters: convertedClusters } as { blocks: string[][], rates: { [root: string]: number }, clusters: { [root: string]: string[] }};
   }
 
-  private mempoolToArrayBuffer(mempool: { [txid: string]: MempoolTransactionExtended }): ArrayBuffer {
+  private mempoolToArrayBuffer(txs: MempoolTransactionExtended[], mempool: { [txid: string]: MempoolTransactionExtended }): ArrayBuffer {
     let len = 4;
     const inputs: { [uid: number]: number[] } = {};
     let validCount = 0;
-    for (const tx of Object.values(mempool)) {
+    for (const tx of txs) {
       if (tx.uid != null) {
         validCount++;
-        const txInputs = tx.vin.map(v => this.getUid(mempool[v.txid])).filter(uid => uid != null) as number[]
+        const txInputs = tx.vin.map(v => this.getUid(mempool[v.txid])).filter(uid => uid != null) as number[];
         inputs[tx.uid] = txInputs;
         len += (10 + txInputs.length) * 4;
       }
@@ -581,7 +612,7 @@ class MempoolBlocks {
     const view = new DataView(buf);
     view.setUint32(0, validCount, false);
     let offset = 4;
-    for (const tx of Object.values(mempool)) {
+    for (const tx of txs) {
       if (tx.uid != null) {
         view.setUint32(offset, tx.uid, false);
         view.setFloat64(offset + 4, tx.fee, false);
@@ -599,6 +630,19 @@ class MempoolBlocks {
     }
     return buf;
   }
+
+  private uidsToArrayBuffer(uids: number[]): ArrayBuffer {
+    let len = (uids.length + 1) * 4;
+    const buf = new ArrayBuffer(len);
+    const view = new DataView(buf);
+    view.setUint32(0, uids.length, false);
+    let offset = 4;
+    for (const uid of uids) {
+      view.setUint32(offset, uid, false);
+      offset += 4;
+    }
+    return buf;
+  }
 }
 
 export default new MempoolBlocks();