Merge pull request #4464 from quake/quake/minor-refactor-tx-pool

chore: minor refactor on tx pool
nervosnetwork · May 19, 2024 · 31e0287 · 31e0287
2 parents c5671a1 + 46d173a
commit 31e0287
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 42 deletions.
diff --git a/tx-pool/src/component/pool_map.rs b/tx-pool/src/component/pool_map.rs
@@ -153,7 +153,7 @@ impl PoolMap {
     }
 
     pub(crate) fn sorted_proposed_iter(&self) -> impl Iterator<Item = &TxEntry> {
-        self.score_sorted_iter_by(vec![Status::Proposed])
+        self.score_sorted_iter_by_status(Status::Proposed)
     }
 
     pub(crate) fn get(&self, id: &ProposalShortId) -> Option<&TxEntry> {
@@ -329,23 +329,19 @@ impl PoolMap {
         conflicts
     }
 
-    // fill proposal txs
-    pub(crate) fn fill_proposals(
+    // find the pending txs sorted by score, and return their proposal short ids
+    pub(crate) fn get_proposals(
         &self,
         limit: usize,
         exclusion: &HashSet<ProposalShortId>,
-        proposals: &mut HashSet<ProposalShortId>,
-        status: Status,
-    ) {
-        for entry in self.score_sorted_iter_by(vec![status]) {
-            if proposals.len() == limit {
-                break;
-            }
-            let id = entry.proposal_short_id();
-            if !exclusion.contains(&id) {
-                proposals.insert(id);
-            }
-        }
+    ) -> HashSet<ProposalShortId> {
+        self.score_sorted_iter_by_status(Status::Pending)
+            .filter_map(|entry| {
+                let id = entry.proposal_short_id();
+                (!exclusion.contains(&id)).then_some(id)
+            })
+            .take(limit)
+            .collect()
     }
 
     pub(crate) fn iter(&self) -> impl Iterator<Item = &PoolEntry> {
@@ -370,15 +366,24 @@ impl PoolMap {
         self.proposed_count = 0;
     }
 
-    pub(crate) fn score_sorted_iter_by(
+    pub(crate) fn score_sorted_iter_by_status(
+        &self,
+        status: Status,
+    ) -> impl Iterator<Item = &TxEntry> {
+        self.entries
+            .iter_by_score()
+            .rev()
+            .filter_map(move |entry| (entry.status == status).then_some(&entry.inner))
+    }
+
+    pub(crate) fn score_sorted_iter_by_statuses(
         &self,
         statuses: Vec<Status>,
     ) -> impl Iterator<Item = &TxEntry> {
         self.entries
             .iter_by_score()
             .rev()
-            .filter(move |entry| statuses.contains(&entry.status))
-            .map(|entry| &entry.inner)
+            .filter_map(move |entry| statuses.contains(&entry.status).then_some(&entry.inner))
     }
 
     fn remove_entry_links(&mut self, id: &ProposalShortId) {

diff --git a/tx-pool/src/component/tests/pending.rs b/tx-pool/src/component/tests/pending.rs
@@ -163,7 +163,7 @@ fn test_remove_entry() {
 }
 
 #[test]
-fn test_fill_proposals() {
+fn test_get_proposals() {
     let mut pool = PoolMap::new(1000);
     let tx1 = build_tx(vec![(&Byte32::zero(), 1), (&h256!("0x1").pack(), 1)], 1);
     let tx2 = build_tx(
@@ -193,30 +193,26 @@ fn test_fill_proposals() {
     let id2 = tx2.proposal_short_id();
     let id3 = tx3.proposal_short_id();
 
-    let mut ret = HashSet::new();
-    pool.fill_proposals(10, &HashSet::new(), &mut ret, Status::Pending);
+    let ret = pool.get_proposals(10, &HashSet::new());
     assert_eq!(
         ret,
         HashSet::from_iter(vec![id1.clone(), id2.clone(), id3.clone()])
     );
 
-    let mut ret = HashSet::new();
-    pool.fill_proposals(1, &HashSet::new(), &mut ret, Status::Pending);
+    let ret = pool.get_proposals(1, &HashSet::new());
     assert_eq!(ret.len(), 1);
 
-    let mut ret = HashSet::new();
-    pool.fill_proposals(2, &HashSet::new(), &mut ret, Status::Pending);
+    let ret = pool.get_proposals(2, &HashSet::new());
     assert_eq!(ret.len(), 2);
 
-    let mut ret = HashSet::new();
     let mut exclusion = HashSet::new();
     exclusion.insert(id2);
-    pool.fill_proposals(2, &exclusion, &mut ret, Status::Pending);
+    let ret = pool.get_proposals(2, &exclusion);
     assert_eq!(ret, HashSet::from_iter(vec![id1, id3]));
 }
 
 #[test]
-fn test_fill_proposals_with_high_score() {
+fn test_get_proposals_with_high_score() {
     let mut pool = PoolMap::new(1000);
     let tx1 = build_tx(vec![(&Byte32::zero(), 1), (&h256!("0x1").pack(), 1)], 1);
     let tx2 = build_tx(
@@ -242,25 +238,21 @@ fn test_fill_proposals_with_high_score() {
     let id2 = tx2.proposal_short_id();
     let id3 = tx3.proposal_short_id();
 
-    let mut ret = HashSet::new();
-    pool.fill_proposals(10, &HashSet::new(), &mut ret, Status::Pending);
+    let ret = pool.get_proposals(10, &HashSet::new());
     assert_eq!(
         ret,
         HashSet::from_iter(vec![id3.clone(), id2.clone(), id1.clone()])
     );
 
-    let mut ret = HashSet::new();
-    pool.fill_proposals(1, &HashSet::new(), &mut ret, Status::Pending);
+    let ret = pool.get_proposals(1, &HashSet::new());
     assert_eq!(ret, HashSet::from_iter(vec![id3.clone()]));
 
-    let mut ret = HashSet::new();
-    pool.fill_proposals(2, &HashSet::new(), &mut ret, Status::Pending);
+    let ret = pool.get_proposals(2, &HashSet::new());
     assert_eq!(ret, HashSet::from_iter(vec![id3.clone(), id2.clone()]));
 
-    let mut ret = HashSet::new();
     let mut exclusion = HashSet::new();
     exclusion.insert(id2);
-    pool.fill_proposals(2, &exclusion, &mut ret, Status::Pending);
+    let ret = pool.get_proposals(2, &exclusion);
     assert_eq!(ret, HashSet::from_iter(vec![id1, id3]));
 }
 

diff --git a/tx-pool/src/pool.rs b/tx-pool/src/pool.rs
@@ -402,10 +402,7 @@ impl TxPool {
         limit: usize,
         exclusion: &HashSet<ProposalShortId>,
     ) -> HashSet<ProposalShortId> {
-        let mut proposals = HashSet::with_capacity(limit);
-        self.pool_map
-            .fill_proposals(limit, exclusion, &mut proposals, Status::Pending);
-        proposals
+        self.pool_map.get_proposals(limit, exclusion)
     }
 
     /// Returns tx from tx-pool or storage corresponding to the id.
@@ -426,7 +423,7 @@ impl TxPool {
     pub(crate) fn get_ids(&self) -> TxPoolIds {
         let pending = self
             .pool_map
-            .score_sorted_iter_by(vec![Status::Pending, Status::Gap])
+            .score_sorted_iter_by_statuses(vec![Status::Pending, Status::Gap])
             .map(|entry| entry.transaction().hash())
             .collect();
 
@@ -442,7 +439,7 @@ impl TxPool {
     pub(crate) fn get_all_entry_info(&self) -> TxPoolEntryInfo {
         let pending = self
             .pool_map
-            .score_sorted_iter_by(vec![Status::Pending, Status::Gap])
+            .score_sorted_iter_by_statuses(vec![Status::Pending, Status::Gap])
             .map(|entry| (entry.transaction().hash(), entry.to_info()))
             .collect();