Test the simulator

Author: niran

crates/simulator/src/core.rs

@@ -1,13 +1,20 @@
 use crate::engine::SimulationEngine;
 use crate::publisher::SimulationPublisher;
 use crate::types::SimulationRequest;
+use async_trait::async_trait;
 use eyre::Result;
-use reth_provider::StateProviderFactory;
 use tracing::{error, info};
 
-/// Core bundle simulator that provides shared simulation logic
-/// Used by both mempool event simulators and ExEx event simulators
-pub struct BundleSimulator<E, P>
+/// Clean trait for bundle simulation without exposing Reth's complex types
+#[async_trait]
+pub trait BundleSimulator: Send + Sync {
+    /// Simulate a bundle execution
+    async fn simulate(&self, request: &SimulationRequest) -> Result<()>;
+}
+
+/// Production bundle simulator for Reth
+/// This is the Reth-specific implementation
+pub struct RethBundleSimulator<E, P>
 where
     E: SimulationEngine,
     P: SimulationPublisher,
@@ -16,34 +23,28 @@ where
     publisher: P,
 }
 
-impl<E, P> BundleSimulator<E, P>
+impl<E, P> RethBundleSimulator<E, P>
 where
     E: SimulationEngine,
     P: SimulationPublisher,
 {
     pub fn new(engine: E, publisher: P) -> Self {
-        Self { engine, publisher }
+        Self { 
+            engine, 
+            publisher,
+        }
     }
+}
 
