apollo_state_sync: add storage_reader_server to the state sync

Author: nadin-Starkware

crates/apollo_state_sync/src/lib.rs

@@ -12,7 +12,13 @@ use apollo_starknet_client::reader::{StarknetFeederGatewayClient, StarknetReader
 use apollo_state_sync_config::config::StateSyncConfig;
 use apollo_state_sync_types::communication::{StateSyncRequest, StateSyncResponse};
 use apollo_state_sync_types::errors::StateSyncError;
-use apollo_state_sync_types::state_sync_types::{StateSyncResult, SyncBlock};
+use apollo_state_sync_types::state_sync_types::{
+    StateSyncResult,
+    StateSyncStorageReaderServerHandler,
+    StateSyncStorageRequest,
+    StateSyncStorageResponse,
+    SyncBlock,
+};
 use apollo_storage::body::BodyStorageReader;
 use apollo_storage::db::TransactionKind;
 use apollo_storage::header::HeaderStorageReader;
@@ -30,6 +36,11 @@ use starknet_types_core::felt::Felt;
 use crate::runner::StateSyncRunner;
 
 const BUFFER_SIZE: usize = 100000;
+type StateSyncStorageReaderServer = StorageReaderServer<
+    StateSyncStorageReaderServerHandler,
+    StateSyncStorageRequest,
+    StateSyncStorageResponse,
+>;
 
 pub fn create_state_sync_and_runner(
     config: StateSyncConfig,
@@ -46,6 +57,10 @@ pub struct StateSync {
     storage_reader: StorageReader,
     new_block_sender: Sender<SyncBlock>,
     starknet_client: Option<Arc<dyn StarknetReader + Send + Sync>>,
+    // TODO(Nadin): Remove #[allow(dead_code)].
+    /// Optional storage reader server for handling remote storage reader queries.
+    #[allow(dead_code)]
+    storage_reader_server: Option<StateSyncStorageReaderServer>,
 }
 
 impl StateSync {
@@ -69,7 +84,12 @@ impl StateSync {
             );
             starknet_client
         });
-        Self { storage_reader, new_block_sender, starknet_client }
+        Self {
+            storage_reader,
+            new_block_sender,
+            starknet_client,
+            storage_reader_server: config.storage_reader_server,
+        }
     }
 }
 

crates/apollo_state_sync/src/runner/mod.rs

@@ -101,10 +101,16 @@ pub struct StateSyncResources {
 }
 
 impl StateSyncResources {
-    pub fn new(storage_config: &StorageConfig) -> Self {
-        let (storage_reader, storage_writer) =
-            open_storage_with_metric(storage_config.clone(), &SYNC_STORAGE_OPEN_READ_TRANSACTIONS)
-                .expect("StateSyncRunner failed opening storage");
+    pub fn new(
+        storage_config: &StorageConfig,
+        storage_reader_server_config: Option<StateSyncStorageReaderServer>,
+    ) -> Self {
+        let (storage_reader, storage_writer) = open_storage_with_metric_and_server(
+            storage_config.clone(),
+            &SYNC_STORAGE_OPEN_READ_TRANSACTIONS,
+            storage_reader_server_config,
+        )
+        .expect("StateSyncRunner failed opening storage");
         let shared_highest_block = Arc::new(RwLock::new(None));
         let pending_data = Arc::new(RwLock::new(PendingData {
             // The pending data might change later to DeprecatedPendingBlock, depending on the
@@ -133,7 +139,7 @@ impl StateSyncRunner {
             network_config,
             revert_config,
             rpc_config,
-            storage_reader_server_config: _,
+            storage_reader_server_config,
         } = config;
 
         let StateSyncResources {
@@ -142,7 +148,7 @@ impl StateSyncRunner {
             shared_highest_block,
             pending_data,
             pending_classes,
-        } = StateSyncResources::new(&storage_config);
+        } = StateSyncResources::new(&storage_config, storage_reader_server_config);
 
         let register_metrics_future = register_metrics(storage_reader.clone()).boxed();
 

crates/apollo_state_sync/src/test.rs

@@ -24,8 +24,12 @@ use crate::StateSync;
 
 fn setup() -> (StateSync, StorageWriter) {
     let ((storage_reader, storage_writer), _) = get_test_storage();
-    let state_sync =
-        StateSync { storage_reader, new_block_sender: channel(0).0, starknet_client: None };
+    let state_sync = StateSync {
+        storage_reader,
+        new_block_sender: channel(0).0,
+        starknet_client: None,
+        storage_reader_server: None,
+    };
     (state_sync, storage_writer)
 }
 
@@ -125,8 +129,12 @@ async fn test_get_block_hash_fallback_to_starknet_client() {
     let starknet_client: Option<Arc<dyn StarknetReader + Send + Sync>> =
         Some(Arc::new(starknet_client));
     let ((storage_reader, _storage_writer), _) = get_test_storage();
-    let mut state_sync =
-        StateSync { storage_reader, new_block_sender: channel(0).0, starknet_client };
+    let mut state_sync = StateSync {
+        storage_reader,
+        new_block_sender: channel(0).0,
+        starknet_client,
+        storage_reader_server: None,
+    };
 
     // The block is not in storage, so it should fall back to starknet_client
     let response = state_sync.handle_request(StateSyncRequest::GetBlockHash(block_number)).await;