apollo_gateway: add GatewayFixedBlockStateReader instead of mempool state reader (#10486)

Author: itamar-starkware

crates/apollo_gateway/src/gateway_fixed_block_state_reader.rs

@@ -0,0 +1,70 @@
+use std::sync::Arc;
+
+use apollo_gateway_types::deprecated_gateway_error::StarknetError;
+use apollo_state_sync_types::communication::{SharedStateSyncClient, StateSyncClient};
+use async_trait::async_trait;
+use starknet_api::block::{BlockInfo, BlockNumber, GasPriceVector, GasPrices};
+use starknet_api::data_availability::L1DataAvailabilityMode;
+use tokio::sync::OnceCell;
+
+pub type StarknetResult<T> = Result<T, StarknetError>;
+
+/// A reader to a fixed block in the synced state of Starknet.
+#[cfg_attr(test, mockall::automock)]
+#[async_trait]
+pub trait GatewayFixedBlockStateReader: Send + Sync {
+    async fn get_block_info(&self) -> StarknetResult<BlockInfo>;
+}
+
+pub struct GatewayFixedBlockSyncStateClient {
+    state_sync_client: Arc<dyn StateSyncClient>,
+    block_number: BlockNumber,
+    block_info_cache: OnceCell<BlockInfo>,
+}
+
+impl GatewayFixedBlockSyncStateClient {
+    pub fn new(state_sync_client: SharedStateSyncClient, block_number: BlockNumber) -> Self {
+        Self { state_sync_client, block_number, block_info_cache: OnceCell::new() }
+    }
+
+    async fn get_block_info_from_sync_client(&self) -> StarknetResult<BlockInfo> {
+        let block = self.state_sync_client.get_block(self.block_number).await.map_err(|e| {
+            StarknetError::internal_with_logging("Failed to get latest block info", e)
+        })?;
+
+        let block_header = block.block_header_without_hash;
+        let block_info = BlockInfo {
+            block_number: block_header.block_number,
+            block_timestamp: block_header.timestamp,
+            sequencer_address: block_header.sequencer.0,
+            gas_prices: GasPrices {
+                eth_gas_prices: GasPriceVector {
+                    l1_gas_price: block_header.l1_gas_price.price_in_wei.try_into()?,
+                    l1_data_gas_price: block_header.l1_data_gas_price.price_in_wei.try_into()?,
+                    l2_gas_price: block_header.l2_gas_price.price_in_wei.try_into()?,
+                },
+                strk_gas_prices: GasPriceVector {
+                    l1_gas_price: block_header.l1_gas_price.price_in_fri.try_into()?,
+                    l1_data_gas_price: block_header.l1_data_gas_price.price_in_fri.try_into()?,
+                    l2_gas_price: block_header.l2_gas_price.price_in_fri.try_into()?,
+                },
+            },
+            use_kzg_da: match block_header.l1_da_mode {
+                L1DataAvailabilityMode::Blob => true,
+                L1DataAvailabilityMode::Calldata => false,
+            },
+        };
+
+        Ok(block_info)
+    }
+}
+
+#[async_trait]
+impl GatewayFixedBlockStateReader for GatewayFixedBlockSyncStateClient {
+    async fn get_block_info(&self) -> StarknetResult<BlockInfo> {
+        self.block_info_cache
+            .get_or_try_init(|| self.get_block_info_from_sync_client())
+            .await
+            .cloned()
+    }
+}

crates/apollo_gateway/src/lib.rs

@@ -1,6 +1,7 @@
 pub mod communication;
 pub mod errors;
 pub mod gateway;
+pub mod gateway_fixed_block_state_reader;
 pub mod metrics;
 pub mod rpc_objects;
 pub mod rpc_state_reader;

crates/apollo_gateway/src/rpc_state_reader.rs

@@ -1,4 +1,5 @@
 use apollo_gateway_config::config::RpcStateReaderConfig;
+use apollo_gateway_types::deprecated_gateway_error::StarknetError;
 use apollo_rpc::CompiledContractClass;
 use apollo_state_sync_types::communication::StateSyncClientResult;
 use async_trait::async_trait;
@@ -21,6 +22,7 @@ use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
 
 use crate::errors::{serde_err_to_state_err, RPCStateReaderError, RPCStateReaderResult};
+use crate::gateway_fixed_block_state_reader::{GatewayFixedBlockStateReader, StarknetResult};
 use crate::rpc_objects::{
     BlockHeader,
     BlockId,
@@ -35,11 +37,7 @@ use crate::rpc_objects::{
     RPC_ERROR_CONTRACT_ADDRESS_NOT_FOUND,
     RPC_ERROR_INVALID_PARAMS,
 };
-use crate::state_reader::{
-    GatewayStateReaderWithCompiledClasses,
-    MempoolStateReader,
-    StateReaderFactory,
-};
+use crate::state_reader::{GatewayStateReaderWithCompiledClasses, StateReaderFactory};
 
 #[derive(Clone)]
 pub struct RpcStateReader {
@@ -100,25 +98,6 @@ impl RpcStateReader {
     }
 }
 
-#[async_trait]
-impl MempoolStateReader for RpcStateReader {
-    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}")))??;
-
-        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)
-    }
-}
-
 impl BlockifierStateReader for RpcStateReader {
     fn get_storage_at(
         &self,
@@ -209,11 +188,40 @@ impl FetchCompiledClasses for RpcStateReader {
 
 impl GatewayStateReaderWithCompiledClasses for RpcStateReader {}
 
+// An alternative impl to get the synced state, not via the state-sync component. To be removed in
+// the future.
+#[async_trait]
+impl GatewayFixedBlockStateReader for RpcStateReader {
+    async fn get_block_info(&self) -> StarknetResult<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| StarknetError::internal_with_logging("JoinError", e))?
+        .map_err(|e| StarknetError::internal_with_logging("RPC error", e))?;
+
+        let block_header: BlockHeader = serde_json::from_value(get_block_with_tx_hashes_result)
+            .map_err(|e| StarknetError::internal_with_logging("Failed to parse block header", e))?;
+        let block_info = block_header.try_into().map_err(|e| {
+            StarknetError::internal_with_logging("Failed to convert block header to BlockInfo", e)
+        })?;
+        Ok(block_info)
+    }
+}
+
 #[async_trait]
 impl StateReaderFactory for RpcStateReaderFactory {
-    async fn get_state_reader_from_latest_block(
+    async fn get_blockifier_state_reader_and_gateway_fixed_block_from_latest_block(
         &self,
-    ) -> StateSyncClientResult<Box<dyn GatewayStateReaderWithCompiledClasses>> {
-        Ok(Box::new(RpcStateReader::from_latest(&self.config)))
+    ) -> StateSyncClientResult<(
+        Box<dyn GatewayStateReaderWithCompiledClasses>,
+        Box<dyn GatewayFixedBlockStateReader>,
+    )> {
+        let reader = RpcStateReader::from_latest(&self.config);
+        let fixed_client = RpcStateReader::from_latest(&self.config);
+        Ok((Box::new(reader), Box::new(fixed_client)))
     }
 }

crates/apollo_gateway/src/rpc_state_reader_test.rs

@@ -10,6 +10,7 @@ use starknet_api::block::{BlockInfo, BlockNumber, GasPricePerToken};
 use starknet_api::contract_class::SierraVersion;
 use starknet_api::{class_hash, contract_address, felt, nonce};
 
+use crate::gateway_fixed_block_state_reader::GatewayFixedBlockStateReader;
 use crate::rpc_objects::{
     BlockHeader,
     BlockId,
@@ -22,7 +23,6 @@ use crate::rpc_objects::{
     RpcSuccessResponse,
 };
 use crate::rpc_state_reader::RpcStateReader;
-use crate::state_reader::MempoolStateReader;
 
 async fn run_rpc_server() -> mockito::ServerGuard {
     mockito::Server::new_async().await

crates/apollo_gateway/src/state_reader.rs

@@ -3,31 +3,28 @@ use async_trait::async_trait;
 use blockifier::execution::contract_class::RunnableCompiledClass;
 use blockifier::state::global_cache::CompiledClasses;
 use blockifier::state::state_api::{StateReader as BlockifierStateReader, StateResult};
-use blockifier::state::state_reader_and_contract_manager::{
-    FetchCompiledClasses,
-    StateReaderAndContractManager,
-};
+use blockifier::state::state_reader_and_contract_manager::FetchCompiledClasses;
 #[cfg(test)]
 use mockall::automock;
-use starknet_api::block::BlockInfo;
 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 {
-    async fn get_block_info(&self) -> StateResult<BlockInfo>;
-}
-
+use crate::gateway_fixed_block_state_reader::GatewayFixedBlockStateReader;
 #[cfg_attr(test, automock)]
 #[async_trait]
 pub trait StateReaderFactory: Send + Sync {
-    async fn get_state_reader_from_latest_block(
+    async fn get_blockifier_state_reader_and_gateway_fixed_block_from_latest_block(
         &self,
-    ) -> StateSyncClientResult<Box<dyn GatewayStateReaderWithCompiledClasses>>;
+    ) -> StateSyncClientResult<(
+        Box<dyn GatewayStateReaderWithCompiledClasses>,
+        Box<dyn GatewayFixedBlockStateReader>,
+    )>;
 }
 
-pub trait GatewayStateReaderWithCompiledClasses: MempoolStateReader + FetchCompiledClasses {}
+// TODO(Arni): Delete this trait, once we replace `dyn GatewayStateReaderWithCompiledClasses` with
+// generics.
+pub trait GatewayStateReaderWithCompiledClasses: FetchCompiledClasses + Send + Sync {}
 
 impl BlockifierStateReader for Box<dyn GatewayStateReaderWithCompiledClasses> {
     fn get_storage_at(
@@ -64,25 +61,3 @@ impl FetchCompiledClasses for Box<dyn GatewayStateReaderWithCompiledClasses> {
         self.as_ref().is_declared(class_hash)
     }
 }
-
-// By default, a Box<dyn Trait> does not implement the trait of the object it contains.
-// 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> {
-    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>>
-{
-    async fn get_block_info(&self) -> StateResult<BlockInfo> {
-        self.state_reader.get_block_info().await
-    }
-}

crates/apollo_gateway/src/state_reader_test.rs

@@ -11,8 +11,6 @@ use apollo_state_sync_types::communication::{
     SharedStateSyncClient,
     StateSyncClientResult,
 };
-use apollo_state_sync_types::state_sync_types::SyncBlock;
-use apollo_test_utils::{get_rng, GetTestInstance};
 use blockifier::blockifier::config::ContractClassManagerConfig;
 use blockifier::execution::contract_class::RunnableCompiledClass;
 use blockifier::state::contract_class_manager::ContractClassManager;
@@ -23,23 +21,13 @@ use blockifier::state::state_reader_and_contract_manager::StateReaderAndContract
 use blockifier::test_utils::initial_test_state::state_reader_and_contract_manager_for_testing;
 use mockall::predicate;
 use rstest::rstest;
-use starknet_api::block::{
-    BlockHeaderWithoutHash,
-    BlockInfo,
-    BlockNumber,
-    BlockTimestamp,
-    GasPricePerToken,
-    GasPriceVector,
-    GasPrices,
-    NonzeroGasPrice,
-};
+use starknet_api::block::BlockNumber;
 use starknet_api::contract_class::ContractClass;
-use starknet_api::core::{ClassHash, SequencerContractAddress};
-use starknet_api::data_availability::L1DataAvailabilityMode;
+use starknet_api::core::ClassHash;
 use starknet_api::state::SierraContractClass;
 use starknet_api::{class_hash, contract_address, felt, nonce, storage_key};
 
-use crate::state_reader::{GatewayStateReaderWithCompiledClasses, MempoolStateReader};
+use crate::state_reader::GatewayStateReaderWithCompiledClasses;
 use crate::sync_state_reader::SyncStateReader;
 
 static DUMMY_CLASS_HASH: LazyLock<ClassHash> = LazyLock::new(|| class_hash!(2_u32));
@@ -249,79 +237,6 @@ fn declared_but_not_in_manager_scenario() -> GetCompiledClassTestScenario {
     }
 }
 
-#[tokio::test]
-async fn test_get_block_info() {
-    let mut mock_state_sync_client = MockStateSyncClient::new();
-    let mock_class_manager_client = MockClassManagerClient::new();
-    let contract_class_manager = ContractClassManager::start(ContractClassManagerConfig::default());
-    let block_number = BlockNumber(1);
-    let block_timestamp = BlockTimestamp(2);
-    let sequencer_address = contract_address!("0x3");
-    let l1_gas_price = GasPricePerToken { price_in_wei: 4_u8.into(), price_in_fri: 5_u8.into() };
-    let l1_data_gas_price =
-        GasPricePerToken { price_in_wei: 6_u8.into(), price_in_fri: 7_u8.into() };
-    let l2_gas_price = GasPricePerToken { price_in_wei: 8_u8.into(), price_in_fri: 9_u8.into() };
-    let l1_da_mode = L1DataAvailabilityMode::get_test_instance(&mut get_rng());
-
-    mock_state_sync_client.expect_get_block().times(1).with(predicate::eq(block_number)).returning(
-        move |_| {
-            Ok(SyncBlock {
-                state_diff: Default::default(),
-                account_transaction_hashes: Default::default(),
-                l1_transaction_hashes: Default::default(),
-                block_header_without_hash: BlockHeaderWithoutHash {
-                    block_number,
-                    l1_gas_price,
-                    l1_data_gas_price,
-                    l2_gas_price,
-                    sequencer: SequencerContractAddress(sequencer_address),
-                    timestamp: block_timestamp,
-                    l1_da_mode,
-                    ..Default::default()
-                },
-            })
-        },
-    );
-
-    let state_reader_and_contract_manager = sync_state_reader_and_contract_manager(
-        Arc::new(mock_state_sync_client),
-        Arc::new(mock_class_manager_client),
-        contract_class_manager.clone(),
-        block_number,
-        tokio::runtime::Handle::current(),
-    );
-    let result = state_reader_and_contract_manager.get_block_info().await.unwrap();
-
-    assert_eq!(
-        result,
-        BlockInfo {
-            block_number,
-            block_timestamp,
-            sequencer_address,
-            gas_prices: GasPrices {
-                eth_gas_prices: GasPriceVector {
-                    l1_gas_price: NonzeroGasPrice::new_unchecked(l1_gas_price.price_in_wei),
-                    l1_data_gas_price: NonzeroGasPrice::new_unchecked(
-                        l1_data_gas_price.price_in_wei
-                    ),
-                    l2_gas_price: NonzeroGasPrice::new_unchecked(l2_gas_price.price_in_wei),
-                },
-                strk_gas_prices: GasPriceVector {
-                    l1_gas_price: NonzeroGasPrice::new_unchecked(l1_gas_price.price_in_fri),
-                    l1_data_gas_price: NonzeroGasPrice::new_unchecked(
-                        l1_data_gas_price.price_in_fri
-                    ),
-                    l2_gas_price: NonzeroGasPrice::new_unchecked(l2_gas_price.price_in_fri),
-                },
-            },
-            use_kzg_da: match l1_da_mode {
-                L1DataAvailabilityMode::Blob => true,
-                L1DataAvailabilityMode::Calldata => false,
-            },
-        }
-    );
-}
-
 #[tokio::test]
 async fn test_get_storage_at() {
     let mut mock_state_sync_client = MockStateSyncClient::new();