Make some mempool functions associated with the `mempool::Storage` type (#2872)

* Make some mempool functions associated with the Mempool type

* Move some functions to methods on mempool::Storage

zebrad/src/components/mempool.rs

@@ -48,8 +48,7 @@ pub use storage::{
 pub use storage::tests::unmined_transactions_in_blocks;
 
 use downloads::{
-    Downloads as TxDownloads, Gossip, TransactionDownloadVerifyError, TRANSACTION_DOWNLOAD_TIMEOUT,
+    Downloads as TxDownloads, Gossip, TRANSACTION_DOWNLOAD_TIMEOUT, TRANSACTION_VERIFY_TIMEOUT,
-    TRANSACTION_VERIFY_TIMEOUT,
 };
 
 type Outbound = Buffer<BoxService<zn::Request, zn::Response, zn::BoxError>, zn::Request>;
@@ -234,22 +233,15 @@ impl Mempool {
         }
     }
 
-    /// Check if transaction should be downloaded and/or verified.
+    /// Remove expired transaction ids from a given list of inserted ones.
-    ///
+    fn remove_expired_from_peer_list(
-    /// If it is already in the mempool (or in its rejected list)
+        send_to_peers_ids: &HashSet<UnminedTxId>,
-    /// then it shouldn't be downloaded/verified.
+        expired_transactions: &HashSet<UnminedTxId>,
-    fn should_download_or_verify(
+    ) -> HashSet<UnminedTxId> {
-        storage: &mut storage::Storage,
+        send_to_peers_ids
-        txid: UnminedTxId,
+            .difference(expired_transactions)
-    ) -> Result<(), MempoolError> {
+            .copied()
-        // Check if the transaction is already in the mempool.
+            .collect()
-        if storage.contains_transaction_exact(&txid) {
-            return Err(MempoolError::InMempool);
-        }
-        if let Some(error) = storage.rejection_error(&txid) {
-            return Err(error);
-        }
-        Ok(())
     }
 }
 
@@ -280,7 +272,7 @@ impl Service<Request> for Mempool {
                             }
                         }
                         Err((txid, e)) => {
-                            reject_if_needed(storage, txid, e);
+                            storage.reject_if_needed(txid, e);
                             // TODO: should we also log the result?
                         }
                     };
@@ -307,10 +299,12 @@ impl Service<Request> for Mempool {
 
                 // Remove expired transactions from the mempool.
                 if let Some(tip_height) = self.latest_chain_tip.best_tip_height() {
-                    let expired_transactions = remove_expired_transactions(storage, tip_height);
+                    let expired_transactions = storage.remove_expired_transactions(tip_height);
                     // Remove transactions that are expired from the peers list
-                    send_to_peers_ids =
+                    send_to_peers_ids = Self::remove_expired_from_peer_list(
-                        remove_expired_from_peer_list(&send_to_peers_ids, &expired_transactions);
+                        &send_to_peers_ids,
+                        &expired_transactions,
+                    );
                 }
 
                 // Send transactions that were not rejected nor expired to peers
@@ -355,7 +349,7 @@ impl Service<Request> for Mempool {
                     let rsp: Vec<Result<(), MempoolError>> = gossiped_txs
                         .into_iter()
                         .map(|gossiped_tx| {
-                            Self::should_download_or_verify(storage, gossiped_tx.id())?;
+                            storage.should_download_or_verify(gossiped_tx.id())?;
                             tx_downloads.download_if_needed_and_verify(gossiped_tx)?;
                             Ok(())
                         })
@@ -385,82 +379,3 @@ impl Service<Request> for Mempool {
         }
     }
 }
