apollo_batcher: full revert commitment flow

Author: yoavGrs

crates/apollo_batcher/src/batcher.rs

@@ -30,6 +30,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};
@@ -39,7 +40,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,
@@ -89,7 +90,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,
@@ -708,7 +713,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 +778,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 +1112,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);
         STORAGE_HEIGHT.decrement(1);
         REVERTED_BLOCKS.increment(1);
@@ -1155,8 +1163,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;
 
@@ -1215,6 +1232,64 @@ 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_results().await;
+        self.write_commitment_results_to_storage(commitment_results)
+            .await
+            // Consider continuing the revert and not panicking here.
+            .expect("Failed to write commitment results to storage.");
+
+        self.validate_revert_task_result(revert_task_result, height).await;
+        self.commitment_manager.decrease_commitment_task_offset();
+        info!("Reverted commitment for height {height}.");
+    }
+
+    async fn validate_revert_task_result(
+        &self,
+        revert_task_output: RevertTaskOutput,
+        request_height: BlockNumber,
+    ) {
+        assert_eq!(
+            revert_task_output.height, request_height,
+            "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)
+                    // Consider continuing the revert and not panicking here.
+                    .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)
@@ -1336,6 +1411,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.
@@ -1362,6 +1439,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/commitment_manager/commitment_manager_impl.rs

@@ -211,6 +211,10 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
         self.commitment_task_offset =
             self.commitment_task_offset.next().expect("Block number overflowed.");
     }
+    pub(crate) fn decrease_commitment_task_offset(&mut self) {
+        self.commitment_task_offset =
+            self.commitment_task_offset.prev().expect("Can't revert before the genesis block.");
+    }
 
     async fn read_commitment_input_and_add_task<R: BatcherStorageReader>(
         &mut self,