apollo_gateway: move instantiate_validator to be async

Author: itamar-starkware

crates/apollo_gateway/src/gateway.rs

@@ -256,9 +256,9 @@ impl ProcessTxBlockingTask {
     }
 
     fn process_tx(self) -> GatewayResult<Nonce> {
-        let mut stateful_transaction_validator = self
-            .stateful_tx_validator_factory
-            .instantiate_validator(self.state_reader_factory.as_ref(), self.runtime.clone())?;
+        let mut stateful_transaction_validator = self.runtime.block_on(
+            self.stateful_tx_validator_factory.instantiate_validator(self.state_reader_factory),
+        )?;
 
         let nonce = stateful_transaction_validator.extract_state_nonce_and_run_validations(
             &self.executable_tx,

crates/apollo_gateway/src/gateway_test.rs

@@ -98,6 +98,7 @@ use crate::state_reader_test_utils::{local_test_state_reader_factory, TestStateR
 use crate::stateful_transaction_validator::{
     MockStatefulTransactionValidatorFactoryTrait,
     MockStatefulTransactionValidatorTrait,
+    StatefulTransactionValidatorTrait,
 };
 use crate::stateless_transaction_validator::MockStatelessTransactionValidatorTrait;
 
@@ -563,9 +564,13 @@ async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_f
         .expect_extract_state_nonce_and_run_validations()
         .return_once(|_, _, _| Err(expected_error));
 
-    mock_stateful_transaction_validator_factory
-        .expect_instantiate_validator()
-        .return_once(|_, _| Ok(Box::new(mock_stateful_transaction_validator)));
+    mock_stateful_transaction_validator_factory.expect_instantiate_validator().return_once(|_| {
+        Box::pin(async {
+            Ok::<Box<dyn StatefulTransactionValidatorTrait>, _>(Box::new(
+                mock_stateful_transaction_validator,
+            ))
+        })
+    });
 
     let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory);
 
@@ -608,7 +613,7 @@ async fn process_tx_returns_error_when_instantiating_validator_fails(
     };
     mock_stateful_transaction_validator_factory
         .expect_instantiate_validator()
-        .return_once(|_, _| Err(expected_error));
+        .return_once(|_| Box::pin(async { Err::<_, _>(expected_error) }));
 
     let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory);
 

crates/apollo_gateway/src/rpc_state_reader.rs

@@ -112,8 +112,8 @@ impl MempoolStateReader for RpcStateReader {
         .await
         .map_err(|e| StateError::StateReadError(format!("JoinError: {e}")))??;
 
-        let block_header: BlockHeader =
-            serde_json::from_value(get_block_with_tx_hashes_result).map_err(serde_err_to_state_err)?;
+        let block_header: BlockHeader = serde_json::from_value(get_block_with_tx_hashes_result)
+            .map_err(serde_err_to_state_err)?;
         let block_info = block_header.try_into()?;
         Ok(block_info)
     }

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use apollo_gateway_config::config::StatefulTransactionValidatorConfig;
 use apollo_gateway_types::deprecated_gateway_error::{
     KnownStarknetErrorCode,
@@ -7,6 +9,7 @@ use apollo_gateway_types::deprecated_gateway_error::{
 use apollo_gateway_types::errors::GatewaySpecError;
 use apollo_mempool_types::communication::SharedMempoolClient;
 use apollo_proc_macros::sequencer_latency_histogram;
+use async_trait::async_trait;
 use blockifier::blockifier::stateful_validator::{
     StatefulValidator,
     StatefulValidatorTrait as BlockifierStatefulValidatorTrait,
@@ -46,12 +49,12 @@ type BlockifierStatefulValidator = StatefulValidator<
     StateReaderAndContractManager<Box<dyn GatewayStateReaderWithCompiledClasses>>,
 >;
 
+#[async_trait]
 #[cfg_attr(test, mockall::automock)]
 pub trait StatefulTransactionValidatorFactoryTrait: Send + Sync {
-    fn instantiate_validator(
+    async fn instantiate_validator(
         &self,
-        state_reader_factory: &dyn StateReaderFactory,
-        runtime: tokio::runtime::Handle,
+        state_reader_factory: Arc<dyn StateReaderFactory>,
     ) -> StatefulTransactionValidatorResult<Box<dyn StatefulTransactionValidatorTrait>>;
 }
 pub struct StatefulTransactionValidatorFactory {
@@ -60,17 +63,18 @@ pub struct StatefulTransactionValidatorFactory {
     pub contract_class_manager: ContractClassManager,
 }
 
+#[async_trait]
 impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFactory {
     // TODO(Ayelet): Move state_reader_factory and chain_info to the struct.
-    fn instantiate_validator(
+    async fn instantiate_validator(
         &self,
-        state_reader_factory: &dyn StateReaderFactory,
-        runtime: tokio::runtime::Handle,
+        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 state_reader = runtime
-            .block_on(state_reader_factory.get_state_reader_from_latest_block())
+        let state_reader = state_reader_factory
+            .get_state_reader_from_latest_block()
+            .await
             .map_err(|err| GatewaySpecError::UnexpectedError {
                 data: format!("Internal server error: {err}"),
             })
@@ -80,7 +84,7 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
                     e,
                 )
             })?;
-        let latest_block_info = runtime.block_on(get_latest_block_info(&state_reader))?;
+        let latest_block_info = get_latest_block_info(&state_reader).await?;
 
         let state_reader_and_contract_manager = StateReaderAndContractManager {
             state_reader,

crates/apollo_gateway/src/stateful_transaction_validator_test.rs

@@ -40,7 +40,6 @@ use starknet_api::transaction::fields::{
 };
 use starknet_api::{declare_tx_args, deploy_account_tx_args, invoke_tx_args, nonce};
 
-use crate::state_reader::{MockStateReaderFactory, StateReaderFactory};
 use crate::state_reader_test_utils::local_test_state_reader_factory;
 use crate::stateful_transaction_validator::{
     StatefulTransactionValidator,
@@ -156,7 +155,7 @@ async fn test_extract_state_nonce_and_run_validations(
 }
 
 #[rstest]
-#[tokio::test(flavor = "multi_thread")]
+#[tokio::test]
 async fn test_instantiate_validator() {
     let stateful_validator_factory = StatefulTransactionValidatorFactory {
         config: StatefulTransactionValidatorConfig::default(),
@@ -166,19 +165,9 @@ async fn test_instantiate_validator() {
     let state_reader_factory =
         local_test_state_reader_factory(CairoVersion::Cairo1(RunnableCairo1::Casm), false);
 
-    let mut mock_state_reader_factory = MockStateReaderFactory::new();
-
-    // Make sure stateful_validator uses the latest block in the initial call.
-    let latest_state_reader = state_reader_factory.get_state_reader_from_latest_block().await;
-    mock_state_reader_factory
-        .expect_get_state_reader_from_latest_block()
-        .return_once(move || latest_state_reader);
+    let validator =
+        stateful_validator_factory.instantiate_validator(Arc::new(state_reader_factory)).await;
 
-    // TODO(Itamar): Remove using runtime when instantiate_validator is async.
-    let validator = tokio::task::block_in_place(|| {
-        stateful_validator_factory
-            .instantiate_validator(&mock_state_reader_factory, tokio::runtime::Handle::current())
-    });
     assert!(validator.is_ok());
 }