-
-/// Remove transactions from the mempool if they have not been mined after a specified height.
-///
-/// https://zips.z.cash/zip-0203#specification
-fn remove_expired_transactions(
-    storage: &mut storage::Storage,
-    tip_height: zebra_chain::block::Height,
-) -> HashSet<UnminedTxId> {
-    let mut txid_set = HashSet::new();
-    // we need a separate set, since reject() takes the original unmined ID,
-    // then extracts the mined ID out of it
-    let mut unmined_id_set = HashSet::new();
-
-    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());
-                unmined_id_set.insert(t.id);
-            }
-        }
-    }
-
-    // expiry height is effecting data, so we match by non-malleable TXID
-    storage.remove_same_effects(&txid_set);
-
-    // also reject it
-    for id in unmined_id_set.iter() {
-        storage.reject(*id, SameEffectsChainRejectionError::Expired.into());
-    }
-
-    unmined_id_set
-}
-
-/// Remove expired transaction ids from a given list of inserted ones.
-fn remove_expired_from_peer_list(
-    send_to_peers_ids: &HashSet<UnminedTxId>,
-    expired_transactions: &HashSet<UnminedTxId>,
-) -> HashSet<UnminedTxId> {
-    send_to_peers_ids
-        .difference(expired_transactions)
-        .copied()
-        .collect()
-}
-
-/// Add a transaction that failed download and verification to the rejected list
-/// if needed, depending on the reason for the failure.
-fn reject_if_needed(
-    storage: &mut storage::Storage,
-    txid: UnminedTxId,
-    e: TransactionDownloadVerifyError,
-) {
-    match e {
-        // Rejecting a transaction already in state would speed up further
-        // download attempts without checking the state. However it would
-        // make the reject list grow forever.
-        // TODO: revisit after reviewing the rejected list cleanup criteria?
-        // TODO: if we decide to reject it, then we need to pass the block hash
-        // to State::Confirmed. This would require the zs::Response::Transaction
-        // to include the hash, which would need to be implemented.
-        TransactionDownloadVerifyError::InState |
-        // An unknown error in the state service, better do nothing
-        TransactionDownloadVerifyError::StateError(_) |
-        // Sync has just started. Mempool shouldn't even be enabled, so will not
-        // happen in practice.
-        TransactionDownloadVerifyError::NoTip |
-        // If download failed, do nothing; the crawler will end up trying to
-        // download it again.
-        TransactionDownloadVerifyError::DownloadFailed(_) |
-        // If it was cancelled then a block was mined, or there was a network
-        // upgrade, etc. No reason to reject it.
-        TransactionDownloadVerifyError::Cancelled => {}
-
-        // Consensus verification failed. Reject transaction to avoid
-        // having to download and verify it again just for it to fail again.
-        TransactionDownloadVerifyError::Invalid(e) => {
-            storage.reject(txid, ExactTipRejectionError::FailedVerification(e).into())
-        }
-    }
-}

zebrad/src/components/mempool/storage.rs

@@ -5,7 +5,7 @@ use thiserror::Error;
 use zebra_chain::transaction::{self, UnminedTx, UnminedTxId};
 
 use self::verified_set::VerifiedSet;
-use super::MempoolError;
+use super::{downloads::TransactionDownloadVerifyError, MempoolError};
 
 #[cfg(any(test, feature = "proptest-impl"))]
 use proptest_derive::Arbitrary;