-    /// Process a simulation request by creating state provider from factory
-    /// Convenience method that handles state provider creation
-    pub async fn simulate<F>(
-        &self,
-        request: &SimulationRequest,
-        state_provider_factory: &F,
-    ) -> Result<()>
-    where
-        F: StateProviderFactory,
-    {
-        // Get state provider for the block
-        // FIXME: We probably want to get the state provider once per block rather than once per
-        // bundle for each block.
-        let state_provider = state_provider_factory
-            .state_by_block_hash(request.block_hash)
-            .map_err(|e| eyre::eyre!("Failed to get state provider: {}", e))?;
-
-        // Run the simulation
-        match self.engine.simulate_bundle(request, &state_provider).await {
+#[async_trait]
+impl<E, P> BundleSimulator for RethBundleSimulator<E, P>
+where
+    E: SimulationEngine + 'static,
+    P: SimulationPublisher + 'static,
+{
+    async fn simulate(&self, request: &SimulationRequest) -> Result<()> {
+        // Run the simulation - engine will get its own state provider
+        match self.engine.simulate_bundle(request).await {
             Ok(result) => {
                 info!(
                     bundle_id = %request.bundle_id,

crates/simulator/src/engine.rs

@@ -101,13 +101,10 @@ where
 #[async_trait]
 pub trait SimulationEngine: Send + Sync {
     /// Simulate a bundle execution
-    async fn simulate_bundle<S>(
+    async fn simulate_bundle(
         &self,
         request: &SimulationRequest,
-        state_provider: &S,
-    ) -> Result<SimulationResult>
-    where
-        S: StateProvider + Send + Sync;
+    ) -> Result<SimulationResult>;
 }
 
 #[derive(Clone)]
@@ -137,14 +134,10 @@ where
     Node: FullNodeComponents,
     <Node as FullNodeComponents>::Evm: ConfigureEvm<NextBlockEnvCtx = OpNextBlockEnvAttributes>,
 {
-    async fn simulate_bundle<S>(
+    async fn simulate_bundle(
         &self,
         request: &SimulationRequest,
-        _state_provider: &S,
-    ) -> Result<SimulationResult>
-    where
-        S: StateProvider + Send + Sync,
-    {
+    ) -> Result<SimulationResult> {
         let start_time = Instant::now();
         let simulation_id = Uuid::new_v4();
 

crates/simulator/src/lib.rs

@@ -15,7 +15,7 @@ use std::sync::Arc;
 use tracing::{error, info};
 
 pub use config::SimulatorNodeConfig;
-pub use core::BundleSimulator;
+pub use core::{BundleSimulator, RethBundleSimulator};
 pub use engine::{RethSimulationEngine, SimulationEngine};
 pub use listeners::{ExExEventListener, MempoolEventListener, MempoolListenerConfig};
 pub use publisher::{SimulationPublisher, TipsSimulationPublisher};
@@ -24,15 +24,14 @@ pub use worker_pool::SimulationWorkerPool;
 
 // Type aliases for concrete implementations
 pub type TipsBundleSimulator<Node> =
-    BundleSimulator<RethSimulationEngine<Node>, TipsSimulationPublisher>;
+    RethBundleSimulator<RethSimulationEngine<Node>, TipsSimulationPublisher>;
 pub type TipsExExEventListener<Node> = ExExEventListener<
     Node,
-    RethSimulationEngine<Node>,
-    TipsSimulationPublisher,
+    TipsBundleSimulator<Node>,
     tips_datastore::PostgresDatastore,
 >;
 pub type TipsMempoolEventListener<Node> =
-    MempoolEventListener<Node, RethSimulationEngine<Node>, TipsSimulationPublisher>;
+    MempoolEventListener<Node, TipsBundleSimulator<Node>>;
 
 // Initialization functions
 
@@ -43,7 +42,7 @@ where
     <Node as FullNodeComponents>::Evm: ConfigureEvm<NextBlockEnvCtx = OpNextBlockEnvAttributes>,
 {
     datastore: Arc<tips_datastore::PostgresDatastore>,
-    simulator: BundleSimulator<RethSimulationEngine<Node>, TipsSimulationPublisher>,
+    simulator: RethBundleSimulator<RethSimulationEngine<Node>, TipsSimulationPublisher>,
 }
 
 /// Initialize common listener components (database, publisher, engine, core simulator)
@@ -81,7 +80,7 @@ where
     let engine = RethSimulationEngine::new(Arc::clone(&provider), evm_config);
     info!("Simulation engine initialized");
 
-    let simulator = BundleSimulator::new(engine, publisher);
+    let simulator = RethBundleSimulator::new(engine, publisher);
     info!("Core bundle simulator initialized");
 
     Ok(CommonListenerComponents {
@@ -119,7 +118,6 @@ where
 
     let worker_pool = SimulationWorkerPool::new(
         Arc::new(common_components.simulator),
-        Arc::clone(&provider),
         config.max_concurrent_simulations,
     );
 
@@ -160,7 +158,6 @@ where
 
     let worker_pool = SimulationWorkerPool::new(
         Arc::new(common_components.simulator),
-        Arc::clone(&provider),
         max_concurrent_simulations,
     );
 
@@ -183,9 +180,7 @@ where
     Node: FullNodeComponents,
     <Node as FullNodeComponents>::Evm: ConfigureEvm<NextBlockEnvCtx = OpNextBlockEnvAttributes>,
 {
-    worker_pool: Arc<
-        SimulationWorkerPool<RethSimulationEngine<Node>, TipsSimulationPublisher, Node::Provider>,
-    >,
+    worker_pool: Arc<SimulationWorkerPool<TipsBundleSimulator<Node>>>,
     exex_listener: TipsExExEventListener<Node>,
     mempool_listener: TipsMempoolEventListener<Node>,
 }
@@ -222,7 +217,6 @@ where
 
         let shared_worker_pool = SimulationWorkerPool::new(
             Arc::new(common_components.simulator),
-            Arc::clone(&provider),
             max_concurrent_simulations,
         );
 

crates/simulator/src/listeners/exex.rs

@@ -1,5 +1,4 @@
-use crate::engine::SimulationEngine;
-use crate::publisher::SimulationPublisher;
+use crate::core::BundleSimulator;
 use crate::types::SimulationRequest;
 use crate::worker_pool::{SimulationTask, SimulationWorkerPool};
 
@@ -62,33 +61,31 @@ where
 
 /// ExEx event listener that processes chain events and queues bundle simulations
 /// Processes chain events (commits, reorgs, reverts) and queues simulation tasks
-pub struct ExExEventListener<Node, E, P, D>
+pub struct ExExEventListener<Node, B, D>
 where
     Node: FullNodeComponents,
-    E: SimulationEngine + Clone + 'static,
-    P: SimulationPublisher + Clone + 'static,
+    B: BundleSimulator + 'static,
     D: tips_datastore::BundleDatastore,
 {
     /// The execution extension context
     ctx: ExExContext<Node>,
     /// Datastore for fetching bundles from mempool
     datastore: Arc<D>,
     /// Shared simulation worker pool
-    worker_pool: Arc<SimulationWorkerPool<E, P, Node::Provider>>,
+    worker_pool: Arc<SimulationWorkerPool<B>>,
 }
 
-impl<Node, E, P, D> ExExEventListener<Node, E, P, D>
+impl<Node, B, D> ExExEventListener<Node, B, D>
 where
     Node: FullNodeComponents,
-    E: SimulationEngine + Clone + 'static,
-    P: SimulationPublisher + Clone + 'static,
+    B: BundleSimulator + 'static,
     D: tips_datastore::BundleDatastore + 'static,
 {
     /// Create a new ExEx event listener
     pub fn new(
         ctx: ExExContext<Node>,
         datastore: Arc<D>,
-        worker_pool: Arc<SimulationWorkerPool<E, P, Node::Provider>>,
+        worker_pool: Arc<SimulationWorkerPool<B>>,
     ) -> Self {
         Self {
             ctx,
@@ -190,12 +187,12 @@ where
     }
 
     /// Process a single block for potential bundle simulations
-    async fn process_block<B>(
+    async fn process_block<Block>(
         &mut self,
-        block: (&B256, &reth_primitives::RecoveredBlock<B>),
+        block: (&B256, &reth_primitives::RecoveredBlock<Block>),
     ) -> Result<()>
     where
-        B: reth_node_api::Block,
+        Block: reth_node_api::Block,
     {
         let (block_hash, sealed_block) = block;
         let block_number = sealed_block.number();

crates/simulator/src/listeners/mempool.rs

@@ -1,5 +1,4 @@
-use crate::engine::SimulationEngine;
-use crate::publisher::SimulationPublisher;
+use crate::core::BundleSimulator;
 use crate::types::SimulationRequest;
 use crate::worker_pool::{SimulationTask, SimulationWorkerPool};
 use alloy_primitives::B256;
@@ -31,11 +30,10 @@ pub struct MempoolListenerConfig {
 }
 
 /// Mempool event listener that processes events and queues simulations
-pub struct MempoolEventListener<Node, E, P>
+pub struct MempoolEventListener<Node, B>
 where
     Node: FullNodeComponents,
-    E: SimulationEngine,
-    P: SimulationPublisher,
+    B: BundleSimulator + 'static,
 {
     /// State provider factory for getting current block info
     provider: Arc<Node::Provider>,
@@ -44,20 +42,19 @@ where
     /// Kafka topic name
     topic: String,
     /// Shared simulation worker pool
-    worker_pool: Arc<SimulationWorkerPool<E, P, Node::Provider>>,
+    worker_pool: Arc<SimulationWorkerPool<B>>,
 }
 
-impl<Node, E, P> MempoolEventListener<Node, E, P>
+impl<Node, B> MempoolEventListener<Node, B>
 where
     Node: FullNodeComponents,
-    E: SimulationEngine + Clone + 'static,
-    P: SimulationPublisher + Clone + 'static,
+    B: BundleSimulator + 'static,
 {
     /// Create a new mempool event listener
     pub fn new(
         provider: Arc<Node::Provider>,
         config: MempoolListenerConfig,
-        worker_pool: Arc<SimulationWorkerPool<E, P, Node::Provider>>,
+        worker_pool: Arc<SimulationWorkerPool<B>>,
     ) -> Result<Self> {
         let consumer: StreamConsumer = ClientConfig::new()
             .set("group.id", &config.kafka_group_id)
@@ -86,8 +83,7 @@ where
     /// Run the mempool event listener
     pub async fn run(self) -> Result<()>
     where
-        E: 'static,
-        P: 'static,
+        B: 'static,
     {
         info!(
             topic = %self.topic,

crates/simulator/src/worker_pool.rs

@@ -1,8 +1,5 @@
 use crate::core::BundleSimulator;
-use crate::engine::SimulationEngine;
-use crate::publisher::SimulationPublisher;
 use crate::types::SimulationRequest;
-use reth_provider::StateProviderFactory;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use tokio::sync::mpsc;
@@ -15,16 +12,12 @@ pub struct SimulationTask {
 }
 
 /// Generic simulation worker pool that can be shared across different simulators
-pub struct SimulationWorkerPool<E, P, S>
+pub struct SimulationWorkerPool<B>
 where
-    E: SimulationEngine,
-    P: SimulationPublisher,
-    S: StateProviderFactory,
+    B: BundleSimulator,
 {
     /// Core bundle simulator
-    simulator: Arc<BundleSimulator<E, P>>,
-    /// State provider factory
-    state_provider_factory: Arc<S>,
+    simulator: Arc<B>,
     /// Channel for sending simulation requests to workers
     simulation_tx: mpsc::Sender<SimulationTask>,
     /// Channel for receiving simulation requests in workers
@@ -37,23 +30,19 @@ where
     max_concurrent: usize,
 }
 
-impl<E, P, S> SimulationWorkerPool<E, P, S>
+impl<B> SimulationWorkerPool<B>
 where
-    E: SimulationEngine + Clone + 'static,
-    P: SimulationPublisher + Clone + 'static,
-    S: reth_provider::StateProviderFactory + Send + Sync + 'static,
+    B: BundleSimulator + 'static,
 {
     /// Create a new simulation worker pool
     pub fn new(
-        simulator: Arc<BundleSimulator<E, P>>,
-        state_provider_factory: Arc<S>,
+        simulator: Arc<B>,
         max_concurrent_simulations: usize,
     ) -> Arc<Self> {
         let (simulation_tx, simulation_rx) = mpsc::channel(1000);
 
         Arc::new(Self {
             simulator,
-            state_provider_factory,
             simulation_tx,
             simulation_rx: Arc::new(tokio::sync::Mutex::new(simulation_rx)),
             latest_block: AtomicU64::new(0),
@@ -145,7 +134,7 @@ where
             // Execute the simulation
             match pool
                 .simulator
-                .simulate(&task.request, pool.state_provider_factory.as_ref())
+                .simulate(&task.request)
                 .await
             {
                 Ok(_) => {