apollo_gateway: move instantiate_validator to be async

Author: itamar-starkware

crates/apollo_gateway/src/gateway.rs

@@ -252,9 +252,10 @@ 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())?;
+        let mut stateful_transaction_validator = self.runtime.block_on(
+            self.stateful_tx_validator_factory
+                .instantiate_validator(self.state_reader_factory.clone()),
+        )?;
 
         let nonce = stateful_transaction_validator.extract_state_nonce_and_run_validations(
             &self.executable_tx,

crates/apollo_gateway/src/gateway_test.rs

@@ -97,6 +97,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;
 
@@ -561,9 +562,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);
 
@@ -606,7 +611,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/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,
@@ -38,29 +41,31 @@ mod stateful_transaction_validator_test;
 
 type BlockifierStatefulValidator = StatefulValidator<Box<dyn MempoolStateReader>>;
 
+#[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,
+        state_reader_factory: Arc<dyn StateReaderFactory>,
     ) -> StatefulTransactionValidatorResult<Box<dyn StatefulTransactionValidatorTrait>>;
 }
 pub struct StatefulTransactionValidatorFactory {
     pub config: StatefulTransactionValidatorConfig,
     pub chain_info: ChainInfo,
 }
 
+#[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,
+        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 current_runtime = tokio::runtime::Handle::current();
-        let state_reader = current_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}"),
             })
@@ -70,7 +75,7 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
                     e,
                 )
             })?;
-        let latest_block_info = current_runtime.block_on(get_latest_block_info(&state_reader))?;
+        let latest_block_info = get_latest_block_info(&state_reader).await?;
 
         let state = CachedState::new(state_reader);
         let mut versioned_constants = VersionedConstants::get_versioned_constants(

crates/apollo_gateway/src/stateful_transaction_validator_test.rs

@@ -174,9 +174,8 @@ async fn test_instantiate_validator() {
         .expect_get_state_reader_from_latest_block()
         .return_once(move || latest_state_reader);
 
-    let validator = tokio::task::block_in_place(|| {
-        stateful_validator_factory.instantiate_validator(&mock_state_reader_factory)
-    });
+    let validator =
+        stateful_validator_factory.instantiate_validator(Arc::new(mock_state_reader_factory)).await;
     assert!(validator.is_ok());
 }