From ea4b4f77e92b1f9c143baf50858d8584c78abb90 Mon Sep 17 00:00:00 2001
From: chirag-bgh <76247491+chirag-bgh@users.noreply.github.com>
Date: Fri, 24 Mar 2023 23:39:02 +0530
Subject: [PATCH] feat: add pooled_transaction_max (#1971)

---
 crates/net/network/src/transactions.rs |  7 +++----
 crates/transaction-pool/src/lib.rs     | 11 +++++++++++
 crates/transaction-pool/src/traits.rs  | 13 +++++++++++++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/crates/net/network/src/transactions.rs b/crates/net/network/src/transactions.rs
index 4e06c5af6b..7b1ae382f3 100644
--- a/crates/net/network/src/transactions.rs
+++ b/crates/net/network/src/transactions.rs
@@ -368,15 +368,14 @@ where
 
                     let mut msg_builder = PooledTransactionsHashesBuilder::new(version);
 
-                    let pooled_txs = self.pool.pooled_transactions();
+                    let pooled_txs =
+                        self.pool.pooled_transactions_max(NEW_POOLED_TRANSACTION_HASHES_SOFT_LIMIT);
                     if pooled_txs.is_empty() {
                         // do not send a message if there are no transactions in the pool
                         return
                     }
 
-                    for pooled_tx in
-                        pooled_txs.into_iter().take(NEW_POOLED_TRANSACTION_HASHES_SOFT_LIMIT)
-                    {
+                    for pooled_tx in pooled_txs.into_iter() {
                         peer.transactions.insert(*pooled_tx.hash());
                         msg_builder.push_pooled(pooled_tx);
                     }
diff --git a/crates/transaction-pool/src/lib.rs b/crates/transaction-pool/src/lib.rs
index 1d2a605e08..e9bdfddc29 100644
--- a/crates/transaction-pool/src/lib.rs
+++ b/crates/transaction-pool/src/lib.rs
@@ -283,6 +283,17 @@ where
         self.pool.pooled_transactions()
     }
 
+    fn pooled_transaction_hashes_max(&self, max: usize) -> Vec<TxHash> {
+        self.pooled_transaction_hashes().into_iter().take(max).collect()
+    }
+
+    fn pooled_transactions_max(
+        &self,
+        max: usize,
+    ) -> Vec<Arc<ValidPoolTransaction<Self::Transaction>>> {
+        self.pooled_transactions().into_iter().take(max).collect()
+    }
+
     fn best_transactions(
         &self,
     ) -> Box<dyn BestTransactions<Item = Arc<ValidPoolTransaction<Self::Transaction>>>> {
diff --git a/crates/transaction-pool/src/traits.rs b/crates/transaction-pool/src/traits.rs
index 645a59df2a..7326503eaa 100644
--- a/crates/transaction-pool/src/traits.rs
+++ b/crates/transaction-pool/src/traits.rs
@@ -79,6 +79,11 @@ pub trait TransactionPool: Send + Sync + Clone {
     /// Consumer: P2P
     fn pooled_transaction_hashes(&self) -> Vec<TxHash>;
 
+    /// Returns only the first `max` hashes of transactions in the pool.
+    ///
+    /// Consumer: P2P
+    fn pooled_transaction_hashes_max(&self, max: usize) -> Vec<TxHash>;
+
     /// Returns the _full_ transaction objects all transactions in the pool.
     ///
     /// Note: This returns a `Vec` but should guarantee that all transactions are unique.
@@ -86,6 +91,14 @@ pub trait TransactionPool: Send + Sync + Clone {
     /// Consumer: P2P
     fn pooled_transactions(&self) -> Vec<Arc<ValidPoolTransaction<Self::Transaction>>>;
 
+    /// Returns only the first `max` transactions in the pool.
+    ///
+    /// Consumer: P2P
+    fn pooled_transactions_max(
+        &self,
+        max: usize,
+    ) -> Vec<Arc<ValidPoolTransaction<Self::Transaction>>>;
+
     /// Returns an iterator that yields transactions that are ready for block production.
     ///
     /// Consumer: Block production