apollo_batcher: cache recent state commitment infos in memory

crates/apollo_batcher/src/batcher.rs

@@ -1623,16 +1623,28 @@ impl Batcher {
     }
 
     pub fn get_state_commitment_infos(
-        &self,
+        &mut self,
         block_number: BlockNumber,
     ) -> BatcherResult<Option<CompressedStateCommitmentInfos>> {
+        self.get_commitment_results_and_write_to_storage()?;
+        if let Some(state_commitment_infos) =
+            self.commitment_manager.recent_state_commitment_infos_cache.get(&block_number)
+        {
+            return Ok(Some(state_commitment_infos.clone()));
+        }
+
         self.storage_reader.get_state_commitment_infos(block_number).map_err(|err| {
             error!("Failed to get state commitment infos from storage: {err}");
             BatcherError::InternalError
         })
     }
 
-    pub fn has_state_commitment_infos(&self, block_number: BlockNumber) -> BatcherResult<bool> {
+    pub fn has_state_commitment_infos(&mut self, block_number: BlockNumber) -> BatcherResult<bool> {
+        self.get_commitment_results_and_write_to_storage()?;
+        if self.commitment_manager.recent_state_commitment_infos_cache.contains(&block_number) {
+            return Ok(true);
+        }
+
         self.storage_reader.has_state_commitment_infos(block_number).map_err(|err| {
             error!("Failed to check state commitment infos existence in storage: {err}");
             BatcherError::InternalError

crates/apollo_batcher/src/batcher_test.rs

@@ -27,7 +27,7 @@ use apollo_batcher_types::batcher_types::{
 };
 use apollo_batcher_types::errors::BatcherError;
 use apollo_class_manager_types::MockClassManagerClient;
-use apollo_committer_types::committer_types::CommitBlockRequest;
+use apollo_committer_types::committer_types::{CommitBlockRequest, ReadPathsAndCommitBlockRequest};
 use apollo_config_manager_types::communication::MockConfigManagerClient;
 use apollo_infra::component_client::ClientError;
 use apollo_infra::component_definitions::ComponentStarter;
@@ -144,6 +144,7 @@ use crate::test_utils::{
     propose_block_input,
     test_contract_nonces,
     test_l1_handler_txs,
+    test_state_commitment_infos,
     test_state_diff,
     test_txs,
     verify_indexed_execution_infos,
@@ -1651,9 +1652,20 @@ async fn revert_block() {
 
     let revert_input = RevertBlockInput { height: LATEST_BLOCK_IN_STORAGE };
 
+    batcher
+        .commitment_manager
+        .recent_state_commitment_infos_cache
+        .put(LATEST_BLOCK_IN_STORAGE, test_state_commitment_infos());
+
     assert_eq!(*(committer_offset.lock().await), INITIAL_HEIGHT);
     batcher.revert_block(revert_input).await.unwrap();
     assert_eq!(*committer_offset.lock().await, LATEST_BLOCK_IN_STORAGE);
+    assert!(
+        !batcher
+            .commitment_manager
+            .recent_state_commitment_infos_cache
+            .contains(&LATEST_BLOCK_IN_STORAGE)
+    );
 
     let metrics = recorder.handle().render();
     assert_eq!(BUILDING_HEIGHT.parse_numeric_metric::<u64>(&metrics), Some(INITIAL_HEIGHT.0 - 1));
@@ -2019,6 +2031,50 @@ async fn get_block_hash_after_reading_commitment_results() {
     );
 }
 
+#[tokio::test]
+async fn get_state_commitment_infos_after_reading_commitment_results() {
+    let mut mock_dependencies = MockDependencies::default();
+    mock_dependencies
+        .storage_reader
+        .expect_get_parent_hash_and_partial_block_hash_components()
+        .with(eq(INITIAL_HEIGHT))
+        .returning(|_| {
+            Ok((
+                Some(BlockHash::default()),
+                Some(PartialBlockHashComponents {
+                    block_number: INITIAL_HEIGHT,
+                    ..Default::default()
+                }),
+            ))
+        });
+    mock_dependencies
+        .storage_writer
+        .expect_set_global_root_and_block_hash()
+        .times(1)
+        .returning(|_, _, _, _| Ok(()));
+
+    let mut batcher = create_batcher(mock_dependencies).await;
+
+    let task = CommitterTaskInput::ReadPathsAndCommitBlock(ReadPathsAndCommitBlockRequest {
+        commit: CommitBlockRequest {
+            height: INITIAL_HEIGHT,
+            state_diff: ThinStateDiff::default(),
+            state_diff_commitment: None,
+        },
+        accessed_keys: Default::default(),
+    });
+    batcher.commitment_manager.tasks_sender.send(task).await.unwrap();
+    wait_for_n_items(&mut batcher.commitment_manager.results_receiver, 1).await;
+
+    // The mock storage reader has no state commitment infos expectations, so both answers come
+    // from the cache.
+    assert_eq!(
+        batcher.get_state_commitment_infos(INITIAL_HEIGHT),
+        Ok(Some(test_state_commitment_infos()))
+    );
+    assert_eq!(batcher.has_state_commitment_infos(INITIAL_HEIGHT), Ok(true));
+}
+
 #[tokio::test]
 async fn get_block_hash_error() {
     let mut mock_dependencies = MockDependencies::default();

crates/apollo_batcher/src/commitment_manager/commitment_manager_impl.rs

@@ -14,6 +14,7 @@ use apollo_committer_types::committer_types::{
 };
 use apollo_committer_types::communication::{CommitterRequestLabelValue, SharedCommitterClient};
 use apollo_storage::accessed_keys::AccessedKeys as StorageAccessedKeys;
+use apollo_storage::state_commitment_infos::CompressedStateCommitmentInfos;
 use lru::LruCache;
 use starknet_api::block::{BlockHash, BlockNumber};
 use starknet_api::block_hash::block_hash_calculator::{
@@ -50,6 +51,7 @@ use crate::metrics::{
 // TODO(Amos): Add this to config.
 const TASK_SEND_RETRY_DELAY: Duration = Duration::from_millis(100);
 const N_RECENT_BLOCK_HASHES: NonZeroUsize = NonZeroUsize::new(30).unwrap();
+const N_RECENT_STATE_COMMITMENT_INFOS: NonZeroUsize = NonZeroUsize::new(20).unwrap();
 
 pub(crate) type CommitmentManagerResult<T> = Result<T, CommitmentManagerError>;
 pub(crate) type ApolloCommitmentManager = CommitmentManager<StateCommitter>;
@@ -65,6 +67,8 @@ pub(crate) struct CommitmentManager<S: StateCommitterTrait> {
     pub(crate) state_committer: S,
     pub(crate) task_timer: TaskTimer,
     pub(crate) recent_block_hashes_cache: LruCache<BlockNumber, BlockHash>,
+    pub(crate) recent_state_commitment_infos_cache:
+        LruCache<BlockNumber, CompressedStateCommitmentInfos>,
 }
 
 impl<S: StateCommitterTrait> CommitmentManager<S> {
@@ -236,6 +240,7 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
                     commitment_results.push(read_path_and_commit_task_result)
                 }
                 CommitterTaskOutput::Revert(revert_task_result) => {
+                    self.recent_state_commitment_infos_cache.pop(&revert_task_result.height);
                     return (commitment_results, revert_task_result);
                 }
             }
@@ -297,6 +302,11 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
                 self.recent_block_hashes_cache.put(height, block_hash);
             }
 
+            // Add commitment infos to cache.
+            if let Some(state_commitment_infos) = state_commitment_infos.clone() {
+                self.recent_state_commitment_infos_cache.put(height, state_commitment_infos);
+            }
+
             // Write the block hash and global root to storage.
             storage_writer.set_global_root_and_block_hash(
                 height,
@@ -373,6 +383,7 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
             state_committer,
             task_timer,
             recent_block_hashes_cache: LruCache::new(N_RECENT_BLOCK_HASHES),
+            recent_state_commitment_infos_cache: LruCache::new(N_RECENT_STATE_COMMITMENT_INFOS),
         }
     }
 

crates/apollo_batcher/src/commitment_manager/commitment_manager_test.rs

@@ -14,11 +14,6 @@ use apollo_committer_types::committer_types::{
 };
 use apollo_committer_types::communication::MockCommitterClient;
 use apollo_storage::accessed_keys::AccessedKeys;
-use apollo_storage::state_commitment_infos::{
-    CompressedPayload,
-    CompressedStateCommitmentInfos,
-    STATE_COMMITMENT_INFOS_VERSION,
-};
 use apollo_storage::StorageResult;
 use assert_matches::assert_matches;
 use mockall::predicate::eq;
@@ -37,6 +32,7 @@ use crate::commitment_manager::errors::CommitmentManagerError;
 use crate::test_utils::{
     get_number_of_items_in_channel_from_receiver,
     get_number_of_items_in_channel_from_sender,
+    test_state_commitment_infos,
     test_state_diff,
     wait_for_condition,
     wait_for_n_items,
@@ -73,10 +69,7 @@ fn mock_dependencies() -> MockDependencies {
         Box::pin(async {
             Ok(ReadPathsAndCommitBlockResponse {
                 global_root: GlobalRoot::default(),
-                state_commitment_infos: CompressedStateCommitmentInfos {
-                    version: STATE_COMMITMENT_INFOS_VERSION,
-                    payload: CompressedPayload(Vec::new()),
-                },
+                state_commitment_infos: test_state_commitment_infos(),
             })
         })
     });
@@ -592,9 +585,13 @@ async fn test_wait_for_revert(mut mock_dependencies: MockDependencies) {
         INITIAL_HEIGHT,
     )
     .await;
+    commitment_manager
+        .recent_state_commitment_infos_cache
+        .put(height, test_state_commitment_infos());
     let (commitment_results, revert_result) = commitment_manager.wait_for_revert_result().await;
     assert_eq!(commitment_results.len(), 2);
     assert_eq!(revert_result.height, height);
+    assert!(!commitment_manager.recent_state_commitment_infos_cache.contains(&height));
 }
 
 #[rstest]

crates/apollo_batcher/src/test_utils.rs

@@ -235,6 +235,13 @@ pub(crate) fn test_contract_nonces() -> HashMap<ContractAddress, Nonce> {
     HashMap::from_iter((0..3u8).map(|i| (contract_address!(i + 33), nonce!(i + 9))))
 }
 
+pub(crate) fn test_state_commitment_infos() -> CompressedStateCommitmentInfos {
+    CompressedStateCommitmentInfos {
+        version: STATE_COMMITMENT_INFOS_VERSION,
+        payload: CompressedPayload(Vec::new()),
+    }
+}
+
 pub(crate) fn test_state_diff() -> ThinStateDiff {
     ThinStateDiff {
         storage_diffs: indexmap! {
@@ -302,10 +309,7 @@ impl Default for MockClients {
             Box::pin(async {
                 Ok(ReadPathsAndCommitBlockResponse {
                     global_root: GlobalRoot::default(),
-                    state_commitment_infos: CompressedStateCommitmentInfos {
-                        version: STATE_COMMITMENT_INFOS_VERSION,
-                        payload: CompressedPayload(Vec::new()),
-                    },
+                    state_commitment_infos: test_state_commitment_infos(),
                 })
             })
         });