Limit the size and age of the ZIP-401 rejected transaction ID list

zebrad/proptest-regressions/components/mempool/storage/tests/prop.txt

@@ -5,3 +5,4 @@
 # It is recommended to check this file in to source control so that
 # everyone who runs the test benefits from these saved cases.
 cc b258d507c0b2aef6c2793bdb3fc29e6367e62fba9df378ea6797e9bc97fd2780 # shrinks to input = RemoveSameEffects { transactions: alloc::vec::Vec<zebra_chain::transaction::unmined::UnminedTx><zebra_chain::transaction::unmined::UnminedTx>, len=2, mined_ids_to_remove: std::collections::hash::set::HashSet<zebra_chain::transaction::hash::Hash><zebra_chain::transaction::hash::Hash>, len=2 }
+cc 4616d813ba54e7b7475a1adb880905dfb05a63b59a18dc079893bf963ae92097 # shrinks to rejection_error = SameEffectsChain(Expired), mut rejection_template = Witnessed(WtxId { id: transaction::Hash("0000000000000000000000000000000000000000000000000000000000000000"), auth_digest: AuthDigest("353b92c552d72e06b512c81f909991b291ca0fec5251d6696755091100000000") })

zebrad/src/components/mempool/storage.rs

@@ -10,13 +10,14 @@
 use std::{
     collections::{HashMap, HashSet},
     mem::size_of,
+    time::Duration,
 };
 
 use thiserror::Error;
 
 use zebra_chain::transaction::{self, UnminedTx, UnminedTxId, VerifiedUnminedTx};
 
-use self::verified_set::VerifiedSet;
+use self::{eviction_list::EvictionList, verified_set::VerifiedSet};
 use super::{config, downloads::TransactionDownloadVerifyError, MempoolError};
 
 #[cfg(any(test, feature = "proptest-impl"))]
@@ -25,11 +26,9 @@ use proptest_derive::Arbitrary;
 #[cfg(test)]
 pub mod tests;
 
+mod eviction_list;
 mod verified_set;
 
-/// The maximum number of verified transactions to store in the mempool.
-const MEMPOOL_SIZE: usize = 4;
-
 /// The size limit for mempool transaction rejection lists.
 ///
 /// > The size of RecentlyEvicted SHOULD never exceed `eviction_memory_entries` entries,
@@ -103,7 +102,6 @@ pub enum RejectionError {
 }
 
 /// Hold mempool verified and rejected mempool transactions.
