From 2384e1b0f643f9e0351ffd28d16b1173b49043ac Mon Sep 17 00:00:00 2001
From: Ayelet Zilber <ayelet@starkware.co>
Date: Thu, 18 Jul 2024 11:11:55 +0300
Subject: [PATCH] feat(mempool): get txs return requested txs after replenshing
 queue

---
 crates/mempool/src/mempool.rs           | 19 +++++++++++++++----
 crates/mempool/src/mempool_test.rs      | 19 +++++++++----------
 crates/mempool/src/transaction_queue.rs |  4 ++++
 3 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs
index 5d142a49..0edf0d23 100644
--- a/crates/mempool/src/mempool.rs
+++ b/crates/mempool/src/mempool.rs
@@ -49,11 +49,22 @@ impl Mempool {
     // library.
     pub fn get_txs(&mut self, n_txs: usize) -> MempoolResult<Vec<ThinTransaction>> {
         let mut eligible_txs: Vec<ThinTransaction> = Vec::with_capacity(n_txs);
-        for tx_hash in self.tx_queue.pop_chunk(n_txs) {
-            let tx = self.tx_pool.remove(tx_hash)?;
-            eligible_txs.push(tx);
+        let mut remaining_txs = n_txs;
+
+        while remaining_txs > 0 && !self.tx_queue.is_empty() {
+            let mut popped_txs = Vec::with_capacity(remaining_txs);
+
+            for tx_hash in self.tx_queue.pop_chunk(remaining_txs) {
+                let tx = self.tx_pool.remove(tx_hash)?;
+                popped_txs.push(tx);
+            }
+
+            self.enqueue_next_eligible_txs(&popped_txs);
+
+            remaining_txs -= popped_txs.len();
+
+            eligible_txs.extend(popped_txs);
         }
-        self.enqueue_next_eligible_txs(&eligible_txs);
 
         Ok(eligible_txs)
     }
diff --git a/crates/mempool/src/mempool_test.rs b/crates/mempool/src/mempool_test.rs
index c3d528f7..e934c8d3 100644
--- a/crates/mempool/src/mempool_test.rs
+++ b/crates/mempool/src/mempool_test.rs
@@ -181,20 +181,20 @@ fn test_get_txs_multi_nonce() {
         add_tx_input!(tx_hash: 1, sender_address: "0x0", tx_nonce: 0_u8, account_nonce: 0_u8).tx;
     let tx_address_0_nonce_1 =
         add_tx_input!(tx_hash: 2, sender_address: "0x0", tx_nonce: 1_u8, account_nonce: 0_u8).tx;
+    let tx_address_0_nonce_2 =
+        add_tx_input!(tx_hash: 3, sender_address: "0x0", tx_nonce: 2_u8, account_nonce: 0_u8).tx;
 
     let queue_txs = [TransactionReference::new(&tx_address_0_nonce_0)];
-    let pool_txs = [tx_address_0_nonce_0.clone(), tx_address_0_nonce_1.clone()];
+    let pool_txs =
+        [tx_address_0_nonce_0.clone(), tx_address_0_nonce_1.clone(), tx_address_0_nonce_2.clone()];
     let mut mempool: Mempool = MempoolState::new(pool_txs, queue_txs).into();
 
     // Test.
-    let txs = mempool.get_txs(2).unwrap();
+    let txs = mempool.get_txs(3).unwrap();
 
-    // Assert that the account's next tx was added the queue.
-    // TODO(Ayelet): all transactions should be returned after replenishing.
-    assert_eq!(txs, &[tx_address_0_nonce_0]);
-    let expected_queue_txs = [TransactionReference::new(&tx_address_0_nonce_1)];
-    let expected_pool_txs = [tx_address_0_nonce_1];
-    let expected_mempool_state = MempoolState::new(expected_pool_txs, expected_queue_txs);
+    // Assert: all transactions are returned.
+    assert_eq!(txs, &[tx_address_0_nonce_0, tx_address_0_nonce_1, tx_address_0_nonce_2]);
+    let expected_mempool_state = MempoolState::new([], []);
     expected_mempool_state.assert_eq_mempool_state(&mempool);
 }
 
@@ -393,7 +393,6 @@ fn test_flow_filling_holes(mut mempool: Mempool) {
     add_tx(&mut mempool, &input_address_0_nonce_0);
     let txs = mempool.get_txs(2).unwrap();
 
-    // TODO(Ayelet): all transactions should be returned after replenishing.
     // Assert: all remaining transactions are returned.
-    assert_eq!(txs, &[input_address_0_nonce_0.tx]);
+    assert_eq!(txs, &[input_address_0_nonce_0.tx, input_address_0_nonce_1.tx]);
 }
diff --git a/crates/mempool/src/transaction_queue.rs b/crates/mempool/src/transaction_queue.rs
index 4ae982fc..e95745d1 100644
--- a/crates/mempool/src/transaction_queue.rs
+++ b/crates/mempool/src/transaction_queue.rs
@@ -63,6 +63,10 @@ impl TransactionQueue {
         }
         false
     }
+
+    pub fn is_empty(&self) -> bool {
+        self.queue.is_empty()
+    }
 }
 
 /// Encapsulates a transaction reference to assess its order (i.e., priority).