Merge pull request #408 from propeller-heads/cluster-test/dc/ENG-5103-test-stale-components

feat(cluster-test): Test stale components

Author: dianacarvalho1

src/evm/protocol/uniswap_v4/state.rs

@@ -1,4 +1,4 @@
-use std::{any::Any, collections::HashMap};
+use std::{any::Any, collections::HashMap, fmt};
 
 use alloy::primitives::{Address, Sign, I256, U256};
 use num_bigint::BigUint;
@@ -40,7 +40,7 @@ use crate::evm::protocol::{
     vm::constants::EXTERNAL_ACCOUNT,
 };
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct UniswapV4State {
     liquidity: u128,
     sqrt_price: U256,
@@ -51,6 +51,19 @@ pub struct UniswapV4State {
     pub hook: Option<Box<dyn HookHandler>>,
 }
 
+impl fmt::Debug for UniswapV4State {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("UniswapV4State")
+            .field("liquidity", &self.liquidity)
+            .field("sqrt_price", &self.sqrt_price)
+            .field("fees", &self.fees)
+            .field("tick", &self.tick)
+            .field("ticks", &self.ticks)
+            .field("tick_spacing", &self.tick_spacing)
+            .finish_non_exhaustive()
+    }
+}
+
 impl PartialEq for UniswapV4State {
     fn eq(&self, other: &Self) -> bool {
         match (&self.hook, &other.hook) {

src/evm/protocol/vm/state.rs

@@ -2,7 +2,7 @@
 use std::{
     any::Any,
     collections::{HashMap, HashSet},
-    fmt::Debug,
+    fmt::{self, Debug},
     str::FromStr,
 };
 
@@ -34,7 +34,7 @@ use crate::evm::{
     },
 };
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct EVMPoolState<D: EngineDatabaseInterface + Clone + Debug>
 where
     <D as DatabaseRef>::Error: Debug,
@@ -70,6 +70,23 @@ where
     adapter_contract: TychoSimulationContract<D>,
 }
 
+impl<D> Debug for EVMPoolState<D>
+where
+    D: EngineDatabaseInterface + Clone + Debug,
+    <D as DatabaseRef>::Error: Debug,
+    <D as EngineDatabaseInterface>::Error: Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("EVMPoolState")
+            .field("id", &self.id)
+            .field("tokens", &self.tokens)
+            .field("balances", &self.balances)
+            .field("involved_contracts", &self.involved_contracts)
+            .field("contract_balances", &self.contract_balances)
+            .finish_non_exhaustive()
+    }
+}
+
 impl<D> EVMPoolState<D>
 where
     D: EngineDatabaseInterface + Clone + Debug + 'static,

src/rfq/protocols/bebop/state.rs

@@ -1,4 +1,4 @@
-use std::{any::Any, collections::HashMap};
+use std::{any::Any, collections::HashMap, fmt};
 
 use async_trait::async_trait;
 use num_bigint::BigUint;
@@ -19,14 +19,23 @@ use crate::rfq::{
     protocols::bebop::{client::BebopClient, models::BebopPriceData},
 };
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct BebopState {
     pub base_token: Token,
     pub quote_token: Token,
     pub price_data: BebopPriceData,
     pub client: BebopClient,
 }
 
