refactor(rpc)!: update mempool interface and test code

LagginTimes · LagginTimes · commit 0c726fb075b8 · 2025-04-24T10:40:36.000Z
diff --git a/crates/bitcoind_rpc/src/lib.rs b/crates/bitcoind_rpc/src/lib.rs
@@ -10,8 +10,9 @@
 #![warn(missing_docs)]
 
 use bdk_core::{BlockId, CheckPoint};
-use bitcoin::{block::Header, Block, BlockHash, Transaction};
+use bitcoin::{block::Header, Block, BlockHash, Transaction, Txid};
 use bitcoincore_rpc::bitcoincore_rpc_json;
+use std::collections::HashSet;
 
 pub mod bip158;
 
@@ -73,6 +74,7 @@ impl<'c, C: bitcoincore_rpc::RpcApi> Emitter<'c, C> {
             last_block: None,
             last_mempool_time: 0,
             last_mempool_tip: None,
+            expected_mempool_txids,
         }
     }
 
@@ -93,7 +95,7 @@ impl<'c, C: bitcoincore_rpc::RpcApi> Emitter<'c, C> {
     /// tracked UTXO which is confirmed at height `h`, but the receiver has only seen up to block
     /// of height `h-1`, we want to re-emit this transaction until the receiver has seen the block
     /// at height `h`.
-    pub fn mempool(&mut self) -> Result<Vec<(Transaction, u64)>, bitcoincore_rpc::Error> {
+    pub fn mempool(&mut self) -> Result<MempoolEvent, bitcoincore_rpc::Error> {
         let client = self.client;
 
         // This is the emitted tip height during the last mempool emission.
@@ -136,17 +138,16 @@ impl<'c, C: bitcoincore_rpc::RpcApi> Emitter<'c, C> {
         let prev_mempool_time = self.last_mempool_time;
         let mut latest_time = prev_mempool_time;
 
-        let txs_to_emit = client
-            .get_raw_mempool_verbose()?
+        let new_txs = raw_mempool
             .into_iter()
             .filter_map({
                 let latest_time = &mut latest_time;
+                let evicted_txids = &mut evicted_txids;
                 move |(txid, tx_entry)| -> Option<Result<_, bitcoincore_rpc::Error>> {
                     let tx_time = tx_entry.time as usize;
                     if tx_time > *latest_time {
                         *latest_time = tx_time;
                     }
-
                     // Avoid emitting transactions that are already emitted if we can guarantee
                     // blocks containing ancestors are already emitted. The bitcoind rpc interface
                     // provides us with the block height that the tx is introduced to the mempool.
@@ -157,14 +158,17 @@ impl<'c, C: bitcoincore_rpc::RpcApi> Emitter<'c, C> {
                     if is_already_emitted && is_within_height {
                         return None;
                     }
-
                     let tx = match client.get_raw_transaction(&txid, None) {
                         Ok(tx) => tx,
-                        // the tx is confirmed or evicted since `get_raw_mempool_verbose`
-                        Err(err) if err.is_not_found_error() => return None,
+                        Err(err) if err.is_not_found_error() => {
+                            // If at tip and the transaction isn't found, mark it as evicted.
+                            if at_tip {
+                                evicted_txids.insert(txid);
+                            }
+                            return None;
+                        }
                         Err(err) => return Some(Err(err)),
                     };
-
                     Some(Ok((tx, tx_time as u64)))
                 }
             })
@@ -414,3 +418,77 @@ impl BitcoindRpcErrorExt for bitcoincore_rpc::Error {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use crate::{bitcoincore_rpc::RpcApi, Emitter};
+    use bdk_bitcoind_rpc::bitcoincore_rpc::bitcoin::Txid;
+    use bdk_chain::local_chain::LocalChain;
+    use bdk_testenv::{anyhow, TestEnv};
+    use bitcoin::{hashes::Hash, Address, Amount, ScriptBuf, WScriptHash};
+    use std::collections::HashSet;
+
+    #[test]
+    fn test_expected_mempool_txids_accumulate_and_remove() -> anyhow::Result<()> {
+        let env = TestEnv::new()?;
+        let chain = LocalChain::from_genesis_hash(env.rpc_client().get_block_hash(0)?).0;
+        let chain_tip = chain.tip();
+        let mut emitter = Emitter::new(env.rpc_client(), chain_tip.clone(), 1, HashSet::new());
+
+        env.mine_blocks(100, None)?;
+        while emitter.next_block()?.is_some() {}
+
+        let spk_to_track = ScriptBuf::new_p2wsh(&WScriptHash::all_zeros());
+        let addr_to_track = Address::from_script(&spk_to_track, bitcoin::Network::Regtest)?;
+        let mut mempool_txids = HashSet::new();
+
+        // Send a tx at different heights and ensure txs are accumulating in expected_mempool_txids.
+        for _ in 0..10 {
+            let sent_txid = env.send(&addr_to_track, Amount::from_sat(1_000))?;
+            mempool_txids.insert(sent_txid);
+            emitter.mempool()?;
+            env.mine_blocks(1, None)?;
+
+            for txid in &mempool_txids {
+                assert!(
+                    emitter.expected_mempool_txids.contains(txid),
+                    "Expected txid {:?} missing",
+                    txid
+                );
+            }
+        }
+
+        // Process each block and check that confirmed txids are removed from from
+        // expected_mempool_txids.
+        while let Some(block_event) = emitter.next_block()? {
+            let confirmed_txids: HashSet<Txid> = block_event
+                .block
+                .txdata
+                .iter()
+                .map(|tx| tx.compute_txid())
+                .collect();
+            mempool_txids = mempool_txids
+                .difference(&confirmed_txids)
+                .copied()
+                .collect::<HashSet<_>>();
+            for txid in confirmed_txids {
+                assert!(
+                    !emitter.expected_mempool_txids.contains(&txid),
+                    "Expected txid {:?} should have been removed",
+                    txid
+                );
+            }
+            for txid in &mempool_txids {
+                assert!(
+                    emitter.expected_mempool_txids.contains(txid),
+                    "Expected txid {:?} missing",
+                    txid
+                );
+            }
+        }
+
+        assert!(emitter.expected_mempool_txids.is_empty());
+
+        Ok(())
+    }
+}
diff --git a/crates/bitcoind_rpc/tests/test_emitter.rs b/crates/bitcoind_rpc/tests/test_emitter.rs
@@ -1,4 +1,4 @@
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::{BTreeMap, BTreeSet, HashSet};
 
 use bdk_bitcoind_rpc::Emitter;
 use bdk_chain::{
@@ -22,7 +22,7 @@ pub fn test_sync_local_chain() -> anyhow::Result<()> {
     let env = TestEnv::new()?;
     let network_tip = env.rpc_client().get_block_count()?;
     let (mut local_chain, _) = LocalChain::from_genesis_hash(env.rpc_client().get_block_hash(0)?);
-    let mut emitter = Emitter::new(env.rpc_client(), local_chain.tip(), 0);
+    let mut emitter = Emitter::new(env.rpc_client(), local_chain.tip(), 0, HashSet::new());
 
     // Mine some blocks and return the actual block hashes.
     // Because initializing `ElectrsD` already mines some blocks, we must include those too when
@@ -156,7 +156,7 @@ fn test_into_tx_graph() -> anyhow::Result<()> {
         index
     });
 
-    let emitter = &mut Emitter::new(env.rpc_client(), chain.tip(), 0);
+    let emitter = &mut Emitter::new(env.rpc_client(), chain.tip(), 0, HashSet::new());
 
     while let Some(emission) = emitter.next_block()? {
         let height = emission.block_height();
@@ -189,7 +189,7 @@ fn test_into_tx_graph() -> anyhow::Result<()> {
         assert!(emitter.next_block()?.is_none());
 
         let mempool_txs = emitter.mempool()?;
-        let indexed_additions = indexed_tx_graph.batch_insert_unconfirmed(mempool_txs);
+        let indexed_additions = indexed_tx_graph.batch_insert_unconfirmed(mempool_txs.new_txs);
         assert_eq!(
             indexed_additions
                 .tx_graph
@@ -252,6 +252,7 @@ fn ensure_block_emitted_after_reorg_is_at_reorg_height() -> anyhow::Result<()> {
             hash: env.rpc_client().get_block_hash(0)?,
         }),
         EMITTER_START_HEIGHT as _,
+        HashSet::new(),
     );
 
     env.mine_blocks(CHAIN_TIP_HEIGHT, None)?;
@@ -328,6 +329,7 @@ fn tx_can_become_unconfirmed_after_reorg() -> anyhow::Result<()> {
             hash: env.rpc_client().get_block_hash(0)?,
         }),
         0,
+        HashSet::new(),
     );
 
     // setup addresses
@@ -419,6 +421,7 @@ fn mempool_avoids_re_emission() -> anyhow::Result<()> {
             hash: env.rpc_client().get_block_hash(0)?,
         }),
         0,
+        HashSet::new(),
     );
 
     // mine blocks and sync up emitter
@@ -437,6 +440,7 @@ fn mempool_avoids_re_emission() -> anyhow::Result<()> {
     // the first emission should include all transactions
     let emitted_txids = emitter
         .mempool()?
+        .new_txs
         .into_iter()
         .map(|(tx, _)| tx.compute_txid())
         .collect::<BTreeSet<Txid>>();
@@ -447,7 +451,7 @@ fn mempool_avoids_re_emission() -> anyhow::Result<()> {
 
     // second emission should be empty
     assert!(
-        emitter.mempool()?.is_empty(),
+        emitter.mempool()?.new_txs.is_empty(),
         "second emission should be empty"
     );
 
@@ -457,7 +461,7 @@ fn mempool_avoids_re_emission() -> anyhow::Result<()> {
     }
     while emitter.next_header()?.is_some() {}
     assert!(
-        emitter.mempool()?.is_empty(),
+        emitter.mempool()?.new_txs.is_empty(),
         "third emission, after chain tip is extended, should also be empty"
     );
 
@@ -484,6 +488,7 @@ fn mempool_re_emits_if_tx_introduction_height_not_reached() -> anyhow::Result<()
             hash: env.rpc_client().get_block_hash(0)?,
         }),
         0,
+        HashSet::new(),
     );
 
     // mine blocks to get initial balance, sync emitter up to tip
@@ -506,6 +511,7 @@ fn mempool_re_emits_if_tx_introduction_height_not_reached() -> anyhow::Result<()
     assert_eq!(
         emitter
             .mempool()?
+            .new_txs
             .into_iter()
             .map(|(tx, _)| tx.compute_txid())
             .collect::<BTreeSet<_>>(),
@@ -515,6 +521,7 @@ fn mempool_re_emits_if_tx_introduction_height_not_reached() -> anyhow::Result<()
     assert_eq!(
         emitter
             .mempool()?
+            .new_txs
             .into_iter()
             .map(|(tx, _)| tx.compute_txid())
             .collect::<BTreeSet<_>>(),
@@ -535,6 +542,7 @@ fn mempool_re_emits_if_tx_introduction_height_not_reached() -> anyhow::Result<()
                 .collect::<BTreeSet<_>>();
             let emitted_txids = emitter
                 .mempool()?
+                .new_txs
                 .into_iter()
                 .map(|(tx, _)| tx.compute_txid())
                 .collect::<BTreeSet<_>>();
@@ -572,6 +580,7 @@ fn mempool_during_reorg() -> anyhow::Result<()> {
             hash: env.rpc_client().get_block_hash(0)?,
         }),
         0,
+        HashSet::new(),
     );
 
     // mine blocks to get initial balance
@@ -593,6 +602,7 @@ fn mempool_during_reorg() -> anyhow::Result<()> {
     assert_eq!(
         emitter
             .mempool()?
+            .new_txs
             .into_iter()
             .map(|(tx, _)| tx.compute_txid())
             .collect::<BTreeSet<_>>(),
@@ -628,6 +638,7 @@ fn mempool_during_reorg() -> anyhow::Result<()> {
             // include mempool txs introduced at reorg height or greater
             let mempool = emitter
                 .mempool()?
+                .new_txs
                 .into_iter()
                 .map(|(tx, _)| tx.compute_txid())
                 .collect::<BTreeSet<_>>();
@@ -643,6 +654,7 @@ fn mempool_during_reorg() -> anyhow::Result<()> {
 
             let mempool = emitter
                 .mempool()?
+                .new_txs
                 .into_iter()
                 .map(|(tx, _)| tx.compute_txid())
                 .collect::<BTreeSet<_>>();
@@ -696,6 +708,7 @@ fn no_agreement_point() -> anyhow::Result<()> {
             hash: env.rpc_client().get_block_hash(0)?,
         }),
         (PREMINE_COUNT - 2) as u32,
+        HashSet::new(),
     );
 
     // mine 101 blocks
@@ -821,10 +834,10 @@ fn test_expect_tx_evicted() -> anyhow::Result<()> {
     assert!(mempool_event.evicted_txids.contains(&txid_1));
 
     // Update graph with evicted tx.
-    let exp_txids = exp_spk_txids.into_iter().map(|(_, txid)| txid);
-    let evicted_txids = mempool_event.evicted_txids(exp_txids);
-    for txid in evicted_txids {
-        let _ = graph.insert_evicted_at(txid, seen_at);
+    for txid in mempool_event.evicted_txids {
+        if graph.graph().get_tx_node(txid).is_some() {
+            let _ = graph.insert_evicted_at(txid, seen_at);
+        }
     }
 
     let canonical_txids = graph
diff --git a/examples/example_bitcoind_rpc_polling/src/main.rs b/examples/example_bitcoind_rpc_polling/src/main.rs
@@ -1,4 +1,5 @@
 use std::{
+    collections::HashSet,
     path::PathBuf,
     sync::{
         atomic::{AtomicBool, Ordering},
@@ -11,10 +12,7 @@ use bdk_bitcoind_rpc::{
     bitcoincore_rpc::{Auth, Client, RpcApi},
     Emitter,
 };
-use bdk_chain::{
-    bitcoin::{Block, Transaction},
-    local_chain, Merge,
-};
+use bdk_chain::{bitcoin::Block, local_chain, Merge};
 use example_cli::{
     anyhow,
     clap::{self, Args, Subcommand},
@@ -36,7 +34,7 @@ const DB_COMMIT_DELAY: Duration = Duration::from_secs(60);
 #[derive(Debug)]
 enum Emission {
     Block(bdk_bitcoind_rpc::BlockEvent<Block>),
-    Mempool(Vec<(Transaction, u64)>),
+    Mempool(bdk_bitcoind_rpc::MempoolEvent),
     Tip(u32),
 }
 
@@ -141,7 +139,7 @@ fn main() -> anyhow::Result<()> {
 
             let chain_tip = chain.lock().unwrap().tip();
             let rpc_client = rpc_args.new_client()?;
-            let mut emitter = Emitter::new(&rpc_client, chain_tip, fallback_height);
+            let mut emitter = Emitter::new(&rpc_client, chain_tip, fallback_height, HashSet::new());
             let mut db_stage = ChangeSet::default();
 
             let mut last_db_commit = Instant::now();
@@ -204,7 +202,7 @@ fn main() -> anyhow::Result<()> {
             let graph_changeset = graph
                 .lock()
                 .unwrap()
-                .batch_insert_relevant_unconfirmed(mempool_txs);
+                .batch_insert_relevant_unconfirmed(mempool_txs.new_txs);
             {
                 let db = &mut *db.lock().unwrap();
                 db_stage.merge(ChangeSet {
@@ -232,7 +230,8 @@ fn main() -> anyhow::Result<()> {
             let (tx, rx) = std::sync::mpsc::sync_channel::<Emission>(CHANNEL_BOUND);
             let emission_jh = std::thread::spawn(move || -> anyhow::Result<()> {
                 let rpc_client = rpc_args.new_client()?;
-                let mut emitter = Emitter::new(&rpc_client, last_cp, fallback_height);
+                let mut emitter =
+                    Emitter::new(&rpc_client, last_cp, fallback_height, HashSet::new());
 
                 let mut block_count = rpc_client.get_block_count()? as u32;
                 tx.send(Emission::Tip(block_count))?;
@@ -287,7 +286,13 @@ fn main() -> anyhow::Result<()> {
                         (chain_changeset, graph_changeset)
                     }
                     Emission::Mempool(mempool_txs) => {
-                        let graph_changeset = graph.batch_insert_relevant_unconfirmed(mempool_txs);
+                        let mut graph_changeset =
+                            graph.batch_insert_relevant_unconfirmed(mempool_txs.new_txs.clone());
+                        for txid in mempool_txs.evicted_txids {
+                            graph_changeset.merge(
+                                graph.insert_evicted_at(txid, mempool_txs.latest_update_time),
+                            );
+                        }
                         (local_chain::ChangeSet::default(), graph_changeset)
                     }
                     Emission::Tip(h) => {
