WIP: Benchmark TxGraph queries

.github/workflows/nightly_docs.yml

@@ -10,7 +10,7 @@ jobs:
       - name: Checkout sources
         uses: actions/checkout@v4
       - name: Set default toolchain
-        run: rustup default nightly-2024-05-12
+        run: rustup default nightly-2024-11-17
       - name: Set profile
         run: rustup set profile minimal
       - name: Update toolchain

.gitignore

@@ -1,4 +1,4 @@
-/target
+**/target
 Cargo.lock
 /.vscode
 

Cargo.toml

@@ -18,10 +18,17 @@ members = [
     "example-crates/example_wallet_esplora_async",
     "example-crates/example_wallet_rpc",
 ]
+exclude = ["bench"]
 
 [workspace.package]
 authors = ["Bitcoin Dev Kit Developers"]
 
 [workspace.lints.clippy]
 print_stdout = "deny"
 print_stderr = "deny"
+
+[workspace.lints.rust.unexpected_cfgs]
+level = "forbid"
+check-cfg = [
+    "cfg(bdk_bench)",
+]

bench/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "bdk_bench"
+version = "0.1.0"
+edition = "2021"
+
+[[bench]]
+name = "bench"
+harness = false
+
+[dependencies]
+bdk_chain = { path = "../crates/chain", features = ["criterion"] }
+criterion = { version = "0.5", default-features = false }

bench/README.md

@@ -0,0 +1,5 @@
+# BDK bench
+
+To run benchmarks in the current directory:
+
+`RUSTFLAGS="--cfg=bdk_bench" cargo bench`

bench/benches/bench.rs

@@ -0,0 +1,11 @@
+extern crate bdk_chain;
+extern crate criterion;
+
+use criterion::{criterion_group, criterion_main};
+
+criterion_group!(benches,
+    bdk_chain::tx_graph::bench::filter_chain_unspents,
+    bdk_chain::tx_graph::bench::list_canonical_txs,
+    bdk_chain::tx_graph::bench::nested_conflicts,
+);
+criterion_main!(benches);

crates/chain/Cargo.toml

@@ -30,6 +30,9 @@ rand = "0.8"
 proptest = "1.2.0"
 bdk_testenv = { path = "../testenv", default-features = false }
 
+[target.'cfg(bdk_bench)'.dependencies]
+criterion = { version = "0.5", optional = true, default-features = false }
+
 
 [features]
 default = ["std", "miniscript"]

crates/chain/src/lib.rs

@@ -68,6 +68,8 @@ pub use bdk_core::*;
 #[allow(unused_imports)]
 #[macro_use]
 extern crate alloc;
+#[cfg(bdk_bench)]
+extern crate criterion;
 #[cfg(feature = "rusqlite")]
 pub extern crate rusqlite;
 #[cfg(feature = "serde")]

crates/chain/src/tx_graph.rs

@@ -1553,3 +1553,311 @@ where
 fn tx_outpoint_range(txid: Txid) -> RangeInclusive<OutPoint> {
     OutPoint::new(txid, u32::MIN)..=OutPoint::new(txid, u32::MAX)
 }
