apollo_gateway: move get_block_info to be async

Author: itamar-starkware

Cargo.lock

@@ -35,7 +35,6 @@ blockifier.workspace = true
 blockifier_test_utils = { workspace = true, optional = true }
 cairo-lang-starknet-classes.workspace = true
 clap.workspace = true
-futures.workspace = true
 lazy_static.workspace = true
 mempool_test_utils.workspace = true
 num-rational.workspace = true

crates/apollo_gateway/Cargo.toml

@@ -100,15 +100,20 @@ impl RpcStateReader {
     }
 }
 
+#[async_trait]
 impl MempoolStateReader for RpcStateReader {
-    fn get_block_info(&self) -> StateResult<BlockInfo> {
+    async fn get_block_info(&self) -> StateResult<BlockInfo> {
         let get_block_params = GetBlockWithTxHashesParams { block_id: self.block_id };
+        let reader = self.clone();
+
+        let get_block_with_tx_hashes_result = tokio::task::spawn_blocking(move || {
+            reader.send_rpc_request("starknet_getBlockWithTxHashes", get_block_params)
+        })
+        .await
+        .map_err(|e| StateError::StateReadError(format!("JoinError: {e}")))??;
 
-        // The response from the rpc is a full block but we only deserialize the header.
-        let block_header: BlockHeader = serde_json::from_value(
-            self.send_rpc_request("starknet_getBlockWithTxHashes", get_block_params)?,
-        )
-        .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/rpc_state_reader.rs

@@ -92,8 +92,7 @@ async fn test_get_block_info() {
     );
 
     let client = RpcStateReader::from_latest(&config);
-    let result =
-        tokio::task::spawn_blocking(move || client.get_block_info()).await.unwrap().unwrap();
+    let result = client.get_block_info().await.unwrap();
     assert_eq!(result, expected_result);
     mock.assert_async().await;
 }

crates/apollo_gateway/src/rpc_state_reader_test.rs

@@ -1,7 +1,6 @@
 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;
 use blockifier::state::state_api::{StateReader as BlockifierStateReader, StateResult};
 use blockifier::state::state_reader_and_contract_manager::{
@@ -15,8 +14,9 @@ use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce};
 use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
 
+#[async_trait]
 pub trait MempoolStateReader: BlockifierStateReader + Send + Sync {
-    fn get_block_info(&self) -> Result<BlockInfo, StateError>;
+    async fn get_block_info(&self) -> StateResult<BlockInfo>;
 }
 
 #[cfg_attr(test, automock)]
@@ -69,18 +69,20 @@ impl FetchCompiledClasses for Box<dyn GatewayStateReaderWithCompiledClasses> {
 // Therefore, for using the Box<dyn GatewayStateReaderWithCompiledClasses>, that the
 // StateReaderFactory creates, we need to implement the MempoolStateReader trait for
 // Box<dyn GatewayStateReaderWithCompiledClasses>.
+#[async_trait]
 impl MempoolStateReader for Box<dyn GatewayStateReaderWithCompiledClasses> {
-    fn get_block_info(&self) -> StateResult<BlockInfo> {
-        self.as_ref().get_block_info()
+    async fn get_block_info(&self) -> StateResult<BlockInfo> {
+        self.as_ref().get_block_info().await
     }
 }
 
 impl GatewayStateReaderWithCompiledClasses for Box<dyn GatewayStateReaderWithCompiledClasses> {}
 
+#[async_trait]
 impl MempoolStateReader
     for StateReaderAndContractManager<Box<dyn GatewayStateReaderWithCompiledClasses>>
 {
-    fn get_block_info(&self) -> StateResult<BlockInfo> {
-        self.state_reader.get_block_info()
+    async fn get_block_info(&self) -> StateResult<BlockInfo> {
+        self.state_reader.get_block_info().await
     }
 }

crates/apollo_gateway/src/state_reader.rs

@@ -103,7 +103,7 @@ async fn test_get_block_info() {
         block_number,
         tokio::runtime::Handle::current(),
     );
-    let result = state_reader_and_contract_manager.get_block_info().unwrap();
+    let result = state_reader_and_contract_manager.get_block_info().await.unwrap();
 
     assert_eq!(
         result,

crates/apollo_gateway/src/state_reader_test.rs

@@ -29,8 +29,9 @@ pub struct TestStateReader {
     pub blockifier_state_reader: DictStateReader,
 }
 
+#[async_trait]
 impl MempoolStateReader for TestStateReader {
-    fn get_block_info(&self) -> Result<BlockInfo, StateError> {
+    async fn get_block_info(&self) -> Result<BlockInfo, StateError> {
         Ok(self.block_info.clone())
     }
 }

crates/apollo_gateway/src/state_reader_test_utils.rs

@@ -80,7 +80,7 @@ impl StatefulTransactionValidatorFactoryTrait for StatefulTransactionValidatorFa
                     e,
                 )
             })?;
-        let latest_block_info = get_latest_block_info(&state_reader)?;
+        let latest_block_info = get_latest_block_info(&state_reader, runtime)?;
 
         let state_reader_and_contract_manager = StateReaderAndContractManager {
             state_reader,
@@ -349,10 +349,12 @@ fn skip_stateful_validations(
     Ok(false)
 }
 
+// TODO(Itamar): Remove using runtime here and make it async.
 pub fn get_latest_block_info(
     state_reader: &dyn MempoolStateReader,
+    runtime: tokio::runtime::Handle,
 ) -> StatefulTransactionValidatorResult<BlockInfo> {
-    state_reader
-        .get_block_info()
+    runtime
+        .block_on(state_reader.get_block_info())
         .map_err(|e| StarknetError::internal_with_logging("Failed to get latest block info", e))
 }

crates/apollo_gateway/src/stateful_transaction_validator.rs

@@ -23,7 +23,6 @@ use blockifier::state::errors::StateError;
 use blockifier::state::global_cache::CompiledClasses;
 use blockifier::state::state_api::{StateReader as BlockifierStateReader, StateResult};
 use blockifier::state::state_reader_and_contract_manager::FetchCompiledClasses;
-use futures::executor::block_on;
 use starknet_api::block::{BlockHash, BlockInfo, BlockNumber, GasPriceVector, GasPrices};
 use starknet_api::contract_class::ContractClass;
 use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce};
@@ -98,9 +97,13 @@ impl SyncStateReader {
     }
 }
 
+#[async_trait]
 impl MempoolStateReader for SyncStateReader {
-    fn get_block_info(&self) -> StateResult<BlockInfo> {
-        let block = block_on(self.state_sync_client.get_block(self.block_number))
+    async fn get_block_info(&self) -> StateResult<BlockInfo> {
+        let block = self
+            .state_sync_client
+            .get_block(self.block_number)
+            .await
             .map_err(|e| StateError::StateReadError(e.to_string()))?;
 
         let block_header = block.block_header_without_hash;

crates/apollo_gateway/src/sync_state_reader.rs

@@ -72,7 +72,7 @@ async fn test_get_block_info() {
         block_number,
         tokio::runtime::Handle::current(),
     );
-    let result = state_sync_reader.get_block_info().unwrap();
+    let result = state_sync_reader.get_block_info().await.unwrap();
 
     assert_eq!(
         result,