@@ -362,4 +362,85 @@ impl Storage {
     pub fn contains_rejected(&self, txid: &UnminedTxId) -> bool {
         self.rejection_error(txid).is_some()
     }
+
+    /// Add a transaction that failed download and verification to the rejected list
+    /// if needed, depending on the reason for the failure.
+    pub fn reject_if_needed(&mut self, txid: UnminedTxId, e: TransactionDownloadVerifyError) {
+        match e {
+            // Rejecting a transaction already in state would speed up further
+            // download attempts without checking the state. However it would
+            // make the reject list grow forever.
+            //
+            // TODO: revisit after reviewing the rejected list cleanup criteria?
+            // TODO: if we decide to reject it, then we need to pass the block hash
+            // to State::Confirmed. This would require the zs::Response::Transaction
+            // to include the hash, which would need to be implemented.
+            TransactionDownloadVerifyError::InState |
+            // An unknown error in the state service, better do nothing
+            TransactionDownloadVerifyError::StateError(_) |
+            // Sync has just started. Mempool shouldn't even be enabled, so will not
+            // happen in practice.
+            TransactionDownloadVerifyError::NoTip |
+            // If download failed, do nothing; the crawler will end up trying to
+            // download it again.
+            TransactionDownloadVerifyError::DownloadFailed(_) |
+            // If it was cancelled then a block was mined, or there was a network
+            // upgrade, etc. No reason to reject it.
+            TransactionDownloadVerifyError::Cancelled => {}
+
+            // Consensus verification failed. Reject transaction to avoid
+            // having to download and verify it again just for it to fail again.
+            TransactionDownloadVerifyError::Invalid(e) => {
+                self.reject(txid, ExactTipRejectionError::FailedVerification(e).into())
+            }
+        }
+    }
+
+    /// Remove transactions from the mempool if they have not been mined after a
+    /// specified height.
+    ///
+    /// https://zips.z.cash/zip-0203#specification
+    pub fn remove_expired_transactions(
+        &mut self,
+        tip_height: zebra_chain::block::Height,
+    ) -> HashSet<UnminedTxId> {
+        let mut txid_set = HashSet::new();
+        // we need a separate set, since reject() takes the original unmined ID,
+        // then extracts the mined ID out of it
+        let mut unmined_id_set = HashSet::new();
+
+        for t in self.transactions() {
+            if let Some(expiry_height) = t.transaction.expiry_height() {
+                if tip_height >= expiry_height {
+                    txid_set.insert(t.id.mined_id());
+                    unmined_id_set.insert(t.id);
+                }
+            }
+        }
+
+        // expiry height is effecting data, so we match by non-malleable TXID
+        self.remove_same_effects(&txid_set);
+
+        // also reject it
+        for id in unmined_id_set.iter() {
+            self.reject(*id, SameEffectsChainRejectionError::Expired.into());
+        }
+
+        unmined_id_set
+    }
+
+    /// Check if transaction should be downloaded and/or verified.
+    ///
+    /// If it is already in the mempool (or in its rejected list)
+    /// then it shouldn't be downloaded/verified.
+    pub fn should_download_or_verify(&mut self, txid: UnminedTxId) -> Result<(), MempoolError> {
+        // Check if the transaction is already in the mempool.
+        if self.contains_transaction_exact(&txid) {
+            return Err(MempoolError::InMempool);
+        }
+        if let Some(error) = self.rejection_error(&txid) {
+            return Err(error);
+        }
+        Ok(())
+    }
 }

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

@@ -1,6 +1,6 @@
 use std::iter;
 
-use crate::components::mempool::{remove_expired_from_peer_list, remove_expired_transactions};
+use crate::components::mempool::Mempool;
 
 use super::{super::*, unmined_transactions_in_blocks};
 
@@ -180,13 +180,13 @@ fn mempool_expired_basic_for_network(network: Network) -> Result<()> {
     assert!(everything_in_mempool.contains(&tx_id));
 
     // remove_expired_transactions() will return what was removed
-    let expired = remove_expired_transactions(&mut storage, Height(1));
+    let expired = storage.remove_expired_transactions(Height(1));
     assert!(expired.contains(&tx_id));
     let everything_in_mempool: HashSet<UnminedTxId> = storage.tx_ids().collect();
     assert_eq!(everything_in_mempool.len(), 0);
 
     // No transaction will be sent to peers
-    let send_to_peers = remove_expired_from_peer_list(&everything_in_mempool, &expired);
+    let send_to_peers = Mempool::remove_expired_from_peer_list(&everything_in_mempool, &expired);
     assert_eq!(send_to_peers.len(), 0);
 
     Ok(())