From 640dd621efd414954c50e62960db26b8a51095db Mon Sep 17 00:00:00 2001
From: zhangsoledad <787953403@qq.com>
Date: Mon, 10 Jan 2022 04:00:30 +0800
Subject: [PATCH] fix: pending pool leak

---
 tx-pool/src/component/pending.rs       | 45 ++++++++++++++++++--------
 tx-pool/src/component/tests/pending.rs | 11 +++++++
 2 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/tx-pool/src/component/pending.rs b/tx-pool/src/component/pending.rs
index 24d3aa5144..463136d882 100644
--- a/tx-pool/src/component/pending.rs
+++ b/tx-pool/src/component/pending.rs
@@ -42,6 +42,26 @@ impl PendingQueue {
         self.inner.len()
     }
 
+    #[cfg(test)]
+    pub(crate) fn outputs_len(&self) -> usize {
+        self.outputs.len()
+    }
+
+    #[cfg(test)]
+    pub(crate) fn header_deps_len(&self) -> usize {
+        self.header_deps.len()
+    }
+
+    #[cfg(test)]
+    pub(crate) fn deps_len(&self) -> usize {
+        self.deps.len()
+    }
+
+    #[cfg(test)]
+    pub(crate) fn inputs_len(&self) -> usize {
+        self.inputs.len()
+    }
+
     pub(crate) fn add_entry(&mut self, entry: TxEntry) -> bool {
         let inputs = entry.transaction().input_pts_iter();
         let tx_short_id = entry.proposal_short_id();
@@ -57,10 +77,9 @@ impl PendingQueue {
                 .or_default()
                 .insert(tx_short_id.clone());
 
-            self.outputs
-                .entry(i.to_owned())
-                .or_default()
-                .insert(tx_short_id.clone());
+            if let Some(outputs) = self.outputs.get_mut(&i) {
+                outputs.insert(tx_short_id.clone());
+            }
         }
 
         // record dep-txid
@@ -70,10 +89,9 @@ impl PendingQueue {
                 .or_default()
                 .insert(tx_short_id.clone());
 
-            self.outputs
-                .entry(d.to_owned())
-                .or_default()
-                .insert(tx_short_id.clone());
+            if let Some(outputs) = self.outputs.get_mut(d) {
+                outputs.insert(tx_short_id.clone());
+            }
         }
 
         // record tx unconsumed output
@@ -189,20 +207,21 @@ impl PendingQueue {
         removed
     }
 
-    pub(crate) fn get_descendants(&self, entry: &TxEntry) -> Vec<ProposalShortId> {
+    pub(crate) fn get_descendants(&self, entry: &TxEntry) -> HashSet<ProposalShortId> {
         let mut entries: VecDeque<&TxEntry> = VecDeque::new();
         entries.push_back(entry);
 
-        let mut descendants = Vec::new();
+        let mut descendants = HashSet::new();
         while let Some(entry) = entries.pop_front() {
             let outputs = entry.transaction().output_pts();
 
             for output in outputs {
                 if let Some(ids) = self.outputs.get(&output) {
-                    descendants.extend(ids.iter().cloned());
                     for id in ids {
-                        if let Some(entry) = self.inner.get(id) {
-                            entries.push_back(entry);
+                        if descendants.insert(id.clone()) {
+                            if let Some(entry) = self.inner.get(id) {
+                                entries.push_back(entry);
+                            }
                         }
                     }
                 }
diff --git a/tx-pool/src/component/tests/pending.rs b/tx-pool/src/component/tests/pending.rs
index cea3f84627..d38ae3c1c0 100644
--- a/tx-pool/src/component/tests/pending.rs
+++ b/tx-pool/src/component/tests/pending.rs
@@ -21,6 +21,9 @@ fn test_basic() {
     assert!(queue.contains_key(&tx1.proposal_short_id()));
     assert!(queue.contains_key(&tx2.proposal_short_id()));
 
+    assert_eq!(queue.inputs_len(), 4);
+    assert_eq!(queue.outputs_len(), 4);
+
     assert_eq!(queue.get(&tx1.proposal_short_id()).unwrap(), &entry1);
     assert_eq!(queue.get_tx(&tx2.proposal_short_id()).unwrap(), &tx2);
 
@@ -112,6 +115,10 @@ fn test_resolve_conflict_header_dep() {
     assert!(queue.add_entry(entry.clone()));
     assert!(queue.add_entry(entry1.clone()));
 
+    assert_eq!(queue.inputs_len(), 3);
+    assert_eq!(queue.header_deps_len(), 1);
+    assert_eq!(queue.outputs_len(), 2);
+
     let mut headers = HashSet::new();
     headers.insert(header);
 
@@ -191,6 +198,10 @@ fn test_fill_proposals() {
     assert!(queue.add_entry(entry2));
     assert!(queue.add_entry(entry3));
 
+    assert_eq!(queue.inputs_len(), 5);
+    assert_eq!(queue.deps_len(), 1);
+    assert_eq!(queue.outputs_len(), 7);
+
     let id1 = tx1.proposal_short_id();
     let id2 = tx2.proposal_short_id();
     let id3 = tx3.proposal_short_id();