Fix the performance issue caused by track_entry_statics

test/src/specs/rpc/truncate.rs

@@ -31,6 +31,7 @@ impl Spec for RpcTruncate {
 
         let tx_pool_info = node.get_tip_tx_pool_info();
         assert!(tx_pool_info.total_tx_size.value() > 0, "tx-pool holds tx2");
+        assert!(tx_pool_info.pending.value() > 0, "tx-pool hods tx2");
 
         // Truncate from `to_truncate`
 

tx-pool/src/component/pool_map.rs

@@ -19,7 +19,6 @@ use ckb_types::{
 };
 use multi_index_map::MultiIndexMap;
 use std::collections::HashSet;
-
 type ConflictEntry = (TxEntry, Reject);
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -71,6 +70,9 @@ pub struct PoolMap {
     pub(crate) total_tx_size: usize,
     // sum of all tx_pool tx's cycles.
     pub(crate) total_tx_cycles: Cycle,
+    pub(crate) pending_count: usize,
+    pub(crate) gap_count: usize,
+    pub(crate) proposed_count: usize,
 }
 
 impl PoolMap {
@@ -82,6 +84,9 @@ impl PoolMap {
             max_ancestors_count,
             total_tx_size: 0,
             total_tx_cycles: 0,
+            pending_count: 0,
+            gap_count: 0,
+            proposed_count: 0,
         }
     }
 
@@ -139,17 +144,12 @@ impl PoolMap {
         self.get_by_id(id).expect("unconsistent pool")
     }
 
-    pub(crate) fn get_by_status(&self, status: Status) -> Vec<&PoolEntry> {
-        self.entries.get_by_status(&status)
-    }
-
     pub(crate) fn pending_size(&self) -> usize {
-        self.entries.get_by_status(&Status::Pending).len()
-            + self.entries.get_by_status(&Status::Gap).len()
+        self.pending_count + self.gap_count
     }
 
     pub(crate) fn proposed_size(&self) -> usize {
-        self.entries.get_by_status(&Status::Proposed).len()
+        self.proposed_count
     }
 
     pub(crate) fn sorted_proposed_iter(&self) -> impl Iterator<Item = &TxEntry> {
@@ -213,19 +213,21 @@ impl PoolMap {
         self.record_entry_edges(&entry)?;
         self.insert_entry(&entry, status);
         self.record_entry_descendants(&entry);
-        self.track_entry_statics();
+        self.track_entry_statics(None, Some(status));
         self.update_stat_for_add_tx(entry.size, entry.cycles);
         Ok((true, evicts))
     }
 
     /// Change the status of the entry, only used for `gap_rtx` and `proposed_rtx`
     pub(crate) fn set_entry(&mut self, short_id: &ProposalShortId, status: Status) {
+        let mut old_status = None;
         self.entries
             .modify_by_id(short_id, |e| {
+                old_status = Some(e.status);
                 e.status = status;
             })
             .expect("unconsistent pool");
-        self.track_entry_statics();
+        self.track_entry_statics(old_status, Some(status));
     }
 
     pub(crate) fn remove_entry(&mut self, id: &ProposalShortId) -> Option<TxEntry> {
@@ -239,6 +241,7 @@ impl PoolMap {
             self.update_descendants_index_key(&entry.inner, EntryOp::Remove);
             self.remove_entry_edges(&entry.inner);
             self.remove_entry_links(id);
+            self.track_entry_statics(Some(entry.status), None);
             self.update_stat_for_remove_tx(entry.inner.size, entry.inner.cycles);
             entry.inner
         })
@@ -362,6 +365,9 @@ impl PoolMap {
         self.links.clear();
         self.total_tx_size = 0;
         self.total_tx_cycles = 0;
+        self.pending_count = 0;
+        self.gap_count = 0;
+        self.proposed_count = 0;
     }
 
     pub(crate) fn score_sorted_iter_by(
@@ -625,20 +631,33 @@ impl PoolMap {
         });
     }
 
-    fn track_entry_statics(&self) {
+    fn track_entry_statics(&mut self, remove: Option<Status>, add: Option<Status>) {
+        match remove {
+            Some(Status::Pending) => self.pending_count -= 1,
+            Some(Status::Gap) => self.gap_count -= 1,
+            Some(Status::Proposed) => self.proposed_count -= 1,
+            _ => {}
+        }
+        match add {
+            Some(Status::Pending) => self.pending_count += 1,
+            Some(Status::Gap) => self.gap_count += 1,
+            Some(Status::Proposed) => self.proposed_count += 1,
+            _ => {}
+        }
+        assert_eq!(
+            self.pending_count + self.gap_count + self.proposed_count,
+            self.entries.len()
+        );
         if let Some(metrics) = ckb_metrics::handle() {
             metrics
                 .ckb_tx_pool_entry
                 .pending
-                .set(self.entries.get_by_status(&Status::Pending).len() as i64);
-            metrics
-                .ckb_tx_pool_entry
-                .gap
-                .set(self.entries.get_by_status(&Status::Gap).len() as i64);
+                .set(self.pending_count as i64);
+            metrics.ckb_tx_pool_entry.gap.set(self.gap_count as i64);
             metrics
                 .ckb_tx_pool_entry
                 .proposed
-                .set(self.proposed_size() as i64);
+                .set(self.proposed_count as i64);
         }
     }
 

tx-pool/src/component/tests/proposed.rs

@@ -3,6 +3,10 @@ use crate::component::tests::util::{
     build_tx, build_tx_with_dep, build_tx_with_header_dep, DEFAULT_MAX_ANCESTORS_COUNT,
     MOCK_CYCLES, MOCK_FEE, MOCK_SIZE,
 };
+use ckb_types::core::{capacity_bytes, ScriptHashType};
+use ckb_types::packed::{CellOutputBuilder, ScriptBuilder};
+use ckb_types::H256;
+use std::time::Instant;
 
 use crate::component::{entry::TxEntry, pool_map::PoolMap};
 use ckb_types::{
@@ -732,3 +736,65 @@ fn test_container_bench_add_limits() {
     pool.clear();
     assert_eq!(pool.size(), 0);
 }
+
+#[test]
+fn test_pool_map_bench() {
+    use rand::Rng;
+    let mut rng = rand::thread_rng();
+
+    let mut pool = PoolMap::new(150);
+
+    let mut instant = Instant::now();
+    let mut time_spend = vec![];
+    for i in 0..20000 {
+        let lock_script1 = ScriptBuilder::default()
+            .code_hash(H256(rand::random()).pack())
+            .hash_type(ScriptHashType::Data.into())
+            .args(Bytes::from(b"lock_script1".to_vec()).pack())
+            .build();
+
+        let type_script1 = ScriptBuilder::default()
+            .code_hash(H256(rand::random()).pack())
+            .hash_type(ScriptHashType::Data.into())
+            .args(Bytes::from(b"type_script1".to_vec()).pack())
+            .build();
+
+        let tx = TransactionBuilder::default()
+            .output(
+                CellOutputBuilder::default()
+                    .capacity(capacity_bytes!(1000).pack())
+                    .lock(lock_script1)
+                    .type_(Some(type_script1).pack())
+                    .build(),
+            )
+            .output_data(Default::default())
+            .build();
+
+        let entry = TxEntry::dummy_resolve(
+            tx,
+            rng.gen_range(0..1000),
+            Capacity::shannons(200),
+            rng.gen_range(0..1000),
+        );
+        if i % 5000 == 0 && i != 0 {
+            eprintln!("i: {}, time: {:?}", i, instant.elapsed());
+            time_spend.push(instant.elapsed());
+            instant = Instant::now();
+        }
+        let status = if rng.gen_range(0..100) >= 30 {
+            Status::Pending
+        } else {
+            Status::Gap
+        };
+        let _ = pool.add_entry(entry, status);
+    }
+    let first = time_spend[0].as_millis();
+    let last = time_spend.last().unwrap().as_millis();
+    let diff = (last as i128 - first as i128).abs();
+    let expect_diff_range = ((first as f64) * 2.0) as i128;
+    eprintln!(
+        "first: {} last: {}, diff: {}, range: {}",
+        first, last, diff, expect_diff_range
+    );
+    assert!(diff < expect_diff_range);
+}

tx-pool/src/pool.rs

@@ -72,15 +72,6 @@ impl TxPool {
         Arc::clone(&self.snapshot)
     }
 
-    fn get_by_status(&self, status: Status) -> Vec<&PoolEntry> {
-        self.pool_map.get_by_status(status)
-    }
-
-    /// Get tx-pool size
-    pub fn status_size(&self, status: Status) -> usize {
-        self.get_by_status(status).len()
-    }
-
     /// Check whether tx-pool enable RBF
     pub fn enable_rbf(&self) -> bool {
         self.config.min_rbf_rate > self.config.min_fee_rate
@@ -254,6 +245,7 @@ impl TxPool {
     // Expire all transaction (and their dependencies) in the pool.
     pub(crate) fn remove_expired(&mut self, callbacks: &Callbacks) {
         let now_ms = ckb_systemtime::unix_time_as_millis();
+
         let removed: Vec<_> = self
             .pool_map
             .iter()