Skip to content

Commit

Permalink
feat: add sequential_transactions_by_sender method for `MockTransac…
Browse files Browse the repository at this point in the history
…tionSet` (paradigmxyz#5741)
  • Loading branch information
tcoratger authored Dec 15, 2023
1 parent 686129c commit ed7e6af
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 47 deletions.
85 changes: 41 additions & 44 deletions crates/transaction-pool/src/pool/pending.rs
Original file line number Diff line number Diff line change
Expand Up @@ -564,15 +564,13 @@ impl<T: TransactionOrdering> Ord for PendingTransaction<T> {

#[cfg(test)]
mod tests {
use std::collections::HashSet;

use reth_primitives::address;

use super::*;
use crate::{
test_utils::{MockOrdering, MockTransaction, MockTransactionFactory, MockTransactionSet},
PoolTransaction,
};
use reth_primitives::{address, TxType};
use std::collections::HashSet;

#[test]
fn test_enforce_basefee() {
Expand Down Expand Up @@ -709,86 +707,85 @@ mod tests {

#[test]
fn truncate_by_sender() {
// this test ensures that we evict from the pending pool by sender
// This test ensures that transactions are removed from the pending pool by sender.
let mut f = MockTransactionFactory::default();
let mut pool = PendingPool::new(MockOrdering::default());

// Addresses for simulated senders A, B, C, and D.
let a = address!("000000000000000000000000000000000000000a");
let b = address!("000000000000000000000000000000000000000b");
let c = address!("000000000000000000000000000000000000000c");
let d = address!("000000000000000000000000000000000000000d");

// TODO: make creating these mock tx chains easier
// create a chain of transactions by sender A, B, C
let a1 = MockTransaction::eip1559().with_sender(a);
let a2 = a1.next();
let a3 = a2.next();
let a4 = a3.next();

let b1 = MockTransaction::eip1559().with_sender(b);
let b2 = b1.next();
let b3 = b2.next();

// C has the same number of txs as B
let c1 = MockTransaction::eip1559().with_sender(c);
let c2 = c1.next();
let c3 = c2.next();

let d1 = MockTransaction::eip1559().with_sender(d);
// Create transaction chains for senders A, B, C, and D.
let a_txs = MockTransactionSet::sequential_transactions_by_sender(a, 4, TxType::EIP1559);
let b_txs = MockTransactionSet::sequential_transactions_by_sender(b, 3, TxType::EIP1559);
let c_txs = MockTransactionSet::sequential_transactions_by_sender(c, 3, TxType::EIP1559);
let d_txs = MockTransactionSet::sequential_transactions_by_sender(d, 1, TxType::EIP1559);

// Set up expected pending transactions.
let expected_pending = vec![
a_txs.transactions[0].clone(),
b_txs.transactions[0].clone(),
c_txs.transactions[0].clone(),
a_txs.transactions[1].clone(),
]
.into_iter()
.map(|tx| (tx.sender(), tx.nonce()))
.collect::<HashSet<_>>();

// just construct a list of all txs to add
let expected_pending = vec![a1.clone(), b1.clone(), c1.clone(), a2.clone()]
.into_iter()
.map(|tx| (tx.sender(), tx.nonce()))
.collect::<HashSet<_>>();
// Set up expected removed transactions.
let expected_removed = vec![
d1.clone(),
c3.clone(),
b3.clone(),
a4.clone(),
c2.clone(),
b2.clone(),
a3.clone(),
d_txs.transactions[0].clone(),
c_txs.transactions[2].clone(),
b_txs.transactions[2].clone(),
a_txs.transactions[3].clone(),
c_txs.transactions[1].clone(),
b_txs.transactions[1].clone(),
a_txs.transactions[2].clone(),
]
.into_iter()
.map(|tx| (tx.sender(), tx.nonce()))
.collect::<HashSet<_>>();
let all_txs = vec![a1, a2, a3, a4.clone(), b1, b2, b3, c1, c2, c3, d1];

// add all the transactions to the pool
// Consolidate all transactions into a single vector.
let all_txs =
[a_txs.into_vec(), b_txs.into_vec(), c_txs.into_vec(), d_txs.into_vec()].concat();

// Add all the transactions to the pool.
for tx in all_txs {
pool.add_transaction(f.validated_arc(tx), 0);
}

// sanity check, make sure everything checks out
// Sanity check, ensuring everything is consistent.
pool.assert_invariants();

// let's set the max total txs to 4, since we remove txs for each sender first, we remove
// in this order:
// Define the maximum total transactions to be 4, removing transactions for each sender.
// Expected order of removal:
// * d1, c3, b3, a4
// * c2, b2, a3
//
// and we are left with:
// Remaining transactions:
// * a1, a2
// * b1
// * c1
let pool_limit = SubPoolLimit { max_txs: 4, max_size: usize::MAX };

// truncate the pool
// Truncate the pool based on the defined limit.
let removed = pool.truncate_pool(pool_limit);
pool.assert_invariants();
assert_eq!(removed.len(), expected_removed.len());

// get the inner txs from the removed txs
// Get the set of removed transactions and compare with the expected set.
let removed =
removed.into_iter().map(|tx| (tx.sender(), tx.nonce())).collect::<HashSet<_>>();
assert_eq!(removed, expected_removed);

// get the pending pool
// Retrieve the current pending transactions after truncation.
let pending = pool.all().collect::<Vec<_>>();
assert_eq!(pending.len(), expected_pending.len());

// get the inner txs from the pending txs
// Get the set of pending transactions and compare with the expected set.
let pending =
pending.into_iter().map(|tx| (tx.sender(), tx.nonce())).collect::<HashSet<_>>();
assert_eq!(pending, expected_pending);
Expand Down
22 changes: 19 additions & 3 deletions crates/transaction-pool/src/test_utils/mock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1145,8 +1145,12 @@ impl MockTransactionSet {
Self { transactions }
}

/// Create a list of dependent transactions with a common sender. The transactions start at the
/// given nonce, and the sender is incremented by the given tx_count.
/// Creates a series of dependent transactions for a given sender and nonce.
///
/// This method generates a sequence of transactions starting from the provided nonce
/// for the given sender.
///
/// The number of transactions created is determined by `tx_count`.
pub fn dependent(sender: Address, from_nonce: u64, tx_count: usize, tx_type: TxType) -> Self {
let mut txs = Vec::with_capacity(tx_count);
let mut curr_tx = MockTransaction::new_from_type(tx_type).with_nonce(from_nonce);
Expand All @@ -1156,7 +1160,19 @@ impl MockTransactionSet {
txs.push(curr_tx.clone());
}

MockTransactionSet::new(txs)
Self::new(txs)
}

/// Creates a chain of transactions for a given sender with a specified count.
///
/// This method generates a sequence of transactions starting from the specified sender
/// and creates a chain of transactions based on the `tx_count`.
pub fn sequential_transactions_by_sender(
sender: Address,
tx_count: usize,
tx_type: TxType,
) -> Self {
Self::dependent(sender, 0, tx_count, tx_type)
}

/// Add transactions to the [MockTransactionSet]
Expand Down

0 comments on commit ed7e6af

Please sign in to comment.