apollo_gateway: move get_state_reader_from_latest_block to be async (#10119)

Author: itamar-starkware

crates/apollo_gateway/src/gateway.rs

@@ -258,7 +258,7 @@ 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())?;
+            .instantiate_validator(self.state_reader_factory.as_ref(), self.runtime.clone())?;
 
         let nonce = stateful_transaction_validator.extract_state_nonce_and_run_validations(
             &self.executable_tx,

crates/apollo_gateway/src/gateway_test.rs

@@ -565,7 +565,7 @@ async fn process_tx_returns_error_when_extract_state_nonce_and_run_validations_f
 
     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 process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory);
 
@@ -608,7 +608,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(|_, _| Err(expected_error));
 
     let process_tx_task = process_tx_task(mock_stateful_transaction_validator_factory);
 

crates/apollo_gateway/src/rpc_state_reader.rs

@@ -1,6 +1,7 @@
 use apollo_gateway_config::config::RpcStateReaderConfig;
 use apollo_rpc::CompiledContractClass;
 use apollo_state_sync_types::communication::StateSyncClientResult;
+use async_trait::async_trait;
 use blockifier::execution::contract_class::{
     CompiledClassV0,
     CompiledClassV1,
@@ -203,8 +204,9 @@ impl FetchCompiledClasses for RpcStateReader {
 
 impl GatewayStateReaderWithCompiledClasses for RpcStateReader {}
 
+#[async_trait]
 impl StateReaderFactory for RpcStateReaderFactory {
-    fn get_state_reader_from_latest_block(
+    async fn get_state_reader_from_latest_block(
         &self,
     ) -> StateSyncClientResult<Box<dyn GatewayStateReaderWithCompiledClasses>> {
         Ok(Box::new(RpcStateReader::from_latest(&self.config)))

crates/apollo_gateway/src/state_reader.rs

@@ -1,4 +1,5 @@
 use apollo_state_sync_types::communication::StateSyncClientResult;
+use async_trait::async_trait;
 use blockifier::execution::contract_class::RunnableCompiledClass;
 use blockifier::state::errors::StateError;
 use blockifier::state::global_cache::CompiledClasses;
@@ -19,8 +20,9 @@ pub trait MempoolStateReader: BlockifierStateReader + Send + Sync {
 }
 
 #[cfg_attr(test, automock)]
+#[async_trait]
 pub trait StateReaderFactory: Send + Sync {
-    fn get_state_reader_from_latest_block(
+    async fn get_state_reader_from_latest_block(
         &self,
     ) -> StateSyncClientResult<Box<dyn GatewayStateReaderWithCompiledClasses>>;
 }

crates/apollo_gateway/src/state_reader_test_utils.rs

@@ -1,4 +1,5 @@
 use apollo_state_sync_types::communication::StateSyncClientResult;
+use axum::async_trait;
 use blockifier::context::BlockContext;
 use blockifier::execution::contract_class::RunnableCompiledClass;
 use blockifier::state::errors::StateError;
@@ -76,8 +77,9 @@ pub struct TestStateReaderFactory {
     pub state_reader: TestStateReader,
 }
 
+#[async_trait]
 impl StateReaderFactory for TestStateReaderFactory {
-    fn get_state_reader_from_latest_block(
+    async fn get_state_reader_from_latest_block(
         &self,
     ) -> StateSyncClientResult<Box<dyn GatewayStateReaderWithCompiledClasses>> {
         Ok(Box::new(self.state_reader.clone()))

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -51,6 +51,7 @@ pub trait StatefulTransactionValidatorFactoryTrait: Send + Sync {
     fn instantiate_validator(
         &self,
         state_reader_factory: &dyn StateReaderFactory,
+        runtime: tokio::runtime::Handle,
     ) -> StatefulTransactionValidatorResult<Box<dyn StatefulTransactionValidatorTrait>>;
 }
 pub struct StatefulTransactionValidatorFactory {
@@ -64,11 +65,12 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
     fn instantiate_validator(
         &self,
         state_reader_factory: &dyn StateReaderFactory,
+        runtime: tokio::runtime::Handle,
     ) -> 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
-            .get_state_reader_from_latest_block()
+        let state_reader = runtime
+            .block_on(state_reader_factory.get_state_reader_from_latest_block())
             .map_err(|err| GatewaySpecError::UnexpectedError {
                 data: format!("Internal server error: {err}"),
             })

crates/apollo_gateway/src/stateful_transaction_validator_test.rs

@@ -156,7 +156,8 @@ async fn test_extract_state_nonce_and_run_validations(
 }
 
 #[rstest]
-fn test_instantiate_validator() {
+#[tokio::test(flavor = "multi_thread")]
+async fn test_instantiate_validator() {
     let stateful_validator_factory = StatefulTransactionValidatorFactory {
         config: StatefulTransactionValidatorConfig::default(),
         chain_info: ChainInfo::create_for_testing(),
@@ -168,12 +169,16 @@ fn test_instantiate_validator() {
     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();
+    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(|| latest_state_reader);
+        .return_once(move || latest_state_reader);
 
-    let validator = stateful_validator_factory.instantiate_validator(&mock_state_reader_factory);
+    // 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());
 }
 

crates/apollo_gateway/src/sync_state_reader.rs

@@ -353,14 +353,18 @@ pub(crate) struct SyncStateReaderFactory {
 }
 
 /// Use any of these factory methods only once per transaction to make sure metrics are accurate.
+#[async_trait]
 impl StateReaderFactory for SyncStateReaderFactory {
-    fn get_state_reader_from_latest_block(
+    // TODO(guy.f): The call to `get_latest_block_number()` is not counted in the storage metrics as
+    // it is done prior to the creation of SharedStateSyncClientMetricWrapper, directly via the
+    // SharedStateSyncClient.
+    async fn get_state_reader_from_latest_block(
         &self,
     ) -> StateSyncClientResult<Box<dyn GatewayStateReaderWithCompiledClasses>> {
         let latest_block_number = self
-            .runtime
-            // TODO(guy.f): Do we want to count this as well?
-            .block_on(self.shared_state_sync_client.get_latest_block_number())?
+            .shared_state_sync_client
+            .get_latest_block_number()
+            .await?
             .ok_or(StateSyncClientError::StateSyncError(StateSyncError::EmptyState))?;
 
         Ok(Box::new(SyncStateReader::from_number(