apollo_gateway: split instantiate_validator

Author: itamar-starkware

crates/apollo_gateway/src/gateway.rs

@@ -257,7 +257,8 @@ impl ProcessTxBlockingTask {
 
     fn process_tx(self) -> GatewayResult<Nonce> {
         let mut stateful_transaction_validator = self.runtime.block_on(
-            self.stateful_tx_validator_factory.instantiate_validator(self.state_reader_factory),
+            self.stateful_tx_validator_factory
+                .instantiate_validator(self.state_reader_factory.as_ref()),
         )?;
 
         let nonce = stateful_transaction_validator.extract_state_nonce_and_run_validations(

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -1,5 +1,3 @@
-use std::sync::Arc;
-
 use apollo_gateway_config::config::StatefulTransactionValidatorConfig;
 use apollo_gateway_types::deprecated_gateway_error::{
     KnownStarknetErrorCode,
@@ -54,7 +52,17 @@ type BlockifierStatefulValidator = StatefulValidator<
 pub trait StatefulTransactionValidatorFactoryTrait: Send + Sync {
     async fn instantiate_validator(
         &self,
-        state_reader_factory: Arc<dyn StateReaderFactory>,
+        state_reader_factory: &dyn StateReaderFactory,
+    ) -> StatefulTransactionValidatorResult<Box<dyn StatefulTransactionValidatorTrait>>;
+
+    async fn get_state_reader_for_validation(
+        &self,
+        state_reader_factory: &dyn StateReaderFactory,
+    ) -> StatefulTransactionValidatorResult<Box<dyn GatewayStateReaderWithCompiledClasses>>;
+
+    async fn create_validator_from_state_reader(
+        &self,
+        state_reader: Box<dyn GatewayStateReaderWithCompiledClasses>,
     ) -> StatefulTransactionValidatorResult<Box<dyn StatefulTransactionValidatorTrait>>;
 }
 pub struct StatefulTransactionValidatorFactory {
@@ -68,11 +76,18 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
     // TODO(Ayelet): Move state_reader_factory and chain_info to the struct.
     async fn instantiate_validator(
         &self,
-        state_reader_factory: Arc<dyn StateReaderFactory>,
+        state_reader_factory: &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 state_reader = state_reader_factory
+        // Compose the split API to preserve original behavior.
+        let state_reader = self.get_state_reader_for_validation(state_reader_factory).await?;
+        self.create_validator_from_state_reader(state_reader).await
+    }
+
+    async fn get_state_reader_for_validation(
+        &self,
+        state_reader_factory: &dyn StateReaderFactory,
+    ) -> StatefulTransactionValidatorResult<Box<dyn GatewayStateReaderWithCompiledClasses>> {
+        state_reader_factory
             .get_state_reader_from_latest_block()
             .await
             .map_err(|err| GatewaySpecError::UnexpectedError {
@@ -83,8 +98,16 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
                     "Failed to get state reader from latest block",
                     e,
                 )
-            })?;
-        let latest_block_info = get_latest_block_info(&state_reader).await?;
+            })
+    }
+
+    async fn create_validator_from_state_reader(
+        &self,
+        state_reader: Box<dyn GatewayStateReaderWithCompiledClasses>,
+    ) -> 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 latest_block_info = get_latest_block_info(state_reader.as_ref()).await?;
 
         let state_reader_and_contract_manager = StateReaderAndContractManager {
             state_reader,

crates/apollo_gateway/src/stateful_transaction_validator_test.rs

@@ -165,8 +165,7 @@ async fn test_instantiate_validator() {
     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(&state_reader_factory).await;
 
     assert!(validator.is_ok());
 }