-#[derive(Default)]
 pub struct Storage {
     /// The set of verified transactions in the mempool. This is a
     /// cache of size [`MEMPOOL_SIZE`].
@@ -125,8 +123,18 @@ pub struct Storage {
     /// These rejections apply until a rollback or network upgrade.
     ///
     /// Any transaction with the same `transaction::Hash` is invalid.
-    chain_rejected_same_effects:
-        HashMap<SameEffectsChainRejectionError, HashSet<transaction::Hash>>,
+    ///
+    /// An [`EvictionList`] is used for both randomly evicted and expired transactions,
+    /// even if it is only needed for the evicted ones. This was done just to simplify
+    /// the existing code; there is no harm in having a timeout for expired transactions
+    /// too since re-checking expired transactions is cheap.
+    // If this code is ever refactored and the lists are split in different fields,
+    // then we can use an `EvictionList` just for the evicted list.
+    chain_rejected_same_effects: HashMap<SameEffectsChainRejectionError, EvictionList>,
+
+    /// The mempool transaction eviction age limit.
+    /// Same as [`Config::eviction_memory_time`].
+    eviction_memory_time: Duration,
 
     /// Max total cost of the verified mempool set, beyond which transactions
     /// are evicted to make room.
@@ -142,9 +150,14 @@ impl Drop for Storage {
 impl Storage {
     #[allow(clippy::field_reassign_with_default)]
     pub(crate) fn new(config: &config::Config) -> Self {
-        let mut default: Storage = Default::default();
-        default.tx_cost_limit = config.tx_cost_limit;
-        default
+        Self {
+            tx_cost_limit: config.tx_cost_limit,
+            eviction_memory_time: config.eviction_memory_time,
+            verified: Default::default(),
+            tip_rejected_exact: Default::default(),
+            tip_rejected_same_effects: Default::default(),
+            chain_rejected_same_effects: Default::default(),
+        }
     }
 
     /// Insert a [`VerifiedUnminedTx`] into the mempool, caching any rejections.
@@ -182,15 +195,26 @@ impl Storage {
             result = Err(rejection_error.into());
         }
 
-        // Once inserted, we evict transactions over the pool size limit.
-        while self.verified.transaction_count() > MEMPOOL_SIZE
-            || self.verified.total_cost() > self.tx_cost_limit
-        {
+        // Once inserted, we evict transactions over the pool size limit per [ZIP-401];
+        //
+        // > On receiving a transaction: (...)
+        // > Calculate its cost. If the total cost of transactions in the mempool including this
+        // > one would `exceed mempooltxcostlimit`, then the node MUST repeatedly call
+        // > EvictTransaction (with the new transaction included as a candidate to evict) until the
+        // > total cost does not exceed `mempooltxcostlimit`.
+        //
+        // [ZIP-401]: https://zips.z.cash/zip-0401
+        while self.verified.total_cost() > self.tx_cost_limit {
+            // > EvictTransaction MUST do the following:
+            // > Select a random transaction to evict, with probability in direct proportion to
+            // > eviction weight. (...) Remove it from the mempool.
             let victim_tx = self
                 .verified
                 .evict_one()
                 .expect("mempool is empty, but was expected to be full");
 
+            // > Add the txid and the current time to RecentlyEvicted, dropping the oldest entry in
+            // > RecentlyEvicted if necessary to keep it to at most `eviction_memory_entries entries`.
             self.reject(
                 victim_tx.transaction.id,
                 SameEffectsChainRejectionError::RandomlyEvicted.into(),
@@ -203,8 +227,6 @@ impl Storage {
             }
         }
 
-        assert!(self.verified.transaction_count() <= MEMPOOL_SIZE);
-
         result
     }
 
@@ -277,11 +299,7 @@ impl Storage {
         if self.tip_rejected_same_effects.len() > MAX_EVICTION_MEMORY_ENTRIES {
             self.tip_rejected_same_effects.clear();
         }
-        for (_, map) in self.chain_rejected_same_effects.iter_mut() {
-            if map.len() > MAX_EVICTION_MEMORY_ENTRIES {
-                map.clear();
-            }
-        }
+        // `chain_rejected_same_effects` limits its size by itself
         self.update_rejected_metrics();
     }
 
@@ -326,12 +344,12 @@ impl Storage {
     ///
     /// Transactions on multiple rejected lists are counted multiple times.
     #[allow(dead_code)]
-    pub fn rejected_transaction_count(&self) -> usize {
+    pub fn rejected_transaction_count(&mut self) -> usize {
         self.tip_rejected_exact.len()
             + self.tip_rejected_same_effects.len()
             + self
                 .chain_rejected_same_effects
-                .iter()
+                .iter_mut()
                 .map(|(_, map)| map.len())
                 .sum::<usize>()
     }
@@ -346,12 +364,12 @@ impl Storage {
                 self.tip_rejected_same_effects.insert(txid.mined_id(), e);
             }
             RejectionError::SameEffectsChain(e) => {
-                // TODO: track evicted victims times, removing those older than
-                // config.eviction_memory_time, as well as FIFO more than
-                // MAX_EVICTION_MEMORY_ENTRIES
+                let eviction_memory_time = self.eviction_memory_time;
                 self.chain_rejected_same_effects
                     .entry(e)
-                    .or_default()
+                    .or_insert_with(|| {
+                        EvictionList::new(MAX_EVICTION_MEMORY_ENTRIES, eviction_memory_time)
+                    })
                     .insert(txid.mined_id());
             }
         }
@@ -374,7 +392,7 @@ impl Storage {
         }
 
         for (error, set) in self.chain_rejected_same_effects.iter() {
-            if set.contains(&txid.mined_id()) {
+            if set.contains_key(&txid.mined_id()) {
                 return Some(error.clone().into());
             }
         }
@@ -486,7 +504,7 @@ impl Storage {
     /// Update metrics related to the rejected lists.
     ///
     /// Must be called every time the rejected lists change.
-    fn update_rejected_metrics(&self) {
+    fn update_rejected_metrics(&mut self) {
         metrics::gauge!(
             "mempool.rejected.transaction.ids",
             self.rejected_transaction_count() as _

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

@@ -0,0 +1,144 @@
+//! [`EvictionList`] represents the transaction eviction list with
+//! efficient operations.
+use std::{
+    collections::{HashMap, VecDeque},
+    time::{Duration, Instant},
+};
+
+use zebra_chain::transaction;
+
+/// An eviction list that allows Zebra to efficiently add entries, get entries,
+/// and remove older entries in the order they were inserted.
+pub struct EvictionList {
+    // Maps each TXID in the list to the most recent instant they were added.
+    unique_entries: HashMap<transaction::Hash, Instant>,
+    // The same as `unique_entries` but in the order they were inserted.
+    ordered_entries: VecDeque<transaction::Hash>,
+    // The maximum size of `unique_entries`.
+    max_size: usize,
+    /// The mempool transaction eviction age limit.
+    /// Same as [`Config::eviction_memory_time`].
+    eviction_memory_time: Duration,
+}
+
+impl EvictionList {
+    /// Create a new [`EvictionList`] with the given maximum size and
+    /// eviction time.
+    pub fn new(max_size: usize, eviction_memory_time: Duration) -> Self {
+        Self {
+            unique_entries: Default::default(),
+            ordered_entries: Default::default(),
+            max_size,
+            eviction_memory_time,
+        }
+    }
+
+    /// Inserts a TXID in the list, keeping track of the time it was inserted.
+    ///
+    /// All entries older than [`EvictionList::eviction_memory_time`] will be removed.
+    ///
+    /// # Panics
+    ///
+    /// If the TXID is already in the list.
+    ///
+    pub fn insert(&mut self, key: transaction::Hash) {
+        // From https://zips.z.cash/zip-0401#specification:
+        // > Nodes SHOULD remove transactions from RecentlyEvicted that were evicted more than
+        // > mempoolevictionmemoryminutes minutes ago. This MAY be done periodically,
+        // > and/or just before RecentlyEvicted is accessed when receiving a transaction.
+        self.prune_old();
+        // > Add the txid and the current time to RecentlyEvicted, dropping the oldest entry
+        // > in RecentlyEvicted if necessary to keep it to at most eviction_memory_entries entries.
+        if self.len() >= self.max_size {
+            self.pop_front();
+        }
+        let value = Instant::now();
+        let old_value = self.unique_entries.insert(key, value);
+        // It should be impossible for an already-evicted transaction to be evicted
+        // again since transactions are not added to the mempool if they are evicted,
+        // and the mempool doesn't allow inserting two transactions with the same
+        // hash (they would conflict).
+        assert_eq!(
+            old_value, None,
+            "an already-evicted transaction should not be evicted again"
+        );
+        self.ordered_entries.push_back(key)
+    }
+
+    /// Checks if the given TXID is in the list.
+    pub fn contains_key(&self, txid: &transaction::Hash) -> bool {
+        if let Some(evicted_at) = self.unique_entries.get(txid) {
+            // Since the list is pruned only in mutable functions, make sure
+            // we take expired items into account.
+            if !self.has_expired(evicted_at) {
+                return true;
+            }
+        }
+        false
+    }
+
+    /// Get the size of the list.
+    //
+    // Note: if this method being mutable becomes an issue, it's possible
+    // to compute the number of expired transactions and subtract,
+    // at the cost of `O(len + expired)` performance each time the method is called.
+    //
+    // Currently the performance is `O(expired)` for the first call, then `O(1)` until the next expiry.
+    pub fn len(&mut self) -> usize {
+        self.prune_old();
+        self.unique_entries.len()
+    }
+
+    /// Clear the list.
+    #[allow(dead_code)]
+    pub fn clear(&mut self) {
+        self.unique_entries.clear();
+        self.ordered_entries.clear();
+    }
+
+    /// Prune TXIDs that are older than `eviction_time` ago.
+    ///
+    // This method is public because ZIP-401 states about pruning:
+    // > This MAY be done periodically,
+    pub fn prune_old(&mut self) {
+        while let Some(txid) = self.front() {
+            let evicted_at = self
+                .unique_entries
+                .get(txid)
+                .unwrap_or_else(|| panic!("all entries should exist in both ordered_entries and unique_entries, missing {:?} in unique_entries", txid));
+            if self.has_expired(evicted_at) {
+                self.pop_front();
+            } else {
+                break;
+            }
+        }
+    }
+
+    /// Get the oldest TXID in the list.
+    fn front(&self) -> Option<&transaction::Hash> {
+        self.ordered_entries.front()
+    }
+
+    /// Removes the first element and returns it, or `None` if the `EvictionList`
+    /// is empty.
+    fn pop_front(&mut self) -> Option<transaction::Hash> {
+        if let Some(key) = self.ordered_entries.pop_front() {
+            let removed = self.unique_entries.remove(&key);
+            assert!(
+                removed.is_some(),
+                "all entries should exist in both ordered_entries and unique_entries, missing {:?} in unique_entries",
+                key
+            );
+            Some(key)
+        } else {
+            None
+        }
+    }
+
+    /// Returns if `evicted_at` is considered expired considering the current
+    /// time and the configured eviction time.
+    fn has_expired(&self, evicted_at: &Instant) -> bool {
+        let now = Instant::now();
+        (now - *evicted_at) > self.eviction_memory_time
+    }
+}