diff --git a/crates/apollo_batcher/src/batcher.rs b/crates/apollo_batcher/src/batcher.rs index 4f1593ff253..bc7e67ce4f5 100644 --- a/crates/apollo_batcher/src/batcher.rs +++ b/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> { + 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 { + pub fn has_state_commitment_infos(&mut self, block_number: BlockNumber) -> BatcherResult { + 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 diff --git a/crates/apollo_batcher/src/batcher_test.rs b/crates/apollo_batcher/src/batcher_test.rs index 1ec789defdf..cbf3e026d1b 100644 --- a/crates/apollo_batcher/src/batcher_test.rs +++ b/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::(&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(); diff --git a/crates/apollo_batcher/src/commitment_manager/commitment_manager_impl.rs b/crates/apollo_batcher/src/commitment_manager/commitment_manager_impl.rs index cd390d2ff57..859e466d919 100644 --- a/crates/apollo_batcher/src/commitment_manager/commitment_manager_impl.rs +++ b/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 = Result; pub(crate) type ApolloCommitmentManager = CommitmentManager; @@ -65,6 +67,8 @@ pub(crate) struct CommitmentManager { pub(crate) state_committer: S, pub(crate) task_timer: TaskTimer, pub(crate) recent_block_hashes_cache: LruCache, + pub(crate) recent_state_commitment_infos_cache: + LruCache, } impl CommitmentManager { @@ -236,6 +240,7 @@ impl CommitmentManager { 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 CommitmentManager { 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 CommitmentManager { 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), } } diff --git a/crates/apollo_batcher/src/commitment_manager/commitment_manager_test.rs b/crates/apollo_batcher/src/commitment_manager/commitment_manager_test.rs index fbaf386489c..ca8f96d2883 100644 --- a/crates/apollo_batcher/src/commitment_manager/commitment_manager_test.rs +++ b/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] diff --git a/crates/apollo_batcher/src/test_utils.rs b/crates/apollo_batcher/src/test_utils.rs index 804556eb38d..53f4fb3c060 100644 --- a/crates/apollo_batcher/src/test_utils.rs +++ b/crates/apollo_batcher/src/test_utils.rs @@ -235,6 +235,13 @@ pub(crate) fn test_contract_nonces() -> HashMap { 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(), }) }) });