Remove transactions in newly committed blocks from the mempool (#2827)

* Add a mempool transaction removal method for mined IDs

And use this method to remove expired transactions,
because all transactions with the same mined ID expire at the same height.

* Remove mined transaction IDs from the mempool

Co-authored-by: Conrado Gouvea <conrado@zfnd.org>

zebrad/src/components/mempool.rs

@@ -210,11 +210,12 @@ impl Service<Request> for Mempool {
                             storage.clear();
                             tx_downloads.cancel_all();
                         }
-                        // Cancel downloads/verifications of transactions with the same
-                        // IDs as recently mined transactions.
+                        // Cancel downloads/verifications/storage of transactions
+                        // with the same mined IDs as recently mined transactions.
                         TipAction::Grow { block } => {
-                            let txid_set = block.transaction_hashes.iter().collect();
-                            tx_downloads.cancel(txid_set);
+                            let mined_ids = block.transaction_hashes.iter().cloned().collect();
+                            tx_downloads.cancel(&mined_ids);
+                            storage.remove_same_effects(&mined_ids);
                         }
                     }
                 }
@@ -309,14 +310,16 @@ fn remove_expired_transactions(
     storage: &mut storage::Storage,
     tip_height: zebra_chain::block::Height,
 ) {
-    let ids = storage.tx_ids().iter().copied().collect();
-    let transactions = storage.transactions(ids);
+    let mut txid_set = HashSet::new();
 
-    for t in transactions {
+    for t in storage.transactions_all() {
         if let Some(expiry_height) = t.transaction.expiry_height() {
             if tip_height >= expiry_height {
-                storage.remove(&t.id);
+                txid_set.insert(t.id.mined_id());
             }
         }
     }
+
+    // expiry height is effecting data, so we match by non-malleable TXID
+    storage.remove_same_effects(&txid_set);
 }

zebrad/src/components/mempool/downloads.rs

@@ -317,7 +317,7 @@ where
 
     /// Cancel download/verification tasks of transactions with the
     /// given transaction hash (see [`UnminedTxId::mined_id`]).
-    pub fn cancel(&mut self, mined_ids: HashSet<&transaction::Hash>) {
+    pub fn cancel(&mut self, mined_ids: &HashSet<transaction::Hash>) {
         // TODO: this can be simplified with [`HashMap::drain_filter`] which
         // is currently nightly-only experimental API.
         let removed_txids: Vec<UnminedTxId> = self

zebrad/src/components/mempool/storage.rs

@@ -5,7 +5,7 @@ use std::{
 
 use zebra_chain::{
     block,
-    transaction::{Transaction, UnminedTx, UnminedTxId},
+    transaction::{self, Transaction, UnminedTx, UnminedTxId},
 };
 use zebra_consensus::error::TransactionError;
 
@@ -104,24 +104,49 @@ impl Storage {
         self.verified.iter().any(|tx| &tx.id == txid)
     }
 
-    /// Remove a [`UnminedTx`] from the mempool via [`UnminedTxId`].  Returns
-    /// whether the transaction was present.
+    /// Remove [`UnminedTx`]es from the mempool via exact [`UnminedTxId`].
     ///
-    /// Removes from the 'verified' set, does not remove from the 'rejected'
-    /// tracking set, if present. Maintains the order in which the other unmined
-    /// transactions have been inserted into the mempool.
-    pub fn remove(&mut self, txid: &UnminedTxId) -> Option<UnminedTx> {
-        // If the txid exists in the verified set and is then deleted,
-        // `retain()` removes it and returns `Some(UnminedTx)`. If it's not
-        // present and nothing changes, returns `None`.
+    /// For v5 transactions, transactions are matched by WTXID, using both the:
+    /// - non-malleable transaction ID, and
+    /// - authorizing data hash.
+    ///
+    /// This matches the exact transaction, with identical blockchain effects, signatures, and proofs.
+    ///
+    /// Returns the number of transactions which were removed.
+    ///
+    /// Removes from the 'verified' set, if present.
+    /// Maintains the order in which the other unmined transactions have been inserted into the mempool.
+    ///
+    /// Does not add or remove from the 'rejected' tracking set.
+    #[allow(dead_code)]
+    pub fn remove_exact(&mut self, exact_wtxids: &HashSet<UnminedTxId>) -> usize {
+        let original_size = self.verified.len();
 
-        match self.verified.clone().iter().find(|tx| &tx.id == txid) {
-            Some(tx) => {
-                self.verified.retain(|tx| &tx.id != txid);
-                Some(tx.clone())
-            }
-            None => None,
-        }
+        self.verified.retain(|tx| !exact_wtxids.contains(&tx.id));
+
+        original_size - self.verified.len()
+    }
+
+    /// Remove [`UnminedTx`]es from the mempool via non-malleable [`transaction::Hash`].
+    ///
+    /// For v5 transactions, transactions are matched by TXID,
+    /// using only the non-malleable transaction ID.
+    /// This matches any transaction with the same effect on the blockchain state,
+    /// even if its signatures and proofs are different.
+    ///
+    /// Returns the number of transactions which were removed.
+    ///
+    /// Removes from the 'verified' set, if present.
+    /// Maintains the order in which the other unmined transactions have been inserted into the mempool.
+    ///
+    /// Does not add or remove from the 'rejected' tracking set.
+    pub fn remove_same_effects(&mut self, mined_ids: &HashSet<transaction::Hash>) -> usize {
+        let original_size = self.verified.len();
+
+        self.verified
+            .retain(|tx| !mined_ids.contains(&tx.id.mined_id()));
+
+        original_size - self.verified.len()
     }
 
     /// Returns the set of [`UnminedTxId`]s in the mempool.
@@ -129,7 +154,7 @@ impl Storage {
         self.verified.iter().map(|tx| tx.id).collect()
     }
 
-    /// Returns the set of [`Transaction`]s matching ids in the mempool.
+    /// Returns the set of [`Transaction`]s matching `tx_ids` in the mempool.
     pub fn transactions(&self, tx_ids: HashSet<UnminedTxId>) -> Vec<UnminedTx> {
         self.verified
             .iter()
@@ -138,6 +163,11 @@ impl Storage {
             .collect()
     }
 
+    /// Returns the set of [`Transaction`]s in the mempool.
+    pub fn transactions_all(&self) -> Vec<UnminedTx> {
+        self.verified.iter().cloned().collect()
+    }
+
     /// Returns `true` if a [`UnminedTx`] matching an [`UnminedTxId`] is in
     /// the mempool rejected list.
     pub fn contains_rejected(&self, txid: &UnminedTxId) -> bool {

zebrad/src/components/mempool/storage/tests/vectors.rs

@@ -1,3 +1,5 @@
+use std::iter;
+
 use super::{super::*, unmined_transactions_in_blocks};
 
 use zebra_chain::parameters::Network;
@@ -5,7 +7,7 @@ use zebra_chain::parameters::Network;
 use color_eyre::eyre::Result;
 
 #[test]
-fn mempool_storage_crud_mainnet() {
+fn mempool_storage_crud_exact_mainnet() {
     zebra_test::init();
 
     let network = Network::Mainnet;
@@ -25,9 +27,39 @@ fn mempool_storage_crud_mainnet() {
     assert!(storage.contains(&unmined_tx.id));
 
     // Remove tx
-    let _ = storage.remove(&unmined_tx.id);
+    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));
+}
+
+#[test]
+fn mempool_storage_crud_same_effects_mainnet() {
+    zebra_test::init();
+
+    let network = Network::Mainnet;
+
+    // Create an empty storage instance
+    let mut storage: Storage = Default::default();
+
+    // Get one (1) unmined transaction
+    let unmined_tx = unmined_transactions_in_blocks(.., network)
+        .next()
+        .expect("at least one unmined transaction");
+
+    // Insert unmined tx into the mempool.
+    let _ = storage.insert(unmined_tx.clone());
+
+    // Check that it is in the mempool, and not rejected.
+    assert!(storage.contains(&unmined_tx.id));
+
+    // Remove tx
+    let removal_count =
+        storage.remove_same_effects(&iter::once(unmined_tx.id.mined_id()).collect());
+
+    // Check that it is /not/ in the mempool.
+    assert_eq!(removal_count, 1);
     assert!(!storage.contains(&unmined_tx.id));
 }