build(blockifier_reexecution): add stability for rpc test (#1712)

Cargo.toml

@@ -188,6 +188,7 @@ rand_distr = "0.4.3"
 regex = "1.10.4"
 replace_with = "0.1.7"
 reqwest = "0.11"
+retry = "2.0.0"
 rstest = "0.17.0"
 rustc-hex = "2.1.0"
 schemars = "0.8.12"

crates/blockifier_reexecution/Cargo.toml

@@ -17,6 +17,7 @@ flate2.workspace = true
 indexmap = { workspace = true, features = ["serde"] }
 papyrus_execution.workspace = true
 pretty_assertions.workspace = true
+retry.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 starknet-core.workspace = true

crates/blockifier_reexecution/src/state_reader/test_state_reader.rs

@@ -26,6 +26,7 @@ use starknet_gateway::rpc_objects::{BlockHeader, GetBlockWithTxHashesParams, Res
 use starknet_gateway::rpc_state_reader::RpcStateReader;
 use starknet_types_core::felt::Felt;
 
+use crate::retry_request;
 use crate::state_reader::compile::{legacy_to_contract_class_v0, sierra_to_contact_class_v1};
 use crate::state_reader::errors::ReexecutionError;
 use crate::state_reader::reexecution_state_reader::ReexecutionStateReader;
@@ -41,6 +42,11 @@ use crate::state_reader::utils::{
     ReexecutionStateMaps,
 };
 
+pub const DEFAULT_RETRY_COUNT: usize = 3;
+pub const DEFAULT_RETRY_WAIT_TIME: u64 = 1000;
+pub const DEFAULT_EXPECTED_ERROR_STRING: &str = "Connection error";
+pub const DEFAULT_RETRY_FAILURE_MESSAGE: &str = "Failed to connect to the RPC node.";
+
 pub type ReexecutionResult<T> = Result<T, ReexecutionError>;
 
 pub type StarknetContractClassMapping = HashMap<ClassHash, StarknetContractClass>;
@@ -103,27 +109,60 @@ impl From<SerializableOfflineReexecutionData> for OfflineReexecutionData {
     }
 }
 
+pub struct RetryConfig {
+    pub(crate) n_retries: usize,
+    pub(crate) retry_interval_milliseconds: u64,
+    pub(crate) expected_error_string: &'static str,
+    pub(crate) retry_failure_message: &'static str,
+}
+
+impl Default for RetryConfig {
+    fn default() -> Self {
+        Self {
+            n_retries: DEFAULT_RETRY_COUNT,
+            retry_interval_milliseconds: DEFAULT_RETRY_WAIT_TIME,
+            expected_error_string: DEFAULT_EXPECTED_ERROR_STRING,
+            retry_failure_message: DEFAULT_RETRY_FAILURE_MESSAGE,
+        }
+    }
+}
+
 pub struct TestStateReader {
-    rpc_state_reader: RpcStateReader,
+    pub(crate) rpc_state_reader: RpcStateReader,
+    pub(crate) retry_config: RetryConfig,
     #[allow(dead_code)]
     contract_class_mapping_dumper: Arc<Mutex<Option<StarknetContractClassMapping>>>,
 }
 
+impl Default for TestStateReader {
+    fn default() -> Self {
+        Self {
+            rpc_state_reader: RpcStateReader::from_latest(&get_rpc_state_reader_config()),
+            retry_config: RetryConfig::default(),
+            contract_class_mapping_dumper: Arc::new(Mutex::new(None)),
+        }
+    }
+}
+
 impl StateReader for TestStateReader {
     fn get_nonce_at(&self, contract_address: ContractAddress) -> StateResult<Nonce> {
-        self.rpc_state_reader.get_nonce_at(contract_address)
+        retry_request!(self.retry_config, || self.rpc_state_reader.get_nonce_at(contract_address))
     }
 
     fn get_storage_at(
         &self,
         contract_address: ContractAddress,
         key: StorageKey,
     ) -> StateResult<Felt> {
-        self.rpc_state_reader.get_storage_at(contract_address, key)
+        retry_request!(self.retry_config, || self
+            .rpc_state_reader
+            .get_storage_at(contract_address, key))
     }
 
     fn get_class_hash_at(&self, contract_address: ContractAddress) -> StateResult<ClassHash> {
-        self.rpc_state_reader.get_class_hash_at(contract_address)
+        retry_request!(self.retry_config, || self
+            .rpc_state_reader
+            .get_class_hash_at(contract_address))
     }
 
     /// Returns the contract class of the given class hash.
@@ -132,7 +171,10 @@ impl StateReader for TestStateReader {
         &self,
         class_hash: ClassHash,
     ) -> StateResult<RunnableContractClass> {
-        match self.get_contract_class(&class_hash)? {
+        let contract_class =
+            retry_request!(self.retry_config, || self.get_contract_class(&class_hash))?;
+
+        match contract_class {
             StarknetContractClass::Sierra(sierra) => {
                 Ok(sierra_to_contact_class_v1(sierra).unwrap().try_into().unwrap())
             }
@@ -156,6 +198,7 @@ impl TestStateReader {
         Self {
             rpc_state_reader: RpcStateReader::from_number(config, block_number),
             contract_class_mapping_dumper,
+            retry_config: RetryConfig::default(),
         }
     }
 

crates/blockifier_reexecution/src/state_reader/utils.rs

@@ -101,3 +101,30 @@ impl TryFrom<ReexecutionStateMaps> for StateMaps {
         })
     }
 }
+
+#[macro_export]
+macro_rules! retry_request {
+    ($retry_config:expr, $closure:expr) => {{
+        retry::retry(
+            retry::delay::Fixed::from_millis($retry_config.retry_interval_milliseconds)
+                .take($retry_config.n_retries),
+            || {
+                match $closure() {
+                    Ok(value) => retry::OperationResult::Ok(value),
+                    // If the error contains the expected_error_string , we want to retry.
+                    Err(e) if e.to_string().contains($retry_config.expected_error_string) => {
+                        retry::OperationResult::Retry(e)
+                    }
+                    // For all other errors, do not retry and return immediately.
+                    Err(e) => retry::OperationResult::Err(e),
+                }
+            },
+        )
+        .map_err(|e| {
+            if e.error.to_string().contains($retry_config.expected_error_string) {
+                panic!("{}: {:?}", $retry_config.retry_failure_message, e.error);
+            }
+            e.error
+        })
+    }};
+}