apollo_gateway: move the state reader factory into the stateful validator factory (#9087)

Author: ArniStarkware

crates/apollo_gateway/src/gateway.rs

@@ -60,7 +60,6 @@ pub struct Gateway {
     pub config: Arc<GatewayConfig>,
     pub stateless_tx_validator: Arc<dyn StatelessTransactionValidatorTrait>,
     pub stateful_tx_validator_factory: Arc<dyn StatefulTransactionValidatorFactoryTrait>,
-    pub state_reader_factory: Arc<dyn StateReaderFactory>,
     pub mempool_client: SharedMempoolClient,
     pub transaction_converter: Arc<dyn TransactionConverterTrait>,
 }
@@ -79,11 +78,11 @@ impl Gateway {
             stateful_tx_validator_factory: Arc::new(StatefulTransactionValidatorFactory {
                 config: config.stateful_tx_validator_config.clone(),
                 chain_info: config.chain_info.clone(),
+                state_reader_factory,
                 contract_class_manager: ContractClassManager::start(
                     config.contract_class_manager_config.clone(),
                 ),
             }),
-            state_reader_factory,
             mempool_client,
             transaction_converter,
         }
@@ -153,7 +152,7 @@ impl Gateway {
 
         let mut stateful_transaction_validator = self
             .stateful_tx_validator_factory
-            .instantiate_validator(self.state_reader_factory.clone())
+            .instantiate_validator()
             .await
             .inspect_err(|e| metric_counters.record_add_tx_failure(e))?;
 

crates/apollo_gateway/src/gateway_test.rs

@@ -93,7 +93,6 @@ use crate::metrics::{
     LABEL_NAME_SOURCE,
     LABEL_NAME_TX_TYPE,
 };
-use crate::state_reader::MockStateReaderFactory;
 use crate::state_reader_test_utils::{local_test_state_reader_factory, TestStateReaderFactory};
 use crate::stateful_transaction_validator::{
     MockStatefulTransactionValidatorFactoryTrait,
@@ -561,15 +560,14 @@ async fn add_tx_returns_error_when_extract_state_nonce_and_run_validations_fails
 
     mock_stateful_transaction_validator_factory
         .expect_instantiate_validator()
-        .return_once(|_| Ok(Box::new(mock_stateful_transaction_validator)));
+        .return_once(|| Ok(Box::new(mock_stateful_transaction_validator)));
 
     let tx_args = invoke_args();
     setup_transaction_converter_mock(&mut mock_dependencies.mock_transaction_converter, &tx_args);
     let gateway = Gateway {
         config: Arc::new(mock_dependencies.config),
         stateless_tx_validator: Arc::new(mock_dependencies.mock_stateless_transaction_validator),
         stateful_tx_validator_factory: Arc::new(mock_stateful_transaction_validator_factory),
-        state_reader_factory: Arc::new(MockStateReaderFactory::new()),
         mempool_client: Arc::new(mock_dependencies.mock_mempool_client),
         transaction_converter: Arc::new(mock_dependencies.mock_transaction_converter),
     };
@@ -614,15 +612,14 @@ async fn add_tx_returns_error_when_instantiating_validator_fails(
     };
     mock_stateful_transaction_validator_factory
         .expect_instantiate_validator()
-        .return_once(|_| Err(expected_error));
+        .return_once(|| Err(expected_error));
 
     let tx_args = invoke_args();
     setup_transaction_converter_mock(&mut mock_dependencies.mock_transaction_converter, &tx_args);
     let gateway = Gateway {
         config: Arc::new(mock_dependencies.config),
         stateless_tx_validator: Arc::new(mock_dependencies.mock_stateless_transaction_validator),
         stateful_tx_validator_factory: Arc::new(mock_stateful_transaction_validator_factory),
-        state_reader_factory: Arc::new(MockStateReaderFactory::new()),
         mempool_client: Arc::new(mock_dependencies.mock_mempool_client),
         transaction_converter: Arc::new(mock_dependencies.mock_transaction_converter),
     };

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -49,25 +49,25 @@ type BlockifierStatefulValidator = StatefulValidator<
 pub trait StatefulTransactionValidatorFactoryTrait: Send + Sync {
     async fn instantiate_validator(
         &self,
-        state_reader_factory: Arc<dyn StateReaderFactory>,
     ) -> StatefulTransactionValidatorResult<Box<dyn StatefulTransactionValidatorTrait>>;
 }
+
 pub struct StatefulTransactionValidatorFactory {
     pub config: StatefulTransactionValidatorConfig,
     pub chain_info: ChainInfo,
+    pub state_reader_factory: Arc<dyn StateReaderFactory>,
     pub contract_class_manager: ContractClassManager,
 }
 
 #[async_trait]
 impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFactory {
-    // TODO(Ayelet): Move state_reader_factory and chain_info to the struct.
     async fn instantiate_validator(
         &self,
-        state_reader_factory: Arc<dyn StateReaderFactory>,
     ) -> StatefulTransactionValidatorResult<Box<dyn StatefulTransactionValidatorTrait>> {
         // TODO(yael 6/5/2024): consider storing the block_info as part of the
         // StatefulTransactionValidator and update it only once a new block is created.
-        let (blockifier_state_reader, gateway_fixed_block_state_reader) = state_reader_factory
+        let (blockifier_state_reader, gateway_fixed_block_state_reader) = self
+            .state_reader_factory
             .get_blockifier_state_reader_and_gateway_fixed_block_from_latest_block()
             .await
             .map_err(|err| GatewaySpecError::UnexpectedError {

crates/apollo_gateway/src/stateful_transaction_validator_test.rs

@@ -129,17 +129,17 @@ async fn test_run_pre_validation_checks(
 #[rstest]
 #[tokio::test]
 async fn test_instantiate_validator() {
+    let state_reader_factory =
+        local_test_state_reader_factory(CairoVersion::Cairo1(RunnableCairo1::Casm), false);
+
     let stateful_validator_factory = StatefulTransactionValidatorFactory {
         config: StatefulTransactionValidatorConfig::default(),
         chain_info: ChainInfo::create_for_testing(),
+        state_reader_factory: Arc::new(state_reader_factory),
         contract_class_manager: ContractClassManager::start(ContractClassManagerConfig::default()),
     };
-    let state_reader_factory =
-        local_test_state_reader_factory(CairoVersion::Cairo1(RunnableCairo1::Casm), false);
-
-    let validator =
-        stateful_validator_factory.instantiate_validator(Arc::new(state_reader_factory)).await;
 
+    let validator = stateful_validator_factory.instantiate_validator().await;
     assert!(validator.is_ok());
 }