+impl fmt::Debug for BebopState {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("BebopState")
+            .field("base_token", &self.base_token)
+            .field("quote_token", &self.quote_token)
+            .finish_non_exhaustive()
+    }
+}
+
 impl BebopState {
     pub fn new(
         base_token: Token,

src/rfq/protocols/hashflow/state.rs

@@ -1,4 +1,4 @@
-use std::{any::Any, collections::HashMap};
+use std::{any::Any, collections::HashMap, fmt};
 
 use async_trait::async_trait;
 use num_bigint::BigUint;
@@ -19,16 +19,25 @@ use crate::rfq::{
     protocols::hashflow::{client::HashflowClient, models::HashflowMarketMakerLevels},
 };
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct HashflowState {
     pub base_token: Token,
     pub quote_token: Token,
     pub levels: HashflowMarketMakerLevels,
-    #[allow(dead_code)]
     pub market_maker: String,
     pub client: HashflowClient,
 }
 
+impl fmt::Debug for HashflowState {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("HashflowState")
+            .field("base_token", &self.base_token)
+            .field("quote_token", &self.quote_token)
+            .field("market_maker", &self.market_maker)
+            .finish_non_exhaustive()
+    }
+}
+
 impl HashflowState {
     pub fn new(
         base_token: Token,

tycho-integration-test/src/main.rs

@@ -3,9 +3,8 @@ mod metrics;
 mod stream_processor;
 
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     fmt::Debug,
-    num::NonZeroUsize,
     sync::{Arc, RwLock},
     time::Duration,
 };
@@ -21,10 +20,10 @@ use alloy_chains::NamedChain;
 use clap::Parser;
 use dotenv::dotenv;
 use itertools::Itertools;
-use lru::LruCache;
 use miette::{miette, IntoDiagnostic, NarratableReportHandler, WrapErr};
 use num_bigint::BigUint;
 use num_traits::{ToPrimitive, Zero};
+use rand::prelude::IndexedRandom;
 use tokio::sync::Semaphore;
 use tracing::{error, info, warn};
 use tracing_subscriber::EnvFilter;
@@ -98,10 +97,14 @@ struct Cli {
     #[arg(long, default_value_t = 5, value_parser = clap::value_parser!(u8).range(1..))]
     parallel_simulations: u8,
 
-    /// Maximum number of simulations to run per protocol update
+    /// Maximum number of simulations (of updated states) to run per update
     #[arg(long, default_value_t = 10, value_parser = clap::value_parser!(u16).range(1..))]
     max_simulations: u16,
 
+    /// Maximum number of simulations (of stale states) to run per update per protocol
+    #[arg(long, default_value_t = 10, value_parser = clap::value_parser!(u16).range(1..))]
+    max_simulations_stale: u16,
+
     /// The RFQ stream will skip messages for this duration (in seconds) after processing a message
     #[arg(long, default_value_t = 600)]
     skip_messages_duration: u64,
@@ -120,6 +123,13 @@ impl Debug for Cli {
     }
 }
 
+#[derive(Default)]
+struct TychoState {
+    states: HashMap<String, Box<dyn ProtocolSim>>,
+    components: HashMap<String, ProtocolComponent>,
+    component_ids_by_protocol: HashMap<String, HashSet<String>>,
+}
+
 #[tokio::main]
 async fn main() -> miette::Result<()> {
     miette::set_hook(Box::new(|_| Box::new(NarratableReportHandler::new())))?;
@@ -192,13 +202,7 @@ async fn run(cli: Cli) -> miette::Result<()> {
         }
     }
 
-    // Assuming a ProtocolComponent instance can be around 1KB (2 tokens, 2 contract_ids, 6 static
-    // attributes) 250,000 entries would use 250MB of memory.
-    // In a 25min test, the cache increased at a rate of ~2 items/minute, or ~3k items/day, so it
-    // would take ~80 days to get full and start dropping the least used items.
-    let protocol_pairs = Arc::new(RwLock::new(LruCache::new(
-        NonZeroUsize::new(250_000).ok_or_else(|| miette!("Invalid NonZeroUsize"))?,
-    )));
+    let tycho_state = Arc::new(RwLock::new(TychoState::default()));
 
     // Process streams updates
     info!("Waiting for first protocol update...");
@@ -214,15 +218,15 @@ async fn run(cli: Cli) -> miette::Result<()> {
 
         let cli = cli.clone();
         let rpc_tools = rpc_tools.clone();
-        let protocol_pairs = protocol_pairs.clone();
+        let tycho_state = tycho_state.clone();
         let permit = semaphore
             .clone()
             .acquire_owned()
             .await
             .into_diagnostic()
             .wrap_err("Failed to acquire permit")?;
         tokio::spawn(async move {
-            if let Err(e) = process_update(cli, chain, rpc_tools, protocol_pairs, &update).await {
+            if let Err(e) = process_update(cli, chain, rpc_tools, tycho_state, &update).await {
                 warn!("{}", format_error_chain(&e));
             }
             drop(permit);
@@ -236,7 +240,7 @@ async fn process_update(
     cli: Arc<Cli>,
     chain: Chain,
     rpc_tools: RPCTools,
-    protocol_pairs: Arc<RwLock<LruCache<String, ProtocolComponent>>>,
+    tycho_state: Arc<RwLock<TychoState>>,
     update: &StreamUpdate,
 ) -> miette::Result<()> {
     info!(
@@ -264,21 +268,24 @@ async fn process_update(
 
     if let UpdateType::Protocol = update.update_type {
         {
-            let mut pairs = protocol_pairs
+            let mut current_state = tycho_state
                 .write()
-                .map_err(|e| miette!("Failed to acquire write lock on protocol pairs: {e}"))?;
-            let prev_size = pairs.len();
+                .map_err(|e| miette!("Failed to acquire write lock on Tycho state: {e}"))?;
             for (id, comp) in update.update.new_pairs.iter() {
-                pairs.put(id.clone(), comp.clone());
-            }
-            let new_size = pairs.len();
-            let cap = pairs.cap().get();
-            if new_size != prev_size {
-                info!(size=%new_size, capacity=%cap, "Protocol components cache updated");
+                current_state
+                    .components
+                    .insert(id.clone(), comp.clone());
+                current_state
+                    .component_ids_by_protocol
+                    .entry(comp.protocol_system.clone())
+                    .or_insert_with(HashSet::new)
+                    .insert(id.clone());
             }
-            if new_size == cap {
-                warn!(size=%new_size, capacity=%cap, "Protocol components cache reached capacity, \
-                least recently used items will be evicted on new insertions");
+            for (id, state) in update.update.states.iter() {
+                // this overwrites existing entries
+                current_state
+                    .states
+                    .insert(id.clone(), state.clone());
             }
         }
         // Record block processing latency
@@ -312,7 +319,7 @@ async fn process_update(
         metrics::record_protocol_sync_state(protocol, sync_state);
     }
 
-    // Process states in parallel
+    // Process updated states in parallel
     let semaphore = Arc::new(Semaphore::new(cli.parallel_simulations as usize));
     let mut tasks = Vec::new();
 
@@ -324,10 +331,11 @@ async fn process_update(
     {
         let component = match update.update_type {
             UpdateType::Protocol => {
-                let mut pairs = protocol_pairs
-                    .write()
-                    .map_err(|e| miette!("Failed to acquire read lock on protocol pairs: {e}"))?;
-                match pairs.get(id) {
+                let states = &tycho_state
+                    .read()
+                    .map_err(|e| miette!("Failed to acquire read lock on Tycho state: {e}"))?
+                    .components;
+                match states.get(id) {
                     Some(comp) => comp.clone(),
                     None => {
                         warn!(id=%id, "Component not found in cached protocol pairs. Potential causes: \
@@ -366,6 +374,76 @@ async fn process_update(
         tasks.push(task);
     }
 
+    // Select states that were not updated in this block to test simulation and execution
+    let selected_ids = {
+        let current_state = tycho_state
+            .read()
+            .map_err(|e| miette!("Failed to acquire write lock on Tycho state: {e}"))?;
+
+        let mut all_selected_ids = Vec::new();
+        for component_ids in current_state
+            .component_ids_by_protocol
+            .values()
+        {
+            // Filter out IDs that are in the current update
+            let available_ids: Vec<_> = component_ids
+                .iter()
+                .filter(|id| !update.update.states.keys().contains(id))
+                .cloned()
+                .collect();
+
+            let protocol_selected_ids: Vec<_> = available_ids
+                .choose_multiple(
+                    &mut rand::rng(),
+                    (cli.max_simulations_stale as usize).min(available_ids.len()),
+                )
+                .cloned()
+                .collect();
+
+            all_selected_ids.extend(protocol_selected_ids);
+        }
+        all_selected_ids
+    };
+
+    for id in selected_ids {
+        let (component, state) = {
+            let current_state = tycho_state
+                .read()
+                .map_err(|e| miette!("Failed to acquire read lock on Tycho state: {e}"))?;
+
+            match (current_state.components.get(&id), current_state.states.get(&id)) {
+                (Some(comp), Some(state)) => (comp.clone(), state.clone()),
+                (None, _) => {
+                    error!(id=%id, "Component not found in saved protocol components.");
+                    continue;
+                }
+                (_, None) => {
+                    error!(id=%id, "State not found in saved protocol states");
+                    continue;
+                }
+            }
+        };
+
+        let block = block.clone();
+        let state_id = id.clone();
+        let state = state.clone_box();
+        let permit = semaphore
+            .clone()
+            .acquire_owned()
+            .await
+            .into_diagnostic()
+            .wrap_err("Failed to acquire permit")?;
+
+        let task = tokio::spawn(async move {
+            let simulation_id = generate_simulation_id(&component.protocol_system, &state_id);
+            let result =
+                process_state(&simulation_id, chain, component, &block, state_id, state).await;
+            drop(permit);
+            result
+        });
+        tasks.push(task);
+    }
+
     let mut block_execution_info = HashMap::new();
 
     for task in tasks {
@@ -504,6 +582,7 @@ async fn process_state(
                     token_in = %token_in.address,
                     token_out = %token_out.address,
                     error = %format_error_chain(&e),
+                    state = ?state,
                     "Get limits operation failed: {}", format_error_chain(&e)
                 );
                 metrics::record_get_limits_failure(&component.protocol_system);
@@ -548,6 +627,7 @@ async fn process_state(
                     token_out = %token_out.address,
                     amount_in = %amount_in,
                     error = %format_error_chain(&e),
+                    state = ?state,
                     "Get amount out operation failed: {}", format_error_chain(&e)
                 );
                 metrics::record_get_amount_out_failure(&component.protocol_system);