From 70806ff543496aa6de7807feff49e7e1370efd20 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]"
 <41898282+github-actions[bot]@users.noreply.github.com>
Date: Tue, 10 Sep 2024 23:48:52 +0000
Subject: [PATCH] [mempool] improve parking lot eviction (#14586) (#14589)

---
 mempool/src/core_mempool/transaction_store.rs |  21 +++-
 mempool/src/counters.rs                       |  31 ++++-
 mempool/src/shared_mempool/tasks.rs           |  10 +-
 mempool/src/tests/common.rs                   |  72 +++++++++++-
 mempool/src/tests/core_mempool_test.rs        | 104 ++++++++++++++++-
 mempool/src/tests/integration_tests.rs        | 106 +++++++++---------
 mempool/src/tests/test_framework.rs           |  33 ++----
 7 files changed, 289 insertions(+), 88 deletions(-)

diff --git a/mempool/src/core_mempool/transaction_store.rs b/mempool/src/core_mempool/transaction_store.rs
index 6f7b1d6547b71..c225a52725707 100644
--- a/mempool/src/core_mempool/transaction_store.rs
+++ b/mempool/src/core_mempool/transaction_store.rs
@@ -357,8 +357,11 @@ impl TransactionStore {
         curr_sequence_number: u64,
     ) -> bool {
         if self.is_full() && self.check_txn_ready(txn, curr_sequence_number) {
-            // try to free some space in Mempool from ParkingLot by evicting a non-ready txn
-            if let Some(txn_pointer) = self.parking_lot_index.get_poppable() {
+            let now = Instant::now();
+            // try to free some space in Mempool from ParkingLot by evicting non-ready txns
+            let mut evicted_txns = 0;
+            let mut evicted_bytes = 0;
+            while let Some(txn_pointer) = self.parking_lot_index.get_poppable() {
                 if let Some(txn) = self
                     .transactions
                     .get_mut(&txn_pointer.sender)
@@ -370,9 +373,23 @@ impl TransactionStore {
                             txn.sequence_info.transaction_sequence_number
                         ))
                     );
+                    evicted_bytes += txn.get_estimated_bytes() as u64;
+                    evicted_txns += 1;
                     self.index_remove(&txn);
+                    if !self.is_full() {
+                        break;
+                    }
+                } else {
+                    error!("Transaction not found in mempool while evicting from parking lot");
+                    break;
                 }
             }
+            if evicted_txns > 0 {
+                counters::CORE_MEMPOOL_PARKING_LOT_EVICTED_COUNT.observe(evicted_txns as f64);
+                counters::CORE_MEMPOOL_PARKING_LOT_EVICTED_BYTES.observe(evicted_bytes as f64);
+                counters::CORE_MEMPOOL_PARKING_LOT_EVICTED_LATENCY
+                    .observe(now.elapsed().as_secs_f64());
+            }
         }
         self.is_full()
     }
