apollo_batcher: full revert commitment flow

Author: yoavGrs

Cargo.lock

@@ -23,6 +23,7 @@ apollo_infra_utils.workspace = true
 apollo_l1_provider_types.workspace = true
 apollo_mempool_types.workspace = true
 apollo_metrics.workspace = true
+apollo_proc_macros.workspace = true
 apollo_reverts.workspace = true
 apollo_starknet_client.workspace = true
 apollo_state_reader.workspace = true

crates/apollo_batcher/Cargo.toml

@@ -26,6 +26,7 @@ use apollo_batcher_types::batcher_types::{
 use apollo_batcher_types::errors::BatcherError;
 use apollo_class_manager_types::transaction_converter::TransactionConverter;
 use apollo_class_manager_types::SharedClassManagerClient;
+use apollo_committer_types::committer_types::RevertBlockResponse;
 use apollo_committer_types::communication::SharedCommitterClient;
 use apollo_infra::component_definitions::{default_component_start_fn, ComponentStarter};
 use apollo_l1_provider_types::errors::{L1ProviderClientError, L1ProviderError};
@@ -35,7 +36,7 @@ use apollo_mempool_types::mempool_types::CommitBlockArgs;
 use apollo_reverts::revert_block;
 use apollo_state_sync_types::state_sync_types::SyncBlock;
 use apollo_storage::block_hash::{BlockHashStorageReader, BlockHashStorageWriter};
-use apollo_storage::global_root::GlobalRootStorageWriter;
+use apollo_storage::global_root::{GlobalRootStorageReader, GlobalRootStorageWriter};
 use apollo_storage::global_root_marker::{
     GlobalRootMarkerStorageReader,
     GlobalRootMarkerStorageWriter,
@@ -86,7 +87,11 @@ use crate::commitment_manager::commitment_manager_impl::{
     ApolloCommitmentManager,
     CommitmentManager,
 };
-use crate::commitment_manager::types::FinalBlockCommitment;
+use crate::commitment_manager::types::{
+    CommitmentTaskOutput,
+    FinalBlockCommitment,
+    RevertTaskOutput,
+};
 use crate::metrics::{
     register_metrics,
     ProposalMetricsHandle,
@@ -478,6 +483,7 @@ impl Batcher {
     }
 
     /// Clear all the proposals from the previous height.
+    #[cfg_attr(any(test, feature = "testing"), apollo_proc_macros::make_visibility(pub))]
     async fn abort_active_height(&mut self) {
         self.abort_active_proposal().await;
         self.executed_proposals.lock().await.clear();
@@ -708,7 +714,7 @@ impl Batcher {
             .expect("The commitment offset unexpectedly doesn't match the given block height.");
 
         // Write ready commitments to storage.
-        self.write_commitment_results_to_storage().await?;
+        self.get_and_write_commitment_results_to_storage().await?;
 
         LAST_SYNCED_BLOCK_HEIGHT.set_lossy(block_number.0);
         SYNCED_TRANSACTIONS.increment(
@@ -773,7 +779,7 @@ impl Batcher {
             .expect("The commitment offset unexpectedly doesn't match the given block height.");
 
         // Write ready commitments to storage.
-        self.write_commitment_results_to_storage().await?;
+        self.get_and_write_commitment_results_to_storage().await?;
 
         let execution_infos = block_execution_artifacts
             .execution_data
@@ -1107,6 +1113,9 @@ impl Batcher {
             self.abort_active_height().await;
         }
 
+        // Wait for the revert commitment to be completed before reverting the storage.
+        self.revert_commitment(height).await;
+
         self.storage_writer.revert_block(height);
         BUILDING_HEIGHT.decrement(1);
         GLOBAL_ROOT_HEIGHT.decrement(1);
@@ -1156,8 +1165,17 @@ impl Batcher {
     }
 
     /// Writes the ready commitment results to storage.
-    async fn write_commitment_results_to_storage(&mut self) -> BatcherResult<()> {
+    async fn get_and_write_commitment_results_to_storage(&mut self) -> BatcherResult<()> {
         let commitment_results = self.commitment_manager.get_commitment_results().await;
+        self.write_commitment_results_to_storage(commitment_results).await?;
+        Ok(())
+    }
+
+    /// Writes the given commitment results to storage.
+    async fn write_commitment_results_to_storage(
+        &mut self,
+        commitment_results: Vec<CommitmentTaskOutput>,
+    ) -> BatcherResult<()> {
         for commitment_task_output in commitment_results.into_iter() {
             let height = commitment_task_output.height;
             info!("Writing commitment results to storage for height {}.", height);
@@ -1217,6 +1235,62 @@ impl Batcher {
         Ok(())
     }
 
+    /// Reverts the commitment for the given height.
+    /// Adds a revert task to the commitment manager channel and waits for the result.
+    /// Writes commitment results to storage and handles the revert task result.
+    async fn revert_commitment(&mut self, height: BlockNumber) {
+        let reversed_state_diff = self
+            .storage_reader
+            .reversed_state_diff(height)
+            .expect("Failed to get reversed state diff from storage.");
+        self.commitment_manager
+            .add_revert_task(height, reversed_state_diff)
+            .await
+            .expect("Failed to add revert task to commitment manager.");
+        let (commitment_results, revert_task_result) =
+            self.commitment_manager.wait_for_revert_result().await;
+        self.write_commitment_results_to_storage(commitment_results)
+            .await
+            .expect("Failed to write commitment results to storage.");
+
+        info!("Revert task result: {revert_task_result:?}");
+        self.validate_revert_task_result(revert_task_result, height).await;
+        info!("Reverted commitment for height {height}.");
+    }
+
+    async fn validate_revert_task_result(
+        &self,
+        revert_task_output: RevertTaskOutput,
+        request_height_to_revert: BlockNumber,
+    ) {
+        assert_eq!(
+            revert_task_output.height, request_height_to_revert,
+            "The task output height does not match the request height."
+        );
+
+        match revert_task_output.response {
+            RevertBlockResponse::RevertedTo(global_root)
+            | RevertBlockResponse::AlreadyReverted(global_root) => {
+                // Verify the global root matches the stored global root.
+                let new_latest_height = revert_task_output
+                    .height
+                    .prev()
+                    .expect("Can't revert before the genesis block.");
+                let stored_global_root = self
+                    .storage_reader
+                    .global_root(new_latest_height)
+                    .expect("Failed to get global root from storage.")
+                    .expect("Global root is not set for height {new_latest_height}.");
+                assert_eq!(
+                    global_root, stored_global_root,
+                    "The given global root does not match the stored global root for height \
+                     {new_latest_height}."
+                );
+            }
+            RevertBlockResponse::Uncommitted => {}
+        }
+    }
+
     pub fn get_block_hash(&self, block_number: BlockNumber) -> BatcherResult<BlockHash> {
         self.storage_reader
             .get_block_hash(block_number)
@@ -1341,6 +1415,8 @@ pub trait BatcherStorageReader: Send + Sync {
     /// Returns the first height the committer has finished calculating commitments for.
     fn global_root_height(&self) -> StorageResult<BlockNumber>;
 
+    fn global_root(&self, height: BlockNumber) -> StorageResult<Option<GlobalRoot>>;
+
     fn get_state_diff(&self, height: BlockNumber) -> StorageResult<Option<ThinStateDiff>>;
 
     /// Returns the state diff that undoes the state diff at the given height.
@@ -1367,6 +1443,10 @@ impl BatcherStorageReader for StorageReader {
         self.begin_ro_txn()?.get_global_root_marker()
     }
 
+    fn global_root(&self, height: BlockNumber) -> StorageResult<Option<GlobalRoot>> {
+        self.begin_ro_txn()?.get_global_root(&height)
+    }
+
     fn get_state_diff(&self, height: BlockNumber) -> StorageResult<Option<ThinStateDiff>> {
         self.begin_ro_txn()?.get_state_diff(height)
     }

crates/apollo_batcher/src/batcher.rs

@@ -48,7 +48,7 @@ use starknet_api::block::{
 use starknet_api::block_hash::block_hash_calculator::PartialBlockHashComponents;
 use starknet_api::block_hash::state_diff_hash::calculate_state_diff_hash;
 use starknet_api::consensus_transaction::InternalConsensusTransaction;
-use starknet_api::core::{ClassHash, CompiledClassHash, Nonce};
+use starknet_api::core::{ClassHash, CompiledClassHash, GlobalRoot, Nonce};
 use starknet_api::state::ThinStateDiff;
 use starknet_api::test_utils::CHAIN_ID_FOR_TESTS;
 use starknet_api::transaction::TransactionHash;
@@ -287,6 +287,15 @@ fn mock_create_builder_for_validate_block(
     );
 }
 
+fn mock_storage_reader_for_revert() -> MockBatcherStorageReader {
+    let mut storage_reader = MockBatcherStorageReader::new();
+    storage_reader.expect_reversed_state_diff().returning(|_| Ok(test_state_diff()));
+    storage_reader.expect_global_root_height().returning(|| Ok(INITIAL_HEIGHT));
+    storage_reader.expect_global_root().returning(|_| Ok(Some(GlobalRoot::default())));
+    storage_reader.expect_state_diff_height().returning(|| Ok(INITIAL_HEIGHT));
+    storage_reader
+}
+
 fn mock_create_builder_for_propose_block(
     block_builder_factory: &mut MockBlockBuilderFactoryTrait,
     output_txs: Vec<InternalConsensusTransaction>,
@@ -777,14 +786,10 @@ async fn multiple_proposals_with_l1_every_n_proposals() {
         );
     }
 
-    let mut storage_writer = MockBatcherStorageWriter::new();
-    storage_writer.expect_revert_block().times(N_PROPOSALS).returning(|_| ());
-
     let mut l1_provider_client = MockL1ProviderClient::new();
     l1_provider_client.expect_start_block().times(N_PROPOSALS).returning(|_, _| Ok(()));
 
     let mock_dependencies = MockDependencies {
-        storage_writer,
         clients: MockClients { block_builder_factory, l1_provider_client, ..Default::default() },
         ..Default::default()
     };
@@ -810,10 +815,7 @@ async fn multiple_proposals_with_l1_every_n_proposals() {
         }
 
         batcher.await_active_proposal(DUMMY_FINAL_N_EXECUTED_TXS).await.unwrap();
-        batcher
-            .revert_block(RevertBlockInput { height: INITIAL_HEIGHT.prev().unwrap() })
-            .await
-            .unwrap();
+        batcher.abort_active_height().await;
     }
 }
 
@@ -1297,14 +1299,17 @@ async fn revert_block() {
         .with(eq(LATEST_BLOCK_IN_STORAGE))
         .returning(|_| ());
 
-    let mock_dependencies = MockDependencies { storage_writer, ..Default::default() };
+    let storage_reader = mock_storage_reader_for_revert();
+    let mock_dependencies =
+        MockDependencies { storage_reader, storage_writer, ..Default::default() };
 
     let mut batcher = create_batcher(mock_dependencies).await;
 
     let metrics = recorder.handle().render();
     assert_eq!(BUILDING_HEIGHT.parse_numeric_metric::<u64>(&metrics), Some(INITIAL_HEIGHT.0));
 
     let revert_input = RevertBlockInput { height: LATEST_BLOCK_IN_STORAGE };
+
     batcher.revert_block(revert_input).await.unwrap();
 
     let metrics = recorder.handle().render();

crates/apollo_batcher/src/batcher_test.rs

@@ -276,8 +276,16 @@ impl Default for MockClients {
             (mock_writer, non_working_candidate_tx_sender, non_working_pre_confirmed_tx_sender)
         });
 
+        let mut committer_client = MockCommitterClient::new();
+        committer_client
+            .expect_commit_block()
+            .returning(|_| Box::pin(async { Ok(CommitBlockResponse::default()) }));
+        committer_client.expect_revert_block().returning(|_| {
+            Box::pin(async { Ok(RevertBlockResponse::RevertedTo(GlobalRoot::default())) })
+        });
+
         Self {
-            committer_client: MockCommitterClient::new(),
+            committer_client,
             l1_provider_client: MockL1ProviderClient::new(),
             mempool_client,
             block_builder_factory,