apollo_gateway: move get_state_reader_from_latest_block to be async

Author: itamar-starkware

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,
@@ -185,8 +186,9 @@ pub struct RpcStateReaderFactory {
     pub config: RpcStateReaderConfig,
 }
 
+#[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 MempoolStateReader>> {
         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::state_api::{StateReader as BlockifierStateReader, StateResult};
@@ -14,8 +15,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 MempoolStateReader>>;
 }

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;
@@ -58,8 +59,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 MempoolStateReader>> {
         Ok(Box::new(self.state_reader.clone()))

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -58,8 +58,8 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
     ) -> 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 = tokio::runtime::Handle::current()
+            .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

@@ -157,7 +157,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,14 @@ 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);
+    let validator = tokio::task::block_in_place(|| {
+        stateful_validator_factory.instantiate_validator(&mock_state_reader_factory)
+    });
     assert!(validator.is_ok());
 }
 

crates/apollo_gateway/src/sync_state_reader.rs

@@ -281,14 +281,16 @@ 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(
+    async fn get_state_reader_from_latest_block(
         &self,
     ) -> StateSyncClientResult<Box<dyn MempoolStateReader>> {
+        // TODO(guy.f): Do we want to count this as well?
         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(