feat(txpool): add support for mock testing (#55)

* chore: some cleanup

* refactor(txpool): simplify layers and add docs

* refactor: more cleanup

* refactor: cleanup and simplifications

* feat(txpool): mock test support

* feat(txpool): more mock testing

* chore: rustfmt

* set basefee correctly

Cargo.lock

@@ -1936,6 +1936,7 @@ dependencies = [
  "linked-hash-map",
  "parking_lot",
  "paste",
+ "rand",
  "reth-primitives",
  "serde",
  "thiserror",

crates/transaction-pool/Cargo.toml

@@ -30,3 +30,4 @@ bitflags = "1.3"
 
 [dev-dependencies]
 paste = "1.0"
+rand = "0.8"

crates/transaction-pool/src/pool/parked.rs

@@ -51,6 +51,16 @@ impl<T: ParkedOrd> ParkedPool<T> {
         self.submission_id = self.submission_id.wrapping_add(1);
         id
     }
+
+    /// Number of transactions in the entire pool
+    pub(crate) fn len(&self) -> usize {
+        self.by_id.len()
+    }
+
+    /// Whether the pool is empty
+    pub(crate) fn is_empty(&self) -> bool {
+        self.by_id.is_empty()
+    }
 }
 
 impl<T: ParkedOrd> Default for ParkedPool<T> {

crates/transaction-pool/src/pool/pending.rs

@@ -142,6 +142,16 @@ impl<T: TransactionOrdering> PendingPool<T> {
     pub(crate) fn get(&self, id: &TransactionId) -> Option<Arc<PendingTransaction<T>>> {
         self.by_id.get(id).cloned()
     }
+
+    /// Number of transactions in the entire pool
+    pub(crate) fn len(&self) -> usize {
+        self.by_id.len()
+    }
+
+    /// Whether the pool is empty
+    pub(crate) fn is_empty(&self) -> bool {
+        self.by_id.is_empty()
+    }
 }
 
 /// A transaction that is ready to be included in a block.

crates/transaction-pool/src/pool/txpool.rs

@@ -80,6 +80,8 @@ pub struct TxPool<T: TransactionOrdering> {
     all_transactions: AllTransactions<T::Transaction>,
 }
 
+// === impl TxPool ===
+
 impl<T: TransactionOrdering> TxPool<T> {
     /// Create a new graph pool instance.
     pub fn new(ordering: Arc<T>) -> Self {
@@ -277,6 +279,27 @@ impl<T: TransactionOrdering> TxPool<T> {
     }
 }
 
+// Additional test impls
+#[cfg(test)]
+#[allow(missing_docs)]
+impl<T: TransactionOrdering> TxPool<T> {
+    pub(crate) fn all(&self) -> &AllTransactions<T::Transaction> {
+        &self.all_transactions
+    }
+
+    pub(crate) fn pending(&self) -> &PendingPool<T> {
+        &self.pending_pool
+    }
+
+    pub(crate) fn base_fee(&self) -> &ParkedPool<BasefeeOrd<T::Transaction>> {
+        &self.basefee_pool
+    }
+
+    pub(crate) fn queued(&self) -> &ParkedPool<QueuedOrd<T::Transaction>> {
+        &self.queued_pool
+    }
+}
+
 /// Container for _all_ transaction in the pool.
 ///
 /// This is the sole entrypoint that's guarding all sub-pools, all sub-pool actions are always

crates/transaction-pool/src/test_util/mock.rs

@@ -6,8 +6,12 @@ use crate::{
     PoolTransaction, TransactionOrdering, ValidPoolTransaction,
 };
 use paste::paste;
+use rand::{
+    distributions::{Uniform, WeightedIndex},
+    prelude::Distribution,
+};
 use reth_primitives::{Address, TxHash, H256, U256};
-use std::{sync::Arc, time::Instant};
+use std::{ops::Range, sync::Arc, time::Instant};
 
 pub type MockTxPool = TxPool<MockOrdering>;
 
@@ -256,6 +260,14 @@ impl MockTransaction {
         let gas = self.get_gas_limit() + 1;
         next.with_gas_limit(gas)
     }
+
+    pub fn is_legacy(&self) -> bool {
+        matches!(self, MockTransaction::Legacy { .. })
+    }
+
+    pub fn is_eip1559(&self) -> bool {
+        matches!(self, MockTransaction::Eip1559 { .. })
+    }
 }
 
 impl PoolTransaction for MockTransaction {
@@ -364,6 +376,39 @@ impl TransactionOrdering for MockOrdering {
     }
 }
 
+/// A configured distribution that can generate transactions
+pub struct MockTransactionDistribution {
+    /// legacy to EIP-1559 ration
+    legacy_ratio: WeightedIndex<u32>,
+    /// generates the gas limit
+    gas_limit_range: Uniform<u64>,
+}
+
+impl MockTransactionDistribution {
+    /// Creates a new generator distribution.
+    ///
+    /// Expects legacy tx in full pct: `30u32` is `30%`.
+    pub fn new(legacy_pct: u32, gas_limit_range: Range<u64>) -> Self {
+        assert!(legacy_pct <= 100, "expect pct");
+
+        let eip_1559 = 100 - legacy_pct;
+        Self {
+            legacy_ratio: WeightedIndex::new([eip_1559, legacy_pct]).unwrap(),
+            gas_limit_range: gas_limit_range.into(),
+        }
+    }
+
+    /// Generates a new transaction
+    pub fn tx(&self, nonce: u64, rng: &mut impl rand::Rng) -> MockTransaction {
+        let tx = if self.legacy_ratio.sample(rng) == 0 {
+            MockTransaction::eip1559()
+        } else {
+            MockTransaction::legacy()
+        };
+        tx.with_nonce(nonce).with_gas_limit(self.gas_limit_range.sample(rng))
+    }
+}
+
 #[test]
 fn test_mock_priority() {
     let o = MockOrdering;

crates/transaction-pool/src/test_util/mod.rs

@@ -2,5 +2,6 @@
 #![allow(missing_docs, unused)]
 
 mod mock;
+mod pool;
 
 pub use mock::*;

crates/transaction-pool/src/test_util/pool.rs

@@ -0,0 +1,230 @@
+//! Test helpers for mocking an entire pool.
+
+use crate::{
+    error::PoolResult,
+    pool::{txpool::TxPool, AddedTransaction},
+    test_util::{
+        MockOrdering, MockTransaction, MockTransactionDistribution, MockTransactionFactory,
+    },
+    TransactionOrdering,
+};
+use rand::Rng;
+use reth_primitives::{Address, U256};
+use serde::{Deserialize, Serialize};
+use std::{
+    collections::HashMap,
+    ops::{Deref, DerefMut},
+    sync::Arc,
+};
+
+/// A wrapped `TxPool` with additional helpers for testing
+pub struct MockPool<T: TransactionOrdering = MockOrdering> {
+    // The wrapped pool.
+    pool: TxPool<T>,
+}
+
+impl MockPool {
+    /// The total size of all subpools
+    fn total_subpool_size(&self) -> usize {
+        self.pool.pending().len() + self.pool.base_fee().len() + self.pool.queued().len()
+    }
+
+    /// Checks that all pool invariants hold.
+    fn enforce_invariants(&self) {
+        assert_eq!(
+            self.pool.len(),
+            self.total_subpool_size(),
+            "Tx in AllTransactions and sum(subpools) must match"
+        );
+    }
+}
+
+impl Default for MockPool {
+    fn default() -> Self {
+        Self { pool: TxPool::new(Arc::new(MockOrdering::default())) }
+    }
+}
+
+impl<T: TransactionOrdering> Deref for MockPool<T> {
+    type Target = TxPool<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.pool
+    }
+}
+
+impl<T: TransactionOrdering> DerefMut for MockPool<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.pool
+    }
+}
+
+/// Simulates transaction execution.
+pub struct MockTransactionSimulator<R: Rng> {
+    /// The pending base fee
+    base_fee: U256,
+    /// Generator for transactions
+    tx_generator: MockTransactionDistribution,
+    /// represents the on chain balance of a sender.
+    balances: HashMap<Address, U256>,
+    /// represents the on chain nonce of a sender.
+    nonces: HashMap<Address, u64>,
+    /// A set of addresses to as senders.
+    senders: Vec<Address>,
+    /// What scenarios to execute.
+    scenarios: Vec<ScenarioType>,
+    /// All previous scenarios executed by a sender.
+    executed: HashMap<Address, ExecutedScenarios>,
+    /// "Validates" generated transactions.
+    validator: MockTransactionFactory,
+    /// The rng instance used to select senders and scenarios.
+    rng: R,
+}
+
+impl<R: Rng> MockTransactionSimulator<R> {
+    /// Returns a new mock instance
+    pub fn new(mut rng: R, config: MockSimulatorConfig) -> Self {
+        let senders = config.addresses(&mut rng);
+        let nonces = senders.iter().copied().map(|a| (a, 0)).collect();
+        let balances = senders.iter().copied().map(|a| (a, config.balance)).collect();
+        Self {
+            base_fee: config.base_fee,
+            balances,
+            nonces,
+            senders,
+            scenarios: config.scenarios,
+            tx_generator: config.tx_generator,
+            executed: Default::default(),
+            validator: Default::default(),
+            rng,
+        }
+    }
+
+    /// Returns a random address from the senders set
+    fn rng_address(&mut self) -> Address {
+        let idx = self.rng.gen_range(0..self.senders.len());
+        self.senders[idx]
+    }
+
+    /// Returns a random scenario from the scenario set
+    fn rng_scenario(&mut self) -> ScenarioType {
+        let idx = self.rng.gen_range(0..self.scenarios.len());
+        self.scenarios[idx].clone()
+    }
+
+    /// Executes the next scenario and applies it to the pool
+    pub fn next(&mut self, pool: &mut MockPool) {
+        let sender = self.rng_address();
+        let scenario = self.rng_scenario();
+        let on_chain_nonce = self.nonces[&sender];
+        let on_chain_balance = self.balances[&sender];
+
+        match scenario {
+            ScenarioType::OnchainNonce => {
+                let tx = self
+                    .tx_generator
+                    .tx(on_chain_nonce, &mut self.rng)
+                    .with_gas_price(self.base_fee);
+                let valid_tx = self.validator.validated(tx);
+
+                let res = pool.add_transaction(valid_tx, on_chain_balance, on_chain_nonce).unwrap();
+
+                // TODO(mattsse): need a way expect based on the current state of the pool and tx
+                // settings
+
+                match res {
+                    AddedTransaction::Pending(_) => {}
+                    AddedTransaction::Parked { .. } => {
+                        panic!("expected pending")
+                    }
+                }
+
+                // TODO(mattsse): check subpools
+            }
+            ScenarioType::HigherNonce { .. } => {
+                unimplemented!()
+            }
+        }
+
+        // make sure everything is set
+        pool.enforce_invariants()
+    }
+}
+
+/// How to configure a new mock transaction stream
+pub struct MockSimulatorConfig {
+    /// How many senders to generate.
+    pub num_senders: usize,
+    // TODO(mattsse): add a way to generate different balances
+    pub balance: U256,
+    /// Scenarios to test
+    pub scenarios: Vec<ScenarioType>,
+    /// The start base fee
+    pub base_fee: U256,
+    /// generator for transactions
+    pub tx_generator: MockTransactionDistribution,
+}
+
+impl MockSimulatorConfig {
+    /// Generates a set of random addresses
+    pub fn addresses(&self, rng: &mut impl rand::Rng) -> Vec<Address> {
+        std::iter::repeat_with(|| Address::random_using(rng)).take(self.num_senders).collect()
+    }
+}
+
+/// Represents
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum ScenarioType {
+    OnchainNonce,
+    HigherNonce { skip: u64 },
+}
+
+/// The actual scenario, ready to be executed
+///
+/// A scenario produces one or more transactions and expects a certain Outcome.
+///
+/// An executed scenario can affect previous executed transactions
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum Scenario {
+    /// Send a tx with the same nonce as on chain.
+    OnchainNonce { nonce: u64 },
+    /// Send a tx with a higher nonce that what the sender has on chain
+    HigherNonce { onchain: u64, nonce: u64 },
+    Multi {
+        // Execute multiple test scenarios
+        scenario: Vec<Scenario>,
+    },
+}
+
+/// Represents an executed scenario
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ExecutedScenario {
+    /// balance at the time of execution
+    balance: U256,
+    /// nonce at the time of execution
+    nonce: u64,
+    /// The executed scenario
+    scenario: Scenario,
+}
+
+/// All executed scenarios by a sender
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ExecutedScenarios {
+    sender: Address,
+    scenarios: Vec<ExecutedScenario>,
+}
+
+#[test]
+fn test_on_chain_nonce_scenario() {
+    let config = MockSimulatorConfig {
+        num_senders: 10,
+        balance: 200_000u64.into(),
+        scenarios: vec![ScenarioType::OnchainNonce],
+        base_fee: 10u64.into(),
+        tx_generator: MockTransactionDistribution::new(30, 10..100),
+    };
+    let mut simulator = MockTransactionSimulator::new(rand::thread_rng(), config);
+    let mut pool = MockPool::default();
+
+    simulator.next(&mut pool);
+}