apollo_committer: add new request handler

Author: ArielElp

crates/apollo_batcher/Cargo.toml

@@ -8,7 +8,7 @@ description = "Block building and transaction batching component for the Starkne
 
 [features]
 cairo_native = ["blockifier/cairo_native"]
-os_input = ["apollo_reverts/os_input", "apollo_storage/os_input"]
+os_input = ["apollo_committer_types/os_input", "apollo_reverts/os_input", "apollo_storage/os_input"]
 testing = []
 
 [lints]

crates/apollo_batcher/src/commitment_manager/commitment_manager_impl.rs

@@ -551,6 +551,16 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
                     COMMITMENT_MANAGER_REVERT_BLOCK_LATENCY.increment(task_duration);
                     COMMITMENT_MANAGER_REVERT_BLOCK_COUNT.increment(1);
                 }
+                #[cfg(feature = "os_input")]
+                CommitterRequestLabelValue::ReadPathsAndCommitBlock => {
+                    debug!(
+                        "Read paths and commit block latency for block {height}: {task_duration} \
+                         milliseconds."
+                    );
+                    // TODO(Ariel): Add dedicated metrics once we use os_input in prod.
+                    COMMITMENT_MANAGER_COMMIT_BLOCK_LATENCY.increment(task_duration);
+                    COMMITMENT_MANAGER_COMMIT_BLOCK_COUNT.increment(1);
+                }
             }
         }
     }

crates/apollo_batcher/src/commitment_manager/types.rs

@@ -107,9 +107,17 @@ impl TaskTimer {
     pub(crate) fn start_timer(&mut self, task: CommitterRequestLabelValue, height: BlockNumber) {
         let instant = Instant::now();
         match task {
-            CommitterRequestLabelValue::CommitBlock => self.commit.insert(height, instant),
-            CommitterRequestLabelValue::RevertBlock => self.revert.insert(height, instant),
-        };
+            CommitterRequestLabelValue::CommitBlock => {
+                self.commit.insert(height, instant);
+            }
+            #[cfg(feature = "os_input")]
+            CommitterRequestLabelValue::ReadPathsAndCommitBlock => {
+                self.commit.insert(height, instant);
+            }
+            CommitterRequestLabelValue::RevertBlock => {
+                self.revert.insert(height, instant);
+            }
+        }
     }
 
     /// Returns the duration of the task in milliseconds.
@@ -120,6 +128,8 @@ impl TaskTimer {
     ) -> Option<u128> {
         let map = match task {
             CommitterRequestLabelValue::CommitBlock => &mut self.commit,
+            #[cfg(feature = "os_input")]
+            CommitterRequestLabelValue::ReadPathsAndCommitBlock => &mut self.commit,
             CommitterRequestLabelValue::RevertBlock => &mut self.revert,
         };
 

crates/apollo_committer/src/communication.rs

@@ -2,14 +2,20 @@ use apollo_committer_types::communication::{CommitterRequest, CommitterResponse}
 use apollo_infra::component_definitions::ComponentRequestHandler;
 use apollo_infra::component_server::{LocalComponentServer, RemoteComponentServer};
 use async_trait::async_trait;
+#[cfg(feature = "os_input")]
+use starknet_committer::db::forest_trait::forest_trait_witnesses::ForestStorageWithWitnesses;
+#[cfg(not(feature = "os_input"))]
 use starknet_committer::db::forest_trait::ForestStorageWithEmptyReadContext;
+#[cfg(feature = "os_input")]
+use starknet_patricia_storage::storage_trait::ImmutableReadOnlyStorage;
 
 use crate::committer::{ApolloCommitter, Committer, StorageConstructor};
 
 pub type LocalCommitterServer =
     LocalComponentServer<ApolloCommitter, CommitterRequest, CommitterResponse>;
 pub type RemoteCommitterServer = RemoteComponentServer<CommitterRequest, CommitterResponse>;
 
+#[cfg(not(feature = "os_input"))]
 #[async_trait]
 impl<S: StorageConstructor, ForestDB: ForestStorageWithEmptyReadContext<Storage = S>>
     ComponentRequestHandler<CommitterRequest, CommitterResponse> for Committer<S, ForestDB>
@@ -25,3 +31,28 @@ impl<S: StorageConstructor, ForestDB: ForestStorageWithEmptyReadContext<Storage
         }
     }
 }
