From 3357c58c41b505b919e2361da138519d23e3a0c4 Mon Sep 17 00:00:00 2001 From: teor Date: Thu, 7 Oct 2021 08:26:36 +1000 Subject: [PATCH] Return transaction iterators from mempool storage Also rename transaction methods for consistency. --- zebrad/src/components/mempool.rs | 10 ++-- zebrad/src/components/mempool/storage.rs | 49 ++++++++++++------- .../mempool/storage/tests/vectors.rs | 12 ++--- zebrad/src/components/mempool/tests/prop.rs | 8 +-- 4 files changed, 45 insertions(+), 34 deletions(-) diff --git a/zebrad/src/components/mempool.rs b/zebrad/src/components/mempool.rs index fbbeaab3..e777dcbf 100644 --- a/zebrad/src/components/mempool.rs +++ b/zebrad/src/components/mempool.rs @@ -181,7 +181,7 @@ impl Mempool { txid: UnminedTxId, ) -> Result<(), MempoolError> { // Check if the transaction is already in the mempool. - if storage.contains(&txid) { + if storage.contains_transaction_exact(&txid) { return Err(MempoolError::InMempool); } if let Some(error) = storage.rejection_error(&txid) { @@ -258,12 +258,12 @@ impl Service for Mempool { tx_downloads, } => match req { Request::TransactionIds => { - let res = storage.tx_ids(); + let res = storage.tx_ids().collect(); async move { Ok(Response::TransactionIds(res)) }.boxed() } Request::TransactionsById(ids) => { - let rsp = Ok(storage.transactions(ids)).map(Response::Transactions); - async move { rsp }.boxed() + let res = storage.transactions_exact(ids).cloned().collect(); + async move { Ok(Response::Transactions(res)) }.boxed() } Request::RejectedTransactionIds(ids) => { let rsp = Ok(storage.rejected_transactions(ids)) @@ -314,7 +314,7 @@ fn remove_expired_transactions( ) { let mut txid_set = HashSet::new(); - for t in storage.transactions_all() { + for t in storage.transactions() { if let Some(expiry_height) = t.transaction.expiry_height() { if tip_height >= expiry_height { txid_set.insert(t.id.mined_id()); diff --git a/zebrad/src/components/mempool/storage.rs b/zebrad/src/components/mempool/storage.rs index d3115071..1a7540b8 100644 --- a/zebrad/src/components/mempool/storage.rs +++ b/zebrad/src/components/mempool/storage.rs @@ -120,12 +120,6 @@ impl Storage { Ok(tx_id) } - /// Returns `true` if a [`UnminedTx`] matching an [`UnminedTxId`] is in - /// the mempool. - pub fn contains(&self, txid: &UnminedTxId) -> bool { - self.verified.iter().any(|tx| &tx.id == txid) - } - /// Remove [`UnminedTx`]es from the mempool via exact [`UnminedTxId`]. /// /// For v5 transactions, transactions are matched by WTXID, using both the: @@ -172,22 +166,39 @@ impl Storage { } /// Returns the set of [`UnminedTxId`]s in the mempool. - pub fn tx_ids(&self) -> Vec { - self.verified.iter().map(|tx| tx.id).collect() - } - - /// Returns the set of [`Transaction`]s matching `tx_ids` in the mempool. - pub fn transactions(&self, tx_ids: HashSet) -> Vec { - self.verified - .iter() - .filter(|tx| tx_ids.contains(&tx.id)) - .cloned() - .collect() + pub fn tx_ids(&self) -> impl Iterator + '_ { + self.verified.iter().map(|tx| tx.id) } /// Returns the set of [`Transaction`]s in the mempool. - pub fn transactions_all(&self) -> Vec { - self.verified.iter().cloned().collect() + pub fn transactions(&self) -> impl Iterator { + self.verified.iter() + } + + /// Returns the number of transactions in the mempool. + #[allow(dead_code)] + pub fn transaction_count(&self) -> usize { + self.verified.len() + } + + /// Returns the set of [`Transaction`]s with exactly matching `tx_ids` in the mempool. + /// + /// This matches the exact transaction, with identical blockchain effects, signatures, and proofs. + pub fn transactions_exact( + &self, + tx_ids: HashSet, + ) -> impl Iterator { + self.verified + .iter() + .filter(move |tx| tx_ids.contains(&tx.id)) + } + + /// Returns `true` if a [`UnminedTx`] exactly matching an [`UnminedTxId`] is in + /// the mempool. + /// + /// This matches the exact transaction, with identical blockchain effects, signatures, and proofs. + pub fn contains_transaction_exact(&self, txid: &UnminedTxId) -> bool { + self.verified.iter().any(|tx| &tx.id == txid) } /// Returns `true` if a [`UnminedTx`] matching the supplied [`UnminedTxId`] is in diff --git a/zebrad/src/components/mempool/storage/tests/vectors.rs b/zebrad/src/components/mempool/storage/tests/vectors.rs index dd17e6f1..31b2ba23 100644 --- a/zebrad/src/components/mempool/storage/tests/vectors.rs +++ b/zebrad/src/components/mempool/storage/tests/vectors.rs @@ -24,14 +24,14 @@ fn mempool_storage_crud_exact_mainnet() { let _ = storage.insert(unmined_tx.clone()); // Check that it is in the mempool, and not rejected. - assert!(storage.contains(&unmined_tx.id)); + assert!(storage.contains_transaction_exact(&unmined_tx.id)); // Remove tx let removal_count = storage.remove_exact(&iter::once(unmined_tx.id).collect()); // Check that it is /not/ in the mempool. assert_eq!(removal_count, 1); - assert!(!storage.contains(&unmined_tx.id)); + assert!(!storage.contains_transaction_exact(&unmined_tx.id)); } #[test] @@ -52,7 +52,7 @@ fn mempool_storage_crud_same_effects_mainnet() { let _ = storage.insert(unmined_tx.clone()); // Check that it is in the mempool, and not rejected. - assert!(storage.contains(&unmined_tx.id)); + assert!(storage.contains_transaction_exact(&unmined_tx.id)); // Remove tx let removal_count = @@ -60,7 +60,7 @@ fn mempool_storage_crud_same_effects_mainnet() { // Check that it is /not/ in the mempool. assert_eq!(removal_count, 1); - assert!(!storage.contains(&unmined_tx.id)); + assert!(!storage.contains_transaction_exact(&unmined_tx.id)); } #[test] @@ -103,12 +103,12 @@ fn mempool_storage_basic_for_network(network: Network) -> Result<()> { // Make sure the last MEMPOOL_SIZE transactions we sent are in the verified for tx in expected_in_mempool { - assert!(storage.contains(&tx.id)); + assert!(storage.contains_transaction_exact(&tx.id)); } // Anything greater should not be in the verified for tx in expected_to_be_rejected { - assert!(!storage.contains(&tx.id)); + assert!(!storage.contains_transaction_exact(&tx.id)); } // Query all the ids we have for rejected, get back `total - MEMPOOL_SIZE` diff --git a/zebrad/src/components/mempool/tests/prop.rs b/zebrad/src/components/mempool/tests/prop.rs index f863040b..a324ecd0 100644 --- a/zebrad/src/components/mempool/tests/prop.rs +++ b/zebrad/src/components/mempool/tests/prop.rs @@ -57,7 +57,7 @@ proptest! { // The first call to `poll_ready` shouldn't clear the storage yet. mempool.dummy_call().await; - prop_assert_eq!(mempool.storage().tx_ids().len(), 1); + prop_assert_eq!(mempool.storage().transaction_count(), 1); // Simulate a chain reset. chain_tip_sender.set_finalized_tip(chain_tip); @@ -65,7 +65,7 @@ proptest! { // This time a call to `poll_ready` should clear the storage. mempool.dummy_call().await; - prop_assert!(mempool.storage().tx_ids().is_empty()); + prop_assert_eq!(mempool.storage().transaction_count(), 0); peer_set.expect_no_requests().await?; state_service.expect_no_requests().await?; @@ -110,7 +110,7 @@ proptest! { // The first call to `poll_ready` shouldn't clear the storage yet. mempool.dummy_call().await; - prop_assert_eq!(mempool.storage().tx_ids().len(), 1); + prop_assert_eq!(mempool.storage().transaction_count(), 1); // Simulate the synchronizer catching up to the network chain tip. mempool.disable(&mut recent_syncs).await; @@ -121,7 +121,7 @@ proptest! { // Enable the mempool again so the storage can be accessed. mempool.enable(&mut recent_syncs).await; - prop_assert!(mempool.storage().tx_ids().is_empty()); + prop_assert_eq!(mempool.storage().transaction_count(), 0); peer_set.expect_no_requests().await?; state_service.expect_no_requests().await?;