diff --git a/mempool/src/counters.rs b/mempool/src/counters.rs
index dba59738223e3..464523419e886 100644
--- a/mempool/src/counters.rs
+++ b/mempool/src/counters.rs
@@ -118,7 +118,7 @@ const RANKING_SCORE_BUCKETS: &[f64] = &[
 
 const TXN_CONSENSUS_PULLED_BUCKETS: &[f64] = &[1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 25.0, 50.0, 100.0];
 
-static TRANSACTION_COUNT_BUCKETS: Lazy<Vec<f64>> = Lazy::new(|| {
+static TXN_COUNT_BUCKETS: Lazy<Vec<f64>> = Lazy::new(|| {
     exponential_buckets(
         /*start=*/ 1.5, /*factor=*/ 1.5, /*count=*/ 20,
     )
@@ -316,6 +316,33 @@ pub static CORE_MEMPOOL_TXN_CONSENSUS_PULLED_BY_BUCKET: Lazy<HistogramVec> = Laz
     .unwrap()
 });
 
+pub static CORE_MEMPOOL_PARKING_LOT_EVICTED_COUNT: Lazy<Histogram> = Lazy::new(|| {
+    register_histogram!(
+        "aptos_core_mempool_parking_lot_evicted_count",
+        "Number of txns evicted from parking lot",
+        TXN_COUNT_BUCKETS.clone()
+    )
+    .unwrap()
+});
+
+pub static CORE_MEMPOOL_PARKING_LOT_EVICTED_BYTES: Lazy<Histogram> = Lazy::new(|| {
+    register_histogram!(
+        "aptos_core_mempool_parking_lot_evicted_bytes",
+        "Bytes of txns evicted from parking lot",
+        exponential_buckets(/*start=*/ 500.0, /*factor=*/ 1.4, /*count=*/ 32).unwrap()
+    )
+    .unwrap()
+});
+
+pub static CORE_MEMPOOL_PARKING_LOT_EVICTED_LATENCY: Lazy<Histogram> = Lazy::new(|| {
+    register_histogram!(
+        "aptos_core_mempool_parking_lot_evicted_latency",
+        "Latency of evicting for each transaction from parking lot",
+        MEMPOOL_LATENCY_BUCKETS.to_vec()
+    )
+    .unwrap()
+});
+
 /// Counter of pending network events to Mempool
 pub static PENDING_MEMPOOL_NETWORK_EVENTS: Lazy<IntCounterVec> = Lazy::new(|| {
     register_int_counter_vec!(
@@ -333,7 +360,7 @@ static MEMPOOL_SERVICE_TXNS: Lazy<HistogramVec> = Lazy::new(|| {
         "aptos_mempool_service_transactions",
         "Number of transactions handled in one request/response between mempool and consensus/state sync",
         &["type"],
-        TRANSACTION_COUNT_BUCKETS.clone()
+        TXN_COUNT_BUCKETS.clone()
     )
     .unwrap()
 });
diff --git a/mempool/src/shared_mempool/tasks.rs b/mempool/src/shared_mempool/tasks.rs
index befefd98d7d09..d1acbe7ce1a2a 100644
--- a/mempool/src/shared_mempool/tasks.rs
+++ b/mempool/src/shared_mempool/tasks.rs
@@ -375,7 +375,15 @@ fn validate_and_add_transactions<NetworkClient, TransactionValidator>(
         .start_timer();
     let validation_results = transactions
         .par_iter()
-        .map(|t| smp.validator.read().validate_transaction(t.0.clone()))
+        .map(|t| {
+            let result = smp.validator.read().validate_transaction(t.0.clone());
+            // Pre-compute the hash and length if the transaction is valid, before locking mempool
+            if result.is_ok() {
+                t.0.committed_hash();
+                t.0.txn_bytes_len();
+            }
+            result
+        })
         .collect::<Vec<_>>();
     vm_validation_timer.stop_and_record();
     {
diff --git a/mempool/src/tests/common.rs b/mempool/src/tests/common.rs
index 9d94785a920fe..87fa0284a5a68 100644
--- a/mempool/src/tests/common.rs
+++ b/mempool/src/tests/common.rs
@@ -15,7 +15,7 @@ use aptos_types::{
     account_address::AccountAddress,
     chain_id::ChainId,
     mempool_status::MempoolStatusCode,
-    transaction::{RawTransaction, Script, SignedTransaction},
+    transaction::{RawTransaction, Script, SignedTransaction, TransactionArgument},
 };
 use once_cell::sync::Lazy;
 use rand::{rngs::StdRng, SeedableRng};
@@ -45,21 +45,83 @@ static ACCOUNTS: Lazy<Vec<AccountAddress>> = Lazy::new(|| {
     ]
 });
 
+static SMALL_SCRIPT: Lazy<Script> = Lazy::new(|| Script::new(vec![], vec![], vec![]));
+
+static LARGE_SCRIPT: Lazy<Script> = Lazy::new(|| {
+    let mut args = vec![];
+    for _ in 0..200 {
+        args.push(TransactionArgument::Address(AccountAddress::random()));
+    }
+    Script::new(vec![], vec![], args)
+});
+
+static HUGE_SCRIPT: Lazy<Script> = Lazy::new(|| {
+    let mut args = vec![];
+    for _ in 0..200_000 {
+        args.push(TransactionArgument::Address(AccountAddress::random()));
+    }
+    Script::new(vec![], vec![], args)
+});
+
 #[derive(Clone, Serialize, Deserialize)]
 pub struct TestTransaction {
-    pub(crate) address: usize,
+    pub(crate) address: AccountAddress,
     pub(crate) sequence_number: u64,
     pub(crate) gas_price: u64,
     pub(crate) account_seqno: u64,
+    pub(crate) script: Option<Script>,
 }
 
 impl TestTransaction {
-    pub(crate) const fn new(address: usize, sequence_number: u64, gas_price: u64) -> Self {
+    pub(crate) fn new(address: usize, sequence_number: u64, gas_price: u64) -> Self {
+        Self {
+            address: TestTransaction::get_address(address),
+            sequence_number,
+            gas_price,
+            account_seqno: 0,
+            script: None,
+        }
+    }
+
+    pub(crate) fn new_with_large_script(
+        address: usize,
+        sequence_number: u64,
+        gas_price: u64,
+    ) -> Self {
+        Self {
+            address: TestTransaction::get_address(address),
+            sequence_number,
+            gas_price,
+            account_seqno: 0,
+            script: Some(LARGE_SCRIPT.clone()),
+        }
+    }
+
+    pub(crate) fn new_with_huge_script(
+        address: usize,
+        sequence_number: u64,
+        gas_price: u64,
+    ) -> Self {
+        Self {
+            address: TestTransaction::get_address(address),
+            sequence_number,
+            gas_price,
+            account_seqno: 0,
+            script: Some(HUGE_SCRIPT.clone()),
+        }
+    }
+
+    pub(crate) fn new_with_address(
+        address: AccountAddress,
+        sequence_number: u64,
+        gas_price: u64,
+    ) -> Self {
         Self {
             address,
             sequence_number,
             gas_price,
             account_seqno: 0,
+            script: None,
         }
     }
 
@@ -87,9 +149,9 @@ impl TestTransaction {
         exp_timestamp_secs: u64,
     ) -> SignedTransaction {
         let raw_txn = RawTransaction::new_script(
-            TestTransaction::get_address(self.address),
+            self.address,
             self.sequence_number,
-            Script::new(vec![], vec![], vec![]),
+            self.script.clone().unwrap_or(SMALL_SCRIPT.clone()),
             max_gas_amount,
             self.gas_price,
             exp_timestamp_secs,
diff --git a/mempool/src/tests/core_mempool_test.rs b/mempool/src/tests/core_mempool_test.rs
index 86b0c33d46316..9b22e515b1d4e 100644
--- a/mempool/src/tests/core_mempool_test.rs
+++ b/mempool/src/tests/core_mempool_test.rs
@@ -14,7 +14,8 @@ use aptos_config::config::{MempoolConfig, NodeConfig};
 use aptos_consensus_types::common::{TransactionInProgress, TransactionSummary};
 use aptos_crypto::HashValue;
 use aptos_types::{
-    mempool_status::MempoolStatusCode, transaction::SignedTransaction, vm_status::DiscardedVMStatus,
+    account_address::AccountAddress, mempool_status::MempoolStatusCode,
+    transaction::SignedTransaction, vm_status::DiscardedVMStatus,
 };
 use itertools::Itertools;
 use maplit::btreemap;
@@ -756,7 +757,9 @@ fn test_capacity_bytes() {
     let mut txns = vec![];
     let last_txn;
     loop {
-        let txn = new_test_mempool_transaction(address, seq_no);
+        let txn = signed_txn_to_mempool_transaction(
+            TestTransaction::new(address, seq_no, 1).make_signed_transaction(),
+        );
         let txn_bytes = txn.get_estimated_bytes();
 
         if size_bytes <= capacity_bytes {
@@ -811,10 +814,9 @@ fn test_capacity_bytes() {
     }
 }
 
-fn new_test_mempool_transaction(address: usize, sequence_number: u64) -> MempoolTransaction {
-    let signed_txn = TestTransaction::new(address, sequence_number, 1).make_signed_transaction();
+fn signed_txn_to_mempool_transaction(txn: SignedTransaction) -> MempoolTransaction {
     MempoolTransaction::new(
-        signed_txn,
+        txn,
         Duration::from_secs(1),
         1,
         TimelineState::NotReady,
@@ -851,6 +853,98 @@ fn test_parking_lot_eviction() {
     assert!(add_txn(&mut pool, TestTransaction::new(0, 2, 1)).is_err());
 }
 
+#[test]
+fn test_parking_lot_eviction_bytes() {
+    // Get the small transaction size
+    let small_txn_size =
+        signed_txn_to_mempool_transaction(TestTransaction::new(1, 1, 1).make_signed_transaction())
+            .get_estimated_bytes();
+
+    let mut config = NodeConfig::generate_random_config();
+    config.mempool.capacity = 100;
+    // Fit 2 small transactions + one additional transaction (by overflowing the capacity bytes)
+    config.mempool.capacity_bytes = 3 * small_txn_size + 1;
+    let mut pool = CoreMempool::new(&config);
+    // Add 2 small transactions to parking lot
+    for address in 0..2 {
+        add_txn(&mut pool, TestTransaction::new(address, 1, 1)).unwrap();
+    }
+    // Add one large transaction that will top off the capacity bytes
+    add_txn(&mut pool, TestTransaction::new_with_large_script(2, 1, 1)).unwrap();
+    // Mempool is full. Insert a small txn for other account.
+    add_txn(&mut pool, TestTransaction::new(3, 0, 1)).unwrap();
+}
+
+#[test]
+fn test_parking_lot_eviction_benchmark() {
+    // Get the small transaction size
+    let small_txn_size =
+        signed_txn_to_mempool_transaction(TestTransaction::new(1, 1, 1).make_signed_transaction())
+            .get_estimated_bytes();
+    let huge_txn_size = signed_txn_to_mempool_transaction(
+        TestTransaction::new_with_huge_script(1, 1, 1).make_signed_transaction(),
+    )
+    .get_estimated_bytes();
+    let num_small_txns = (huge_txn_size / small_txn_size) * 2;
+
+    let mut config = NodeConfig::generate_random_config();
+    config.mempool.capacity_per_user = 200;
+    config.mempool.capacity = 4_000_000;
+    // ~5 MB
+    config.mempool.capacity_bytes = num_small_txns * small_txn_size + 1;
+    let mut pool = CoreMempool::new(&config);
+
+    // // Add one huge transaction that will evict all transactions from parking lot
+    // let huge_signed_txn = TestTransaction::new_with_huge_script(0, 1, 1).make_signed_transaction();
+    // // Pre-compute these values, as shared mempool would do
+    // huge_signed_txn.committed_hash();
+    // huge_signed_txn.txn_bytes_len();
+    //
+    // let now = Instant::now();
+    // add_signed_txn(&mut pool, huge_signed_txn.clone()).unwrap();
+    // // Flush the huge transaction
+    // add_txn(&mut pool, TestTransaction::new(1, 0, 1)).unwrap();
+
+    let accounts: Vec<_> = (0..num_small_txns)
+        .map(|_| AccountAddress::random())
+        .collect();
+    // Fill up parking lot to capacity
+    for account in accounts {
+        for seq_num in 1..2 {
+            add_txn(
+                &mut pool,
+                TestTransaction::new_with_address(account, seq_num, 1),
+            )
+            .unwrap();
+        }
+    }
+    // Add one huge transaction that will cause mempool to be (beyond) full
+    let huge_signed_txn = TestTransaction::new_with_huge_script(0, 0, 1).make_signed_transaction();
+    add_signed_txn(&mut pool, huge_signed_txn).unwrap();
+    assert_eq!(pool.get_parking_lot_size(), num_small_txns);
+
+    // Add one huge transaction that will evict many transactions from parking lot
+    let huge_signed_txn = TestTransaction::new_with_huge_script(1, 0, 1).make_signed_transaction();
+    // Pre-compute these values, as shared mempool would do
+    huge_signed_txn.committed_hash();
+    huge_signed_txn.txn_bytes_len();
+    let now = Instant::now();
+    add_signed_txn(&mut pool, huge_signed_txn).unwrap();
+    let time_to_evict_ms = now.elapsed().as_millis();
+
+    let has_remainder = huge_txn_size % small_txn_size != 0;
+    let num_expected_evicted = num_small_txns / 2 + has_remainder as usize;
+    assert_eq!(
+        pool.get_parking_lot_size(),
+        num_small_txns - num_expected_evicted
+    );
+    assert!(
+        time_to_evict_ms < 300,
+        "Parking lot eviction of {} should take less than 300 ms on a reasonable machine",
+        num_expected_evicted
+    );
+}
+
 #[test]
 fn test_parking_lot_evict_only_for_ready_txn_insertion() {
     let mut config = NodeConfig::generate_random_config();
diff --git a/mempool/src/tests/integration_tests.rs b/mempool/src/tests/integration_tests.rs
index 21f79a5bdd0b3..8c1661c23a7b7 100644
--- a/mempool/src/tests/integration_tests.rs
+++ b/mempool/src/tests/integration_tests.rs
@@ -19,12 +19,14 @@ use aptos_network::{
     transport::ConnectionMetadata,
     ProtocolId,
 };
+use once_cell::sync::Lazy;
 use std::time::Duration;
 
 const ALL_PROTOCOLS: [ProtocolId; 1] = [ProtocolId::MempoolDirectSend];
-static ALL_TXNS: &[TestTransaction] = &[test_transaction(0), test_transaction(1)];
-static TXN_1: &[TestTransaction] = &[test_transaction(0)];
-static TXN_2: &[TestTransaction] = &[test_transaction(1)];
+static ALL_TXNS: Lazy<Vec<TestTransaction>> =
+    Lazy::new(|| vec![test_transaction(0), test_transaction(1)]);
+static TXN_1: Lazy<Vec<TestTransaction>> = Lazy::new(|| vec![test_transaction(0)]);
+static TXN_2: Lazy<Vec<TestTransaction>> = Lazy::new(|| vec![test_transaction(1)]);
 
 fn inbound_node_combinations() -> [(MempoolNode, (PeerNetworkId, ConnectionMetadata)); 6] {
     [
@@ -94,10 +96,10 @@ async fn single_inbound_node_test() {
         node.receive_message(
             ProtocolId::MempoolDirectSend,
             other_peer_network_id,
-            ALL_TXNS,
+            &ALL_TXNS,
         )
         .await;
-        node.assert_only_txns_in_mempool(ALL_TXNS);
+        node.assert_only_txns_in_mempool(&ALL_TXNS);
     }
 }
 
@@ -106,20 +108,20 @@ async fn single_inbound_node_test() {
 async fn single_outbound_node_test() {
     for (mut node, (other_peer_network_id, other_metadata)) in outbound_node_combinations() {
         // Add transactions
-        node.assert_txns_not_in_mempool(TXN_1);
-        node.add_txns_via_client(TXN_1).await;
+        node.assert_txns_not_in_mempool(&TXN_1);
+        node.add_txns_via_client(&TXN_1).await;
 
         // After we connect, all messages should be received and broadcast upstream
         node.connect_self(other_peer_network_id.network_id(), other_metadata);
-        node.send_broadcast_and_receive_ack(other_peer_network_id, TXN_1)
+        node.send_broadcast_and_receive_ack(other_peer_network_id, &TXN_1)
             .await;
-        node.assert_only_txns_in_mempool(TXN_1);
+        node.assert_only_txns_in_mempool(&TXN_1);
 
         // Adding more txns should also broadcast them upstream
-        node.add_txns_via_client(TXN_2).await;
-        node.send_broadcast_and_receive_ack(other_peer_network_id, TXN_2)
+        node.add_txns_via_client(&TXN_2).await;
+        node.send_broadcast_and_receive_ack(other_peer_network_id, &TXN_2)
             .await;
-        node.assert_only_txns_in_mempool(ALL_TXNS);
+        node.assert_only_txns_in_mempool(&ALL_TXNS);
     }
 }
 
@@ -139,12 +141,12 @@ async fn vfn_middle_man_test() {
     node.connect_self(fn_peer_network_id.network_id(), fn_metadata);
 
     // Incoming transactions should be accepted
-    node.receive_message(ProtocolId::MempoolDirectSend, fn_peer_network_id, ALL_TXNS)
+    node.receive_message(ProtocolId::MempoolDirectSend, fn_peer_network_id, &ALL_TXNS)
         .await;
-    node.assert_only_txns_in_mempool(ALL_TXNS);
+    node.assert_only_txns_in_mempool(&ALL_TXNS);
 
     // And they should be forwarded upstream
-    node.send_broadcast_and_receive_ack(validator_peer_network_id, ALL_TXNS)
+    node.send_broadcast_and_receive_ack(validator_peer_network_id, &ALL_TXNS)
         .await;
 }
 
@@ -156,15 +158,15 @@ async fn test_skip_ack_rebroadcast() {
         validator_mock_connection(ConnectionOrigin::Inbound, &ALL_PROTOCOLS);
 
     // Add transactions
-    node.assert_txns_not_in_mempool(ALL_TXNS);
-    node.add_txns_via_client(ALL_TXNS).await;
+    node.assert_txns_not_in_mempool(&ALL_TXNS);
+    node.add_txns_via_client(&ALL_TXNS).await;
 
     node.connect_self(other_peer_network_id.network_id(), other_metadata.clone());
 
     // After we drop, we should rebroadcast successfully
     node.drop_next_network_msg(other_peer_network_id.network_id())
         .await;
-    node.send_broadcast_and_receive_ack(other_peer_network_id, ALL_TXNS)
+    node.send_broadcast_and_receive_ack(other_peer_network_id, &ALL_TXNS)
         .await;
 }
 
@@ -175,9 +177,9 @@ async fn test_interrupt_in_sync_inbound() {
     for (mut node, (other_peer_network_id, other_metadata)) in inbound_node_combinations() {
         // First txn is received
         node.connect_self(other_peer_network_id.network_id(), other_metadata.clone());
-        node.receive_message(ProtocolId::MempoolDirectSend, other_peer_network_id, TXN_1)
+        node.receive_message(ProtocolId::MempoolDirectSend, other_peer_network_id, &TXN_1)
             .await;
-        node.assert_only_txns_in_mempool(TXN_1);
+        node.assert_only_txns_in_mempool(&TXN_1);
 
         // Drop the connection, and a reconnect should merge txns
         node.disconnect_self(other_peer_network_id.network_id(), other_metadata.clone());
@@ -186,10 +188,10 @@ async fn test_interrupt_in_sync_inbound() {
         node.receive_message(
             ProtocolId::MempoolDirectSend,
             other_peer_network_id,
-            ALL_TXNS,
+            &ALL_TXNS,
         )
         .await;
-        node.assert_only_txns_in_mempool(ALL_TXNS);
+        node.assert_only_txns_in_mempool(&ALL_TXNS);
     }
 }
 
@@ -201,12 +203,12 @@ async fn test_ready_txns() {
         validator_mock_connection(ConnectionOrigin::Inbound, &ALL_PROTOCOLS);
 
     // Add 2nd txn
-    node.assert_txns_not_in_mempool(TXN_2);
-    node.add_txns_via_client(TXN_2).await;
+    node.assert_txns_not_in_mempool(&TXN_2);
+    node.add_txns_via_client(&TXN_2).await;
 
     // No txns should be sent or ready
     node.connect_self(other_peer_network_id.network_id(), other_metadata.clone());
-    node.assert_txns_not_in_mempool(ALL_TXNS);
+    node.assert_txns_not_in_mempool(&ALL_TXNS);
     node.wait_for_no_msg(
         other_peer_network_id.network_id(),
         Duration::from_millis(100),
@@ -214,11 +216,11 @@ async fn test_ready_txns() {
     .await;
 
     // Adding earlier txns should fill in the gaps, and now it should send all
-    node.add_txns_via_client(TXN_1).await;
-    node.assert_txns_in_mempool(ALL_TXNS);
-    node.send_broadcast_and_receive_ack(other_peer_network_id, ALL_TXNS)
+    node.add_txns_via_client(&TXN_1).await;
+    node.assert_txns_in_mempool(&ALL_TXNS);
+    node.send_broadcast_and_receive_ack(other_peer_network_id, &ALL_TXNS)
         .await;
-    node.assert_only_txns_in_mempool(ALL_TXNS);
+    node.assert_only_txns_in_mempool(&ALL_TXNS);
 }
 
 /// Test that in the validator network, messages won't be sent back to the original sender
@@ -230,16 +232,16 @@ async fn test_broadcast_self_txns() {
 
     // Other node sends earlier txn
     node.connect_self(other_peer_network_id.network_id(), other_metadata.clone());
-    node.receive_message(ProtocolId::MempoolDirectSend, other_peer_network_id, TXN_1)
+    node.receive_message(ProtocolId::MempoolDirectSend, other_peer_network_id, &TXN_1)
         .await;
-    node.assert_txns_in_mempool(TXN_1);
+    node.assert_txns_in_mempool(&TXN_1);
 
     // Add txns to current node
-    node.add_txns_via_client(TXN_2).await;
-    node.assert_txns_in_mempool(ALL_TXNS);
+    node.add_txns_via_client(&TXN_2).await;
+    node.assert_txns_in_mempool(&ALL_TXNS);
 
     // Txns should be sent to other node (but not the earlier txn)
-    node.send_broadcast_and_receive_ack(other_peer_network_id, TXN_2)
+    node.send_broadcast_and_receive_ack(other_peer_network_id, &TXN_2)
         .await;
 }
 
@@ -254,12 +256,12 @@ async fn test_update_gas_price() {
         validator_mock_connection(ConnectionOrigin::Outbound, &ALL_PROTOCOLS);
 
     // Get first txn
-    node.add_txns_via_client(TXN_1).await;
-    node.assert_txns_in_mempool(TXN_1);
+    node.add_txns_via_client(&TXN_1).await;
+    node.assert_txns_in_mempool(&TXN_1);
 
     // Send to other node
     node.connect_self(other_peer_network_id.network_id(), other_metadata.clone());
-    node.send_broadcast_and_receive_ack(other_peer_network_id, TXN_1)
+    node.send_broadcast_and_receive_ack(other_peer_network_id, &TXN_1)
         .await;
 
     // Update txn
@@ -279,16 +281,16 @@ async fn test_mempool_full_rebroadcast() {
         validator_mock_connection(ConnectionOrigin::Outbound, &ALL_PROTOCOLS);
 
     // Get first txn
-    node.add_txns_via_client(ALL_TXNS).await;
-    node.assert_txns_in_mempool(ALL_TXNS);
+    node.add_txns_via_client(&ALL_TXNS).await;
+    node.assert_txns_in_mempool(&ALL_TXNS);
 
     // Send to other node (which is full)
     node.connect_self(other_peer_network_id.network_id(), other_metadata.clone());
-    node.send_broadcast_and_receive_retry(other_peer_network_id, ALL_TXNS)
+    node.send_broadcast_and_receive_retry(other_peer_network_id, &ALL_TXNS)
         .await;
 
     // Txn should be sent again later
-    node.send_broadcast_and_receive_ack(other_peer_network_id, ALL_TXNS)
+    node.send_broadcast_and_receive_ack(other_peer_network_id, &ALL_TXNS)
         .await;
 }
 
@@ -300,20 +302,20 @@ async fn test_rebroadcast_retry_is_empty() {
         validator_mock_connection(ConnectionOrigin::Outbound, &ALL_PROTOCOLS);
 
     // Get first txn
-    node.add_txns_via_client(TXN_1).await;
-    node.assert_txns_in_mempool(TXN_1);
+    node.add_txns_via_client(&TXN_1).await;
+    node.assert_txns_in_mempool(&TXN_1);
 
     // Send to other node (which is full)
     node.connect_self(other_peer_network_id.network_id(), other_metadata.clone());
-    node.send_broadcast_and_receive_retry(other_peer_network_id, TXN_1)
+    node.send_broadcast_and_receive_retry(other_peer_network_id, &TXN_1)
         .await;
 
     // Add txn2. In the meantime, txn1 was committed.
-    node.add_txns_via_client(TXN_2).await;
-    node.commit_txns(TXN_1).await;
+    node.add_txns_via_client(&TXN_2).await;
+    node.commit_txns(&TXN_1).await;
 
     // Txn should be sent again later
-    node.send_broadcast_and_receive_ack(other_peer_network_id, TXN_2)
+    node.send_broadcast_and_receive_ack(other_peer_network_id, &TXN_2)
         .await;
 }
 
@@ -343,7 +345,7 @@ async fn fn_to_val_test() {
 
         // NOTE: Always return node at end, or it will be dropped and channels closed
         let pfn_future = async move {
-            pfn.add_txns_via_client(ALL_TXNS).await;
+            pfn.add_txns_via_client(&ALL_TXNS).await;
             pfn.connect(pfn_vfn_network, vfn_metadata);
 
             // Forward to VFN
@@ -370,13 +372,13 @@ async fn fn_to_val_test() {
                 val.send_next_network_msg(vfn_val_network).await;
             }
 
-            val.wait_on_txns_in_mempool(ALL_TXNS).await;
+            val.wait_on_txns_in_mempool(&ALL_TXNS).await;
             val
         };
 
         let (pfn, vfn, val) = futures::future::join3(pfn_future, vfn_future, val_future).await;
-        pfn.assert_only_txns_in_mempool(ALL_TXNS);
-        vfn.assert_only_txns_in_mempool(ALL_TXNS);
-        val.assert_only_txns_in_mempool(ALL_TXNS);
+        pfn.assert_only_txns_in_mempool(&ALL_TXNS);
+        vfn.assert_only_txns_in_mempool(&ALL_TXNS);
+        val.assert_only_txns_in_mempool(&ALL_TXNS);
     }
 }
diff --git a/mempool/src/tests/test_framework.rs b/mempool/src/tests/test_framework.rs
index ddddbb944f21b..30260e6b69557 100644
--- a/mempool/src/tests/test_framework.rs
+++ b/mempool/src/tests/test_framework.rs
@@ -235,12 +235,7 @@ impl MempoolNode {
                 .collect();
             let expected: Vec<_> = txns
                 .iter()
-                .map(|txn| {
-                    (
-                        TestTransaction::get_address(txn.address),
-                        txn.sequence_number,
-                    )
-                })
+                .map(|txn| (txn.address, txn.sequence_number))
                 .collect();
             Err((actual, expected))
         } else {
@@ -392,12 +387,7 @@ impl MempoolNode {
                         .collect();
                     let expected_txns: Vec<_> = expected_txns
                         .iter()
-                        .map(|txn| {
-                            (
-                                TestTransaction::get_address(txn.address),
-                                txn.sequence_number,
-                            )
-                        })
+                        .map(|txn| (txn.address, txn.sequence_number))
                         .collect();
 
                     panic!(
@@ -420,12 +410,7 @@ impl MempoolNode {
                         .collect();
                     let expected_txns: Vec<_> = expected_txns
                         .iter()
-                        .map(|txn| {
-                            (
-                                TestTransaction::get_address(txn.address),
-                                txn.sequence_number,
-                            )
-                        })
+                        .map(|txn| (txn.address, txn.sequence_number))
                         .collect();
 
                     panic!(
@@ -686,8 +671,14 @@ fn mpsc_channel<T>() -> (
 }
 
 /// Creates a single [`TestTransaction`] with the given `seq_num`.
-pub const fn test_transaction(seq_num: u64) -> TestTransaction {
-    TestTransaction::new(1, seq_num, 1)
+pub fn test_transaction(seq_num: u64) -> TestTransaction {
+    TestTransaction {
+        address: TestTransaction::get_address(1),
+        sequence_number: seq_num,
+        gas_price: 1,
+        account_seqno: 0,
+        script: None,
+    }
 }
 
 /// Tells us if a [`SignedTransaction`] block contains only the [`TestTransaction`]s
@@ -723,7 +714,7 @@ pub fn block_contains_any_transaction(
 fn block_contains_transaction(block: &[SignedTransaction], txn: &TestTransaction) -> bool {
     block.iter().any(|signed_txn| {
         signed_txn.sequence_number() == txn.sequence_number
-            && signed_txn.sender() == TestTransaction::get_address(txn.address)
+            && signed_txn.sender() == txn.address
             && signed_txn.gas_unit_price() == txn.gas_price
     })
 }