+
+#[cfg(any(test, bdk_bench))]
+mod bench_util {
+    use bdk_core::{CheckPoint, ConfirmationBlockTime};
+    use bitcoin::{
+        absolute, constants, hashes::Hash, secp256k1::Secp256k1, transaction, BlockHash, Network,
+        TxIn,
+    };
+    use miniscript::{Descriptor, DescriptorPublicKey};
+
+    use super::*;
+    use crate::keychain_txout::KeychainTxOutIndex;
+    use crate::local_chain::LocalChain;
+    use crate::IndexedTxGraph;
+
+    #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+    pub enum Keychain {
+        External,
+    }
+
+    const EXTERNAL: &str = "tr([ab28dc00/86h/1h/0h]tpubDCdDtzAMZZrkwKBxwNcGCqe4FRydeD9rfMisoi7qLdraG79YohRfPW4YgdKQhpgASdvh612xXNY5xYzoqnyCgPbkpK4LSVcH5Xv4cK7johH/0/*)";
+
+    pub fn parse_descriptor(s: &str) -> Descriptor<DescriptorPublicKey> {
+        <Descriptor<DescriptorPublicKey>>::parse_descriptor(&Secp256k1::new(), s)
+            .unwrap()
+            .0
+    }
+
+    /// New tx guaranteed to have at least one output
+    pub fn new_tx(lt: u32) -> Transaction {
+        Transaction {
+            version: transaction::Version::TWO,
+            lock_time: absolute::LockTime::from_consensus(lt),
+            input: vec![],
+            output: vec![TxOut::NULL],
+        }
+    }
+
+    pub fn spk_at_index(txout_index: &KeychainTxOutIndex<Keychain>, index: u32) -> ScriptBuf {
+        txout_index
+            .get_descriptor(Keychain::External)
+            .unwrap()
+            .at_derivation_index(index)
+            .unwrap()
+            .script_pubkey()
+    }
+
+    type KeychainTxGraph = IndexedTxGraph<ConfirmationBlockTime, KeychainTxOutIndex<Keychain>>;
+
+    /// Initialize indexed tx-graph with one keychain and a local chain. Also insert the
+    /// first ancestor tx.
+    pub fn init_graph_chain() -> (KeychainTxGraph, LocalChain) {
+        let genesis = constants::genesis_block(Network::Regtest).block_hash();
+        let block_0 = BlockId {
+            height: 0,
+            hash: genesis,
+        };
+        // chain tip 100
+        let mut cp = CheckPoint::new(block_0);
+        let chain_tip = BlockId {
+            height: 100,
+            hash: BlockHash::all_zeros(),
+        };
+        cp = cp.push(chain_tip).unwrap();
+        let chain = LocalChain::from_tip(cp).unwrap();
+
+        let desc = parse_descriptor(EXTERNAL);
+        let mut index = KeychainTxOutIndex::new(10);
+        index.insert_descriptor(Keychain::External, desc).unwrap();
+        let mut graph = IndexedTxGraph::new(index);
+
+        // insert funding tx (coinbase) confirmed at chain tip
+        add_ancestor_tx(&mut graph, chain_tip, 0);
+
+        (graph, chain)
+    }
+
+    /// Add ancestor tx confirmed at `block_id` with `locktime` (used for uniqueness).
+    /// The transaction always pays 1 BTC to SPK 0.
+    pub fn add_ancestor_tx(graph: &mut KeychainTxGraph, block_id: BlockId, locktime: u32) {
+        let spk_0 = spk_at_index(&graph.index, 0);
+        let tx = Transaction {
+            input: vec![TxIn::default()],
+            output: vec![TxOut {
+                value: Amount::ONE_BTC,
+                script_pubkey: spk_0,
+            }],
+            ..new_tx(locktime)
+        };
+        let txid = tx.compute_txid();
+        let _ = graph.insert_tx(tx);
+        let _ = graph.insert_anchor(
+            txid,
+            ConfirmationBlockTime {
+                block_id,
+                confirmation_time: 100,
+            },
+        );
+    }
+
+    /// Add `n` conflicts to `graph` that spend the given `previous_output`, incrementing
+    /// the tx last-seen each time.
+    pub fn add_conflicts(n: u32, graph: &mut KeychainTxGraph, previous_output: OutPoint) {
+        let spk_1 = spk_at_index(&graph.index, 1);
+        for i in 1..n + 1 {
+            let tx = Transaction {
+                input: vec![TxIn {
+                    previous_output,
+                    ..Default::default()
+                }],
+                output: vec![TxOut {
+                    value: Amount::ONE_BTC - Amount::from_sat(i as u64 * 10),
+                    script_pubkey: spk_1.clone(),
+                }],
+                ..new_tx(i)
+            };
+            let update = TxUpdate {
+                txs: vec![Arc::new(tx)],
+                ..Default::default()
+            };
+            let _ = graph.apply_update_at(update, Some(i as u64));
+        }
+    }
+
+    /// Apply a chain of `n` unconfirmed txs where each subsequent tx spends the output
+    /// of the previous one.
+    pub fn chain_unconfirmed(n: u32, graph: &mut KeychainTxGraph, mut previous_output: OutPoint) {
+        for i in 0..n {
+            // create tx
+            let tx = Transaction {
+                input: vec![TxIn {
+                    previous_output,
+                    ..Default::default()
+                }],
+                ..new_tx(i)
+            };
+            let txid = tx.compute_txid();
+            let update = TxUpdate {
+                txs: vec![Arc::new(tx)],
+                ..Default::default()
+            };
+            let _ = graph.apply_update_at(update, Some(21));
+            // store the next prevout
+            previous_output = OutPoint::new(txid, 0);
+        }
+    }
+
+    /// Insert `n` txs where
+    /// - half spend ancestor A
+    /// - half spend ancestor B
+    /// - and one spends both
+    pub fn add_nested_conflicts(n: u32, graph: &mut KeychainTxGraph, chain: &LocalChain) {
+        // add ancestor B
+        add_ancestor_tx(graph, chain.tip().block_id(), 1);
+
+        let outpoints: Vec<_> = graph.index.outpoints().iter().map(|(_, op)| *op).collect();
+        assert!(outpoints.len() >= 2);
+        let op_a = outpoints[0];
+        let op_b = outpoints[1];
+
+        for i in 0..n {
+            let tx = if i == n / 2 {
+                // tx spends both A, B
+                Transaction {
+                    input: vec![
+                        TxIn {
+                            previous_output: op_a,
+                            ..Default::default()
+                        },
+                        TxIn {
+                            previous_output: op_b,
+                            ..Default::default()
+                        },
+                    ],
+                    ..new_tx(i)
+                }
+            } else if i % 2 == 1 {
+                // tx spends A
+                Transaction {
+                    input: vec![TxIn {
+                        previous_output: op_a,
+                        ..Default::default()
+                    }],
+                    ..new_tx(i)
+                }
+            } else {
+                // tx spends B
+                Transaction {
+                    input: vec![TxIn {
+                        previous_output: op_b,
+                        ..Default::default()
+                    }],
+                    ..new_tx(i)
+                }
+            };
+
+            let update = TxUpdate {
+                txs: vec![Arc::new(tx)],
+                ..Default::default()
+            };
+            let _ = graph.apply_update_at(update, Some(i as u64));
+        }
+    }
+
+    #[test]
+    fn test_add_conflicts() {
+        let (mut graph, chain) = init_graph_chain();
+        let txouts: Vec<_> = graph.graph().all_txouts().collect();
+        assert_eq!(txouts.len(), 1);
+        let prevout = txouts.first().unwrap().0;
+        add_conflicts(3, &mut graph, prevout);
+
+        let unspent = graph.graph().filter_chain_unspents(
+            &chain,
+            chain.tip().block_id(),
+            graph.index.outpoints().clone(),
+        );
+        assert_eq!(unspent.count(), 1);
+    }
+
+    #[test]
+    fn test_chain_unconfirmed() {
+        let (mut graph, _) = init_graph_chain();
+        let (prevout, _txout) = graph
+            .graph()
+            .all_txouts()
+            .find(|(_, txout)| txout.value == Amount::ONE_BTC)
+            .expect("initial graph should have txout");
+        chain_unconfirmed(5, &mut graph, prevout);
+        assert_eq!(graph.graph().txs.len(), 6); // 1 onchain + 5 unconfirmed
+    }
+
+    #[test]
+    #[rustfmt::skip]
+    fn test_nested() {
+        let (mut graph, chain) = init_graph_chain();
+        let chain_tip = chain.tip().block_id();
+        let n = 5;
+        add_nested_conflicts(n, &mut graph, &chain);
+
+        let op = graph.index.outpoints().clone();
+        assert_eq!(graph.graph().full_txs().count() as u32, 2 + n); // 2 onchain + n unconfirmed
+        assert_eq!(graph.graph().filter_chain_txouts(&chain, chain_tip, op).count(), 2); // 2 onchain
+        assert_eq!(graph.graph().list_canonical_txs(&chain, chain_tip).count(), 2 + 2); // 2 onchain + 2 unconfirmed
+    }
+}
+
+/// Bench
+#[allow(missing_docs)]
+#[cfg(bdk_bench)]
+pub mod bench {
+    use std::hint::black_box;
+
+    use criterion::Criterion;
+
+    use super::bench_util::*;
+
+    #[inline(never)]
+    pub fn filter_chain_unspents(bench: &mut Criterion) {
+        let (mut graph, chain) = black_box(init_graph_chain());
+        let prevout = graph.graph().all_txouts().next().unwrap().0;
+        black_box(add_conflicts(1000, &mut graph, prevout));
+
+        bench.bench_function("filter_chain_unspents", |b| {
+            b.iter(|| {
+                let unspent = graph.graph().filter_chain_unspents(
+                    &chain,
+                    chain.tip().block_id(),
+                    graph.index.outpoints().clone(),
+                );
+                assert_eq!(unspent.count(), 1);
+            })
+        });
+    }
+
+    #[inline(never)]
+    pub fn list_canonical_txs(bench: &mut Criterion) {
+        let (mut graph, chain) = black_box(init_graph_chain());
+        let prevout = graph.graph().all_txouts().next().unwrap().0;
+        black_box(chain_unconfirmed(100, &mut graph, prevout));
+
+        bench.bench_function("list_canonical_txs", |b| {
+            b.iter(|| {
+                let txs = graph
+                    .graph()
+                    .list_canonical_txs(&chain, chain.tip().block_id());
+                // 1 onchain + 100 unconfirmed
+                assert_eq!(txs.count(), 1 + 100);
+            })
+        });
+    }
+
+    #[inline(never)]
+    pub fn nested_conflicts(bench: &mut Criterion) {
+        let (mut graph, chain) = black_box(init_graph_chain());
+        black_box(add_nested_conflicts(2000, &mut graph, &chain));
+        let graph = graph.graph();
+        let chain_tip = chain.tip().block_id();
+
+        bench.bench_function("nested_conflicts", |b| {
+            b.iter(|| {
+                let txs = graph.list_canonical_txs(&chain, chain_tip);
+                // 2 onchain + 2 unconfirmed
+                assert_eq!(txs.count(), 4);
+            })
+        });
+    }
+}