Merge branch 'mempool-multimap'

iljakuklic · iljakuklic · commit 8f226c7d592e · 2023-08-11T11:51:37.000+01:00
diff --git a/mempool/src/pool/mod.rs b/mempool/src/pool/mod.rs
@@ -763,9 +763,8 @@ impl<M: MemoryUsageEstimator> Mempool<M> {
         let expired_ids: Vec<_> = self
             .store
             .txs_by_creation_time
-            .values()
-            .flat_map(Deref::deref)
-            .map(|entry_id| self.store.txs_by_id.get(entry_id).expect("entry should exist"))
+            .iter()
+            .map(|(_time, id)| self.store.txs_by_id.get(id).expect("entry should exist"))
             .filter(|entry| {
                 let now = self.clock.get_time();
                 let expired = now.saturating_sub(entry.creation_time()) > self.max_tx_age;
@@ -794,9 +793,8 @@ impl<M: MemoryUsageEstimator> Mempool<M> {
             let removed_id = self
                 .store
                 .txs_by_descendant_score
-                .values()
-                .flat_map(Deref::deref)
-                .copied()
+                .iter()
+                .map(|(_score, entry)| *entry.deref())
                 .next()
                 .expect("pool not empty");
             let removed = self.store.txs_by_id.get(&removed_id).expect("tx with id should exist");
@@ -930,9 +928,8 @@ impl<M: MemoryUsageEstimator> Mempool<M> {
     pub fn get_all(&self) -> Vec<SignedTransaction> {
         self.store
             .txs_by_descendant_score
-            .values()
-            .flat_map(Deref::deref)
-            .map(|id| self.store.get_entry(id).expect("entry").transaction().clone())
+            .iter()
+            .map(|(_score, id)| self.store.get_entry(id).expect("entry").transaction().clone())
             .collect()
     }
 
@@ -972,7 +969,7 @@ impl<M: MemoryUsageEstimator> Mempool<M> {
         let tx_source = TransactionSourceForConnect::for_mempool(&best_index);
 
         let block_timestamp = tx_accumulator.block_timestamp();
-        let tx_id_iter = self.store.txs_by_ancestor_score.values().flat_map(Deref::deref).rev();
+        let tx_id_iter = self.store.txs_by_ancestor_score.iter().map(|(_, id)| id).rev();
         let mut tx_iter = tx_id_iter
             .filter_map(|tx_id| {
                 let tx = self.store.txs_by_id.get(tx_id)?.deref();
@@ -1118,11 +1115,8 @@ impl<M: MemoryUsageEstimator> Mempool<M> {
             .txs_by_descendant_score
             .iter()
             .rev()
-            .find(|(_score, txs)| {
-                total_size += txs
-                    .iter()
-                    .map(|tx_id| self.store.txs_by_id.get(tx_id).map_or(0, |tx| tx.size()))
-                    .sum::<usize>();
+            .find(|(_score, tx_id)| {
+                total_size += self.store.txs_by_id.get(tx_id).map_or(0, |tx| tx.size());
                 (total_size / 1_000_000) >= in_top_x_mb
             })
             .map_or_else(
diff --git a/mempool/src/pool/store/mod.rs b/mempool/src/pool/store/mod.rs
@@ -17,7 +17,7 @@ mod mem_usage;
 
 use std::{
     cmp::Ordering,
-    collections::{btree_map::Entry::Occupied, BTreeMap, BTreeSet},
+    collections::{BTreeMap, BTreeSet},
     ops::Deref,
 };
 
@@ -74,7 +74,7 @@ type StrictDropPolicy = mem_usage::NoOpDropPolicy;
 
 type TrackedMap<K, V> = Tracked<BTreeMap<K, V>>;
 type TrackedSet<K> = Tracked<BTreeSet<K>>;
-type TrackedTxIdMultiMap<K> = TrackedMap<K, TrackedSet<Id<Transaction>>>;
+type TrackedTxIdMultiMap<K> = TrackedSet<(K, Id<Transaction>)>;
 
 #[derive(Debug)]
 pub struct MempoolStore {
@@ -182,9 +182,9 @@ impl MempoolStore {
         }
 
         let expected_size = map_size_deep(&self.txs_by_id)
-            + map_size_deep(&self.txs_by_descendant_score)
-            + map_size_deep(&self.txs_by_ancestor_score)
-            + map_size_deep(&self.txs_by_creation_time)
+            + self.txs_by_descendant_score.indirect_memory_usage()
+            + self.txs_by_ancestor_score.indirect_memory_usage()
+            + self.txs_by_creation_time.indirect_memory_usage()
             + self.spender_txs.indirect_memory_usage()
             + self.txs_by_seq_no.indirect_memory_usage()
             + self.seq_nos_by_tx.indirect_memory_usage();
@@ -194,11 +194,7 @@ impl MempoolStore {
             "Memory size tracker out of sync",
         );
 
-        let entries = self
-            .txs_by_descendant_score
-            .values()
-            .flat_map(|ids| ids.deref())
-            .collect::<Vec<_>>();
+        let entries: Vec<_> = self.txs_by_descendant_score.iter().map(|(_, id)| id).collect();
 
         for id in self.txs_by_id.keys() {
             assert_eq!(
@@ -337,11 +333,8 @@ impl MempoolStore {
         self.add_to_ancestor_score_index(&entry);
         self.mem_tracker.modify(
             &mut self.txs_by_creation_time,
-            |txs_by_creation_time, tracker| {
-                tracker.modify(
-                    txs_by_creation_time.entry(entry.creation_time()).or_default(),
-                    |entry, _| entry.insert(tx_id),
-                );
+            |txs_by_creation_time, _tracker| {
+                txs_by_creation_time.insert((entry.creation_time(), tx_id));
             },
         );
 
@@ -358,11 +351,8 @@ impl MempoolStore {
         self.refresh_ancestors(entry);
         self.mem_tracker.modify(
             &mut self.txs_by_descendant_score,
-            |by_desc_score, tracker| {
-                tracker.modify(
-                    by_desc_score.entry(entry.descendant_score()).or_default(),
-                    |map_entry, _| map_entry.insert(*entry.tx_id()),
-                )
+            |by_desc_score, _tracker| {
+                by_desc_score.insert((entry.descendant_score(), *entry.tx_id()));
             },
         );
     }
@@ -372,55 +362,37 @@ impl MempoolStore {
         // because such children would be orphans.
         // When we implement disconnecting a block, we'll need to clean up the mess we're leaving
         // here.
-        self.mem_tracker.modify(&mut self.txs_by_ancestor_score, |anc_score, tracker| {
-            tracker.modify(
-                anc_score.entry(entry.ancestor_score()).or_default(),
-                |e, _| e.insert(*entry.tx_id()),
-            )
-        });
+        self.mem_tracker
+            .modify(&mut self.txs_by_ancestor_score, |by_anc_score, _tracker| {
+                by_anc_score.insert((entry.ancestor_score(), *entry.tx_id()));
+            });
     }
 
     fn refresh_ancestors(&mut self, entry: &TxMempoolEntry) {
         // Since the ancestors of `entry` have had their descendant score modified, their ordering
         // in txs_by_descendant_score may no longer be correct. We thus remove all ancestors and
         // reinsert them, taking the new, updated fees into account
         let ancestors = entry.unconfirmed_ancestors(self);
-        self.mem_tracker.modify(&mut self.txs_by_descendant_score, |by_ds, tracker| {
-            for entries in by_ds.values_mut() {
-                tracker.modify(entries, |e, _| e.retain(|id| !ancestors.contains(id)));
-            }
+        self.mem_tracker.modify(&mut self.txs_by_descendant_score, |by_ds, _tracker| {
+            by_ds.retain(|(_score, e)| !ancestors.contains(e));
             for ancestor_id in ancestors.0 {
                 let ancestor =
                     self.txs_by_id.get(&ancestor_id).expect("Inconsistent mempool state");
-                tracker.modify(
-                    by_ds.entry(ancestor.descendant_score()).or_default(),
-                    |e, _| e.insert(ancestor_id),
-                );
+                by_ds.insert((ancestor.descendant_score(), ancestor_id));
             }
-
-            by_ds.retain(|_score, txs| !txs.is_empty())
         });
     }
 
     /// refresh descendants with new ancestor scores
     fn refresh_descendants(&mut self, entry: &TxMempoolEntry) {
         let descendants = entry.unconfirmed_descendants(self);
-        self.mem_tracker.modify(&mut self.txs_by_ancestor_score, |by_as, tracker| {
-            for entries in by_as.values_mut() {
-                tracker.modify(entries, |e, _| e.retain(|id| !descendants.contains(id)));
-            }
+        self.mem_tracker.modify(&mut self.txs_by_ancestor_score, |by_as, _tracker| {
+            by_as.retain(|(_score, e)| !descendants.contains(e));
             for descendant_id in descendants.0 {
                 let descendant =
                     self.txs_by_id.get(&descendant_id).expect("Inconsistent mempool state");
-                tracker.modify(
-                    by_as.entry(descendant.ancestor_score()).or_default(),
-                    |e, _| e.insert(descendant_id),
-                );
+                by_as.insert((descendant.ancestor_score(), descendant_id));
             }
-
-            tracker.modify(&mut self.txs_by_descendant_score, |by_ds, _| {
-                by_ds.retain(|_score, txs| !txs.is_empty())
-            });
         })
     }
 
@@ -458,7 +430,7 @@ impl MempoolStore {
             self.drop_tx(&entry);
             Some(entry)
         } else {
-            assert!(!self.txs_by_descendant_score.values().any(|by_ds| by_ds.contains(tx_id)));
+            assert!(!self.txs_by_descendant_score.iter().any(|(_, id)| id == tx_id));
             assert!(!self.spender_txs.iter().any(|(_, id)| *id == *tx_id));
             None
         }
@@ -480,41 +452,21 @@ impl MempoolStore {
 
     fn remove_from_ancestor_score_index(&mut self, entry: &TxMempoolEntry) {
         self.refresh_descendants(entry);
-        self.mem_tracker.modify(&mut self.txs_by_ancestor_score, |by_as, tracker| {
-            let map_entry = by_as.entry(entry.ancestor_score()).and_modify(|entries| {
-                tracker.modify(entries, |e, _| e.remove(entry.tx_id()));
-            });
-
-            match map_entry {
-                Occupied(entries) if entries.get().is_empty() => drop(entries.remove_entry()),
-                _ => (),
-            };
+        self.mem_tracker.modify(&mut self.txs_by_ancestor_score, |by_as, _tracker| {
+            by_as.remove(&(entry.ancestor_score(), *entry.tx_id()));
         })
     }
 
     fn remove_from_descendant_score_index(&mut self, entry: &TxMempoolEntry) {
         self.refresh_ancestors(entry);
-        self.mem_tracker.modify(&mut self.txs_by_descendant_score, |by_ds, tracker| {
-            let map_entry = by_ds.entry(entry.descendant_score()).and_modify(|entries| {
-                tracker.modify(entries, |e, _| e.remove(entry.tx_id()));
-            });
-
-            match map_entry {
-                Occupied(entries) if entries.get().is_empty() => drop(entries.remove_entry()),
-                _ => {}
-            };
+        self.mem_tracker.modify(&mut self.txs_by_descendant_score, |by_ds, _tracker| {
+            by_ds.remove(&(entry.descendant_score(), *entry.tx_id()));
         })
     }
 
     fn remove_from_creation_time_index(&mut self, entry: &TxMempoolEntry) {
-        self.mem_tracker.modify(&mut self.txs_by_creation_time, |by_ct, tracker| {
-            by_ct.entry(entry.creation_time()).and_modify(|entries| {
-                tracker.modify(entries, |e, _| e.remove(entry.tx_id()));
-            });
-
-            if by_ct.get(&entry.creation_time()).expect("key must exist").is_empty() {
-                by_ct.remove(&entry.creation_time());
-            }
+        self.mem_tracker.modify(&mut self.txs_by_creation_time, |by_ct, _tracker| {
+            by_ct.remove(&(entry.creation_time(), *entry.tx_id()));
         })
     }
 
diff --git a/mempool/src/pool/tests/mod.rs b/mempool/src/pool/tests/mod.rs
@@ -1335,7 +1335,6 @@ async fn ancestor_score(#[case] seed: Seed) -> anyhow::Result<()> {
         entry_b.fees_with_ancestors(),
         (entry_b.fee() + entry_tx.fee()).unwrap()
     );
-    assert!(!mempool.store.txs_by_ancestor_score.contains_key(&tx_b_fee.into()));
     log::debug!(
         "BEFORE REMOVAL raw txs_by_ancestor_score {:?}",
         mempool.store.txs_by_ancestor_score
@@ -1364,8 +1363,6 @@ async fn ancestor_score(#[case] seed: Seed) -> anyhow::Result<()> {
         entry_c.ancestor_score()
     );
 
-    assert!(!mempool.store.txs_by_ancestor_score.contains_key(&tx_c_fee.into()));
-
     check_txs_sorted_by_ancestor_score(&mempool);
     mempool.store.assert_valid();
 
@@ -1376,8 +1373,8 @@ fn check_txs_sorted_by_ancestor_score<E>(mempool: &Mempool<E>) {
     let txs_by_ancestor_score = mempool
         .store
         .txs_by_descendant_score
-        .values()
-        .flat_map(Deref::deref)
+        .iter()
+        .map(|(_score, id)| id)
         .collect::<Vec<_>>();
     for i in 0..(txs_by_ancestor_score.len() - 1) {
         log::debug!("i =  {}", i);
@@ -1477,15 +1474,13 @@ async fn descendant_score(#[case] seed: Seed) -> anyhow::Result<()> {
         entry_b.fees_with_descendants(),
         (entry_b.fee() + entry_c.fee()).unwrap()
     );
-    assert!(!mempool.store.txs_by_descendant_score.contains_key(&tx_b_fee.into()));
     log::debug!(
         "raw_txs_by_descendant_score {:?}",
         mempool.store.txs_by_descendant_score
     );
     check_txs_sorted_by_descendant_sore(&mempool);
 
     mempool.store.remove_tx(entry_c.tx_id(), MempoolRemovalReason::Block);
-    assert!(!mempool.store.txs_by_descendant_score.contains_key(&tx_c_fee.into()));
     let entry_b = mempool.store.txs_by_id.get(&tx_b_id).expect("tx_b");
     assert_eq!(entry_b.fees_with_descendants(), entry_b.fee());
 
@@ -1499,8 +1494,8 @@ fn check_txs_sorted_by_descendant_sore<M>(mempool: &Mempool<M>) {
     let txs_by_descendant_score = mempool
         .store
         .txs_by_descendant_score
-        .values()
-        .flat_map(Deref::deref)
+        .iter()
+        .map(|(_score, id)| id)
         .collect::<Vec<_>>();
     for i in 0..(txs_by_descendant_score.len() - 1) {
         log::debug!("i =  {}", i);
diff --git a/mempool/src/pool/tests/orphans.rs b/mempool/src/pool/tests/orphans.rs
@@ -221,7 +221,7 @@ async fn transaction_graph_subset_permutation(#[case] seed: Seed) {
     let mut rng = make_seedable_rng(seed);
 
     // Generate a valid graph of transactions
-    let num_txs = rng.gen_range(15..40);
+    let num_txs = rng.gen_range(15..90);
     let time = TimeGetter::default().get_time();
     let full_tx_sequence: Vec<_> =
         generate_transaction_graph(&mut rng, time).take(num_txs).collect();
@@ -251,6 +251,13 @@ async fn transaction_graph_subset_permutation(#[case] seed: Seed) {
             let _ = mempool.add_transaction(tx, TxOrigin::TEST).expect("tx add");
         });
 
+        log::info!(
+            "Stats: count {}, memory {}, encoded size {}",
+            mempool.store.txs_by_id.len(),
+            mempool.memory_usage(),
+            mempool.store.txs_by_id.values().map(|e| e.size()).sum::<usize>(),
+        );
+
         // Check the final state of each transaction in the original sequence
         results.push(all_tx_ids.iter().map(|id| TxStatus::fetch(&mempool, id)).collect());
     }