+
+#[cfg(feature = "os_input")]
+#[async_trait]
+impl<S, ForestDB> ComponentRequestHandler<CommitterRequest, CommitterResponse>
+    for Committer<S, ForestDB>
+where
+    S: StorageConstructor + ImmutableReadOnlyStorage + 'static,
+    ForestDB: ForestStorageWithWitnesses<Storage = S>,
+{
+    async fn handle_request(&mut self, request: CommitterRequest) -> CommitterResponse {
+        match request {
+            CommitterRequest::CommitBlock(commit_block_request) => {
+                CommitterResponse::CommitBlock(self.commit_block(commit_block_request).await)
+            }
+            CommitterRequest::RevertBlock(revert_block_request) => {
+                CommitterResponse::RevertBlock(self.revert_block(revert_block_request).await)
+            }
+            CommitterRequest::ReadPathsAndCommitBlock(req) => {
+                CommitterResponse::ReadPathsAndCommitBlock(
+                    self.read_paths_and_commit_block(req).await,
+                )
+            }
+        }
+    }
+}

crates/apollo_committer_types/src/communication.rs

@@ -21,6 +21,8 @@ use crate::committer_types::{
     RevertBlockRequest,
     RevertBlockResponse,
 };
+#[cfg(feature = "os_input")]
+use crate::committer_types::{ReadPathsAndCommitBlockRequest, ReadPathsAndCommitBlockResponse};
 use crate::errors::{CommitterClientError, CommitterClientResult, CommitterResult};
 
 pub type LocalCommitterClient = LocalComponentClient<CommitterRequest, CommitterResponse>;
@@ -43,6 +45,14 @@ pub trait CommitterClient: Send + Sync {
         &self,
         input: RevertBlockRequest,
     ) -> CommitterClientResult<RevertBlockResponse>;
+
+    #[cfg(feature = "os_input")]
+    /// Applies the state diff, collects merged Patricia witnesses for OS input, and persists replay
+    /// data (digest + payload).
+    async fn read_paths_and_commit_block(
+        &self,
+        input: ReadPathsAndCommitBlockRequest,
+    ) -> CommitterClientResult<ReadPathsAndCommitBlockResponse>;
 }
 
 #[derive(Serialize, Deserialize, Clone, AsRefStr, EnumDiscriminants)]
@@ -54,6 +64,8 @@ pub trait CommitterClient: Send + Sync {
 pub enum CommitterRequest {
     CommitBlock(CommitBlockRequest),
     RevertBlock(RevertBlockRequest),
+    #[cfg(feature = "os_input")]
+    ReadPathsAndCommitBlock(ReadPathsAndCommitBlockRequest),
 }
 
 impl_debug_for_infra_requests_and_responses!(CommitterRequest);
@@ -64,6 +76,8 @@ impl PrioritizedRequest for CommitterRequest {}
 pub enum CommitterResponse {
     CommitBlock(CommitterResult<CommitBlockResponse>),
     RevertBlock(CommitterResult<RevertBlockResponse>),
+    #[cfg(feature = "os_input")]
+    ReadPathsAndCommitBlock(CommitterResult<ReadPathsAndCommitBlockResponse>),
 }
 
 impl_debug_for_infra_requests_and_responses!(CommitterResponse);
@@ -109,4 +123,21 @@ where
             Direct
         )
     }
+
+    #[cfg(feature = "os_input")]
+    async fn read_paths_and_commit_block(
+        &self,
+        input: ReadPathsAndCommitBlockRequest,
+    ) -> CommitterClientResult<ReadPathsAndCommitBlockResponse> {
+        let request = CommitterRequest::ReadPathsAndCommitBlock(input);
+        handle_all_response_variants!(
+            self,
+            request,
+            CommitterResponse,
+            ReadPathsAndCommitBlock,
+            CommitterClientError,
+            CommitterError,
+            Direct
+        )
+    }
 }

crates/apollo_committer_types/src/test_utils.rs

@@ -10,6 +10,8 @@ use crate::committer_types::{
     RevertBlockRequest,
     RevertBlockResponse,
 };
+#[cfg(feature = "os_input")]
+use crate::committer_types::{ReadPathsAndCommitBlockRequest, ReadPathsAndCommitBlockResponse};
 use crate::communication::{CommitterClient, MockCommitterClient};
 use crate::errors::CommitterClientResult;
 
@@ -39,6 +41,15 @@ impl CommitterClient for MockCommitterClientWithOffset {
         self.set_offset(input.height).await;
         self.inner.revert_block(input).await
     }
+
+    #[cfg(feature = "os_input")]
+    async fn read_paths_and_commit_block(
+        &self,
+        input: ReadPathsAndCommitBlockRequest,
+    ) -> CommitterClientResult<ReadPathsAndCommitBlockResponse> {
+        self.set_offset(input.commit.height).await;
+        self.inner.read_paths_and_commit_block(input).await
+    }
 }
 
 impl MockCommitterClientWithOffset {