apollo_storage,apollo_batcher: spawn storage reader server independently

Author: nadin-Starkware

crates/apollo_batcher/src/batcher.rs

@@ -69,6 +69,7 @@ use starknet_api::core::{ContractAddress, GlobalRoot, Nonce};
 use starknet_api::state::{StateNumber, ThinStateDiff};
 use starknet_api::transaction::TransactionHash;
 use tokio::sync::Mutex;
+use tokio::task::JoinHandle;
 use tracing::{debug, error, info, instrument, trace, Instrument};
 
 use crate::block_builder::{
@@ -173,11 +174,10 @@ pub struct Batcher {
     /// This is returned by the decision_reached function.
     prev_proposal_commitment: Option<(BlockNumber, ProposalCommitment)>,
 
-    /// Optional storage reader server for handling remote storage reader queries.
-    /// Kept alive to maintain the server running.
-    #[allow(dead_code)]
-    storage_reader_server: Option<GenericStorageReaderServer>,
     commitment_manager: ApolloCommitmentManager,
+
+    /// Task handle for the storage reader server, if enabled.
+    storage_reader_server_handle: Option<JoinHandle<()>>,
 }
 
 impl Batcher {
@@ -192,8 +192,8 @@ impl Batcher {
         transaction_converter: TransactionConverter,
         block_builder_factory: Box<dyn BlockBuilderFactoryTrait>,
         pre_confirmed_block_writer_factory: Box<dyn PreconfirmedBlockWriterFactoryTrait>,
-        storage_reader_server: Option<GenericStorageReaderServer>,
         commitment_manager: ApolloCommitmentManager,
+        storage_reader_server_handle: Option<JoinHandle<()>>,
     ) -> Self {
         Self {
             config,
@@ -214,8 +214,8 @@ impl Batcher {
             // Allow the first few proposals to be without L1 txs while system starts up.
             proposals_counter: 1,
             prev_proposal_commitment: None,
-            storage_reader_server,
             commitment_manager,
+            storage_reader_server_handle,
         }
     }
 
@@ -1255,6 +1255,15 @@ fn log_txs_execution_result(
     }
 }
 
+impl Drop for Batcher {
+    fn drop(&mut self) {
+        // Abort the storage reader server task if it was spawned.
+        if let Some(handle) = self.storage_reader_server_handle.take() {
+            handle.abort();
+        }
+    }
+}
+
 pub async fn create_batcher(
     config: BatcherConfig,
     committer_client: SharedCommitterClient,
@@ -1271,6 +1280,9 @@ pub async fn create_batcher(
         )
         .expect("Failed to open batcher's storage");
 
+    let storage_reader_server_handle =
+        GenericStorageReaderServer::spawn_if_enabled(storage_reader_server);
+
     let execute_config = &config.block_builder_config.execute_config;
     let worker_pool = Arc::new(WorkerPool::start(execute_config));
     let pre_confirmed_block_writer_factory = Box::new(PreconfirmedBlockWriterFactory {
@@ -1310,8 +1322,8 @@ pub async fn create_batcher(
         transaction_converter,
         block_builder_factory,
         pre_confirmed_block_writer_factory,
-        storage_reader_server,
         commitment_manager,
+        storage_reader_server_handle,
     )
 }
 

crates/apollo_batcher/src/batcher_test.rs

@@ -194,8 +194,8 @@ async fn create_batcher_impl<R: BatcherStorageReader + 'static>(
         TransactionConverter::new(clients.class_manager_client, CHAIN_ID_FOR_TESTS.clone()),
         Box::new(clients.block_builder_factory),
         Box::new(clients.pre_confirmed_block_writer_factory),
-        None,
         commitment_manager,
+        None,
     );
     // Call post-creation functionality (e.g., metrics registration).
     batcher.start().await;

crates/apollo_storage/Cargo.toml

@@ -36,6 +36,7 @@ starknet-types-core = { workspace = true, features = ["papyrus-serialization"] }
 starknet_api.workspace = true
 tempfile = { workspace = true, optional = true }
 thiserror.workspace = true
+tokio.workspace = true
 tower = { workspace = true, optional = true }
 tracing = { workspace = true, features = ["log"] }
 validator = { workspace = true, features = ["derive"] }

crates/apollo_storage/src/storage_reader_server.rs

@@ -160,6 +160,22 @@ where
             StorageError::IOError(io::Error::other(e))
         })
     }
+
+    /// Spawns the storage reader server in a background task if it's enabled.
+    pub fn spawn_if_enabled(server: Option<Self>) -> Option<tokio::task::JoinHandle<()>>
+    where
+        RequestHandler: Send + Sync + 'static,
+        Request: Send + Sync + 'static,
+        Response: Send + Sync + 'static,
+    {
+        server.map(|server| {
+            tokio::spawn(async move {
+                if let Err(e) = server.run().await {
+                    tracing::error!("Storage reader server error: {:?}", e);
+                }
+            })
+        })
+    }
 }
 
 /// Axum handler for storage query requests.