Allow removing a tx from mempool storage

zebrad/src/components/mempool/storage.rs

@@ -1,4 +1,7 @@
-use std::collections::{HashMap, HashSet, VecDeque};
+use std::{
+    collections::{HashMap, HashSet, VecDeque},
+    hash::Hash,
+};
 
 use zebra_chain::{
     block,
@@ -91,6 +94,27 @@ impl Storage {
         self.verified.iter().any(|tx| &tx.id == txid)
     }
 
+    /// Remove a [`UnminedTx`] from the mempool via [`UnminedTxId`].  Returns
+    /// whether the transaction was present.
+    ///
+    /// 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.
+    #[allow(dead_code)]
+    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`.
+
+        return match self.verified.binary_search_by_key(txid, |&tx| tx.id.hash()) {
+            Ok(tx) => {
+                self.verified.retain(|x| &x.id != txid);
+                Some(tx)
+            }
+            Err(e) => None,
+        };
+    }
+
     /// Returns the set of [`UnminedTxId`]s in the mempool.
     pub fn tx_ids(&self) -> Vec<UnminedTxId> {
         self.verified.iter().map(|tx| tx.id).collect()

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

@@ -6,6 +6,34 @@ use zebra_chain::{
 
 use color_eyre::eyre::Result;
 
+#[test]
+fn mempool_storage_crud_mainnet() {
+    zebra_test::init();
+
+    let network = Network::Mainnet;
+
+    // Create an empty storage instance
+    let mut storage: Storage = Default::default();
+
+    // Get transactions from the first 10 blocks of the Zcash blockchain
+    let (total_transactions, unmined_transactions) = unmined_transactions_in_blocks(10, network);
+
+    // Get one (1) unmined transaction
+    let unmined_tx = unmined_transactions[0];
+
+    // Insert unmined tx into the mempool.
+    storage.insert(unmined_tx);
+
+    // Check that it is in the mempool, and not rejected.
+    assert!(storage.contains(&unmined_tx.id));
+
+    // Remove tx
+    let _ = storage.remove(&unmined_tx.id);
+
+    // Check that it is /not/ in the mempool.
+    assert!(!storage.contains(&unmined_tx.id));
+}
+
 #[test]
 fn mempool_storage_basic() -> Result<()> {
     zebra_test::init();