From 788f144e5b7435d8b8340801f80c57cda8c7c6d5 Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Tue, 11 Feb 2025 16:14:47 -0600
Subject: [PATCH] bound by map len, add test

---
 .../transaction_state_container.rs            | 27 ++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/core/src/banking_stage/transaction_scheduler/transaction_state_container.rs b/core/src/banking_stage/transaction_scheduler/transaction_state_container.rs
index f57b4a10350c60..ceca4b5a7e26a8 100644
--- a/core/src/banking_stage/transaction_scheduler/transaction_state_container.rs
+++ b/core/src/banking_stage/transaction_scheduler/transaction_state_container.rs
@@ -140,7 +140,15 @@ impl<Tx: TransactionWithMeta> StateContainer<Tx> for TransactionStateContainer<T
         for id in priority_ids {
             self.priority_queue.push(id);
         }
-        let num_dropped = self.priority_queue.len().saturating_sub(self.capacity);
+
+        // The number of items in the `id_to_transaction_state` map is
+        // greater than or equal to the number of elements in the queue.
+        // To avoid the map going over capacity, we use the length of the
+        // map here instead of the queue.
+        let num_dropped = self
+            .id_to_transaction_state
+            .len()
+            .saturating_sub(self.capacity);
 
         for _ in 0..num_dropped {
             let priority_id = self.priority_queue.pop_min().expect("queue is not empty");
@@ -475,5 +483,22 @@ mod tests {
         assert_eq!(container.pop().unwrap().priority, 12);
         assert_eq!(container.pop().unwrap().priority, 11);
         assert!(container.pop().is_none());
+
+        // Container now has no items in the queue, but still has 5 items in the map.
+        // If we attempt to push additional transactions to the queue, they
+        // are rejected regardless of their priority.
+        let priority = u64::MAX;
+        let (_transaction_ttl, packet, priority, cost) = test_transaction(priority);
+        let id = container
+            .try_insert_map_only_with_data(packet.original_packet().data(..).unwrap(), |data| {
+                packet_parser(data, priority, cost)
+            })
+            .unwrap();
+        let priority_id = TransactionPriorityId::new(priority, id);
+        assert_eq!(
+            container.push_ids_into_queue(std::iter::once(priority_id)),
+            1
+        );
+        assert!(container.pop().is_none());
     }
 }