feat(apollo_reverts): add revert blocks functionality

Cargo.toml

@@ -4,6 +4,7 @@
 resolver = "2"
 
 members = [
+  "crates/apollo_reverts",
   "crates/blockifier",
   "crates/blockifier_reexecution",
   "crates/blockifier_test_utils",
@@ -84,6 +85,7 @@ alloy-sol-types = "0.8.3"
 alloy-transport = "0.3.5"
 alloy-transport-http = "0.3.5"
 anyhow = "1.0.44"
+apollo_reverts = { path = "crates/apollo_reverts", version = "0.0.0" }
 ark-ec = "0.4.2"
 ark-ff = "0.4.0-alpha.7"
 ark-secp256k1 = "0.4.0"

commitlint.config.js

@@ -20,6 +20,7 @@ const Configuration = {
     rules: {
         'scope-empty': [2, 'never'],
         'scope-enum': [2, 'always', [
+            "apollo_reverts",
             'blockifier',
             'blockifier_reexecution',
             'blockifier_test_utils',

config/sequencer/default_config.json

@@ -889,15 +889,15 @@
     "privacy": "Public",
     "value": "consensus_proposals"
   },
-  "consensus_manager_config.revert_up_to_and_including": {
-    "description": "The batcher will revert blocks up to this block number (including). Use this configurations carefully to prevent significant revert operations and data loss.",
-    "privacy": "Private",
-    "value": 18446744073709551615
+  "consensus_manager_config.revert_config.revert_up_to_and_including": {
+    "description": "The component will revert blocks up to this block number (including).",
+    "pointer_target": "revert_config.revert_up_to_and_including",
+    "privacy": "Public"
   },
-  "consensus_manager_config.revert_up_to_and_including.#is_none": {
-    "description": "Flag for an optional field.",
-    "privacy": "TemporaryValue",
-    "value": true
+  "consensus_manager_config.revert_config.should_revert": {
+    "description": "If set true, the component would revert blocks and do nothing else.",
+    "pointer_target": "revert_config.should_revert",
+    "privacy": "Public"
   },
   "consensus_manager_config.votes_topic": {
     "description": "The topic for consensus votes.",
@@ -1149,6 +1149,16 @@
     "param_type": "String",
     "privacy": "TemporaryValue"
   },
+  "revert_config.revert_up_to_and_including": {
+    "description": "The component will revert blocks up to this block number (including).",
+    "privacy": "TemporaryValue",
+    "value": 18446744073709551615
+  },
+  "revert_config.should_revert": {
+    "description": "If set true, the component would revert blocks and do nothing else.",
+    "privacy": "TemporaryValue",
+    "value": false
+  },
   "state_sync_config.central_sync_client_config.#is_none": {
     "description": "Flag for an optional field.",
     "privacy": "TemporaryValue",

crates/apollo_reverts/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "apollo_reverts"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+futures.workspace = true
+papyrus_config.workspace = true
+papyrus_storage.workspace = true
+serde.workspace = true
+starknet_api.workspace = true
+tracing.workspace = true
+validator.workspace = true

crates/apollo_reverts/src/lib.rs

@@ -0,0 +1,116 @@
+use std::collections::BTreeMap;
+use std::future::Future;
+
+use futures::future::pending;
+use futures::never::Never;
+use papyrus_config::dumping::{ser_param, SerializeConfig};
+use papyrus_config::{ParamPath, ParamPrivacyInput, SerializedParam};
+use papyrus_storage::base_layer::BaseLayerStorageWriter;
+use papyrus_storage::body::BodyStorageWriter;
+use papyrus_storage::class_manager::ClassManagerStorageWriter;
+use papyrus_storage::header::HeaderStorageWriter;
+use papyrus_storage::state::StateStorageWriter;
+use papyrus_storage::StorageWriter;
+use serde::{Deserialize, Serialize};
+use starknet_api::block::BlockNumber;
+use tracing::info;
+use validator::Validate;
+
+#[derive(Clone, Debug, Serialize, Deserialize, Validate, PartialEq)]
+pub struct RevertConfig {
+    pub revert_up_to_and_including: BlockNumber,
+    pub should_revert: bool,
+}
+
+impl Default for RevertConfig {
+    fn default() -> Self {
+        Self {
+            // Use u64::MAX as a placeholder to prevent setting this value to
+            // a low block number by mistake, which will cause significant revert operations.
+            revert_up_to_and_including: BlockNumber(u64::MAX),
+            should_revert: false,
+        }
+    }
+}
+
+impl SerializeConfig for RevertConfig {
+    fn dump(&self) -> BTreeMap<ParamPath, SerializedParam> {
+        BTreeMap::from_iter([
+            ser_param(
+                "revert_up_to_and_including",
+                &self.revert_up_to_and_including,
+                "The component will revert blocks up to this block number (including).",
+                // Use this configurations carefully to prevent significant revert operations and
+                // data loss
+                ParamPrivacyInput::Public,
+            ),
+            ser_param(
+                "should_revert",
+                &self.should_revert,
+                "If set true, the component would revert blocks and do nothing else.",
+                ParamPrivacyInput::Public,
+            ),
+        ])
+    }
+}
+
+pub async fn revert_blocks_and_eternal_pending<Fut>(
+    mut storage_height_marker: BlockNumber,
+    revert_up_to_and_including: BlockNumber,
+    mut revert_block_fn: impl FnMut(BlockNumber) -> Fut,
+    component_name: &str,
+) -> Never
+where
+    Fut: Future<Output = ()>,
+{
+    if storage_height_marker <= revert_up_to_and_including {
+        panic!(
+            "{component_name}'s storage height marker {storage_height_marker} is not larger than \
+             the target block {revert_up_to_and_including}. No reverts are needed."
+        );
+    }
+
+    info!(
+        "Reverting {component_name} from block {storage_height_marker} to block \
+         {revert_up_to_and_including}"
+    );
+
+    while storage_height_marker > revert_up_to_and_including {
+        storage_height_marker = storage_height_marker.prev().expect(
+            "A block number that's greater than another block number should return Some on prev",
+        );
+        info!("Reverting {component_name} block number {storage_height_marker}.");
+        revert_block_fn(storage_height_marker).await;
+        info!("Successfully reverted {component_name} block number {storage_height_marker}.");
+    }
+
+    info!(
+        "Done reverting {component_name} up to block {revert_up_to_and_including} including. \
+         Starting eternal pending."
+    );
+    pending().await
+}
+
+/// Reverts everything related to the block, will succeed even if there is partial information for
+/// the block.
+// This function will panic if the storage reader fails to revert.
+pub fn revert_block(storage_writer: &mut StorageWriter, current_block_number: BlockNumber) {
+    storage_writer
+        .begin_rw_txn()
+        .unwrap()
+        .revert_header(current_block_number)
+        .unwrap()
+        .0
+        .revert_body(current_block_number)
+        .unwrap()
+        .0
+        .revert_state_diff(current_block_number)
+        .unwrap()
+        .0
+        .try_revert_class_manager_marker(current_block_number)
+        .unwrap()
+        .try_revert_base_layer_marker(current_block_number)
+        .unwrap()
+        .commit()
+        .unwrap();
+}

crates/starknet_batcher/Cargo.toml

@@ -9,6 +9,7 @@ repository.workspace = true
 workspace = true
 
 [dependencies]
+apollo_reverts.workspace = true
 async-trait.workspace = true
 blockifier.workspace = true
 chrono.workspace = true

crates/starknet_batcher/src/batcher.rs

@@ -1,6 +1,7 @@
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
+use apollo_reverts::revert_block;
 use async_trait::async_trait;
 use blockifier::state::contract_class_manager::ContractClassManager;
 #[cfg(test)]
@@ -650,6 +651,7 @@ impl Batcher {
     }
 
     #[instrument(skip(self), err)]
+    // This function will panic if there is a storage failure to revert the block.
     pub async fn revert_block(&mut self, input: RevertBlockInput) -> BatcherResult<()> {
         info!("Reverting block at height {}.", input.height);
         let height = self.get_height_from_storage()?.prev().ok_or(
@@ -671,10 +673,7 @@ impl Batcher {
             self.abort_active_height().await;
         }
 
-        self.storage_writer.revert_block(height).map_err(|err| {
-            error!("Failed to revert block at height {}: {}", height, err);
-            BatcherError::InternalError
-        })?;
+        self.storage_writer.revert_block(height);
         STORAGE_HEIGHT.decrement(1);
         Ok(())
     }
@@ -733,7 +732,7 @@ pub trait BatcherStorageWriterTrait: Send + Sync {
         state_diff: ThinStateDiff,
     ) -> papyrus_storage::StorageResult<()>;
 
-    fn revert_block(&mut self, height: BlockNumber) -> papyrus_storage::StorageResult<()>;
+    fn revert_block(&mut self, height: BlockNumber);
 }
 
 impl BatcherStorageWriterTrait for papyrus_storage::StorageWriter {
@@ -746,10 +745,9 @@ impl BatcherStorageWriterTrait for papyrus_storage::StorageWriter {
         self.begin_rw_txn()?.append_state_diff(height, state_diff)?.commit()
     }
 
-    fn revert_block(&mut self, height: BlockNumber) -> papyrus_storage::StorageResult<()> {
-        let (txn, reverted_state_diff) = self.begin_rw_txn()?.revert_state_diff(height)?;
-        trace!("Reverted state diff: {:#?}", reverted_state_diff);
-        txn.commit()
+    // This function will panic if there is a storage failure to revert the block.
+    fn revert_block(&mut self, height: BlockNumber) {
+        revert_block(self, height);
     }
 }
 

crates/starknet_batcher/src/batcher_test.rs

@@ -815,7 +815,8 @@ async fn revert_block() {
         .expect_revert_block()
         .times(1)
         .with(eq(LATEST_BLOCK_IN_STORAGE))
-        .returning(|_| Ok(()));
+        .returning(|_| ());
+
     let mut batcher = create_batcher(mock_dependencies).await;
 
     let metrics = recorder.handle().render();

crates/starknet_consensus_manager/Cargo.toml

@@ -9,6 +9,7 @@ repository.workspace = true
 workspace = true
 
 [dependencies]
+apollo_reverts.workspace = true
 async-trait.workspace = true
 papyrus_config.workspace = true
 papyrus_network.workspace = true

crates/starknet_consensus_manager/src/config.rs

@@ -1,11 +1,7 @@
 use std::collections::BTreeMap;
 
-use papyrus_config::dumping::{
-    append_sub_config_name,
-    ser_optional_param,
-    ser_param,
-    SerializeConfig,
-};
+use apollo_reverts::RevertConfig;
+use papyrus_config::dumping::{append_sub_config_name, ser_param, SerializeConfig};
 use papyrus_config::{ParamPath, ParamPrivacyInput, SerializedParam};
 use papyrus_network::NetworkConfig;
 use serde::{Deserialize, Serialize};
@@ -24,7 +20,7 @@ pub struct ConsensusManagerConfig {
     #[validate]
     pub network_config: NetworkConfig,
     pub cende_config: CendeConfig,
-    pub revert_up_to_and_including: Option<BlockNumber>,
+    pub revert_config: RevertConfig,
     pub votes_topic: String,
     pub proposals_topic: String,
     pub broadcast_buffer_size: usize,
@@ -59,20 +55,11 @@ impl SerializeConfig for ConsensusManagerConfig {
                 ParamPrivacyInput::Public,
             ),
         ]);
-        config.extend(ser_optional_param(
-            &self.revert_up_to_and_including,
-            // Use u64::MAX as a placeholder to prevent setting this value to
-            // a low block number by mistake, which will cause significant revert operations.
-            BlockNumber(u64::MAX),
-            "revert_up_to_and_including",
-            "The batcher will revert blocks up to this block number (including). Use this \
-             configurations carefully to prevent significant revert operations and data loss.",
-            ParamPrivacyInput::Private,
-        ));
         config.extend(append_sub_config_name(self.consensus_config.dump(), "consensus_config"));
         config.extend(append_sub_config_name(self.context_config.dump(), "context_config"));
         config.extend(append_sub_config_name(self.cende_config.dump(), "cende_config"));
         config.extend(append_sub_config_name(self.network_config.dump(), "network_config"));
+        config.extend(append_sub_config_name(self.revert_config.dump(), "revert_config"));
         config
     }
 }
@@ -84,7 +71,7 @@ impl Default for ConsensusManagerConfig {
             context_config: ContextConfig::default(),
             cende_config: CendeConfig::default(),
             network_config: NetworkConfig::default(),
-            revert_up_to_and_including: None,
+            revert_config: RevertConfig::default(),
             votes_topic: "consensus_votes".to_string(),
             proposals_topic: "consensus_proposals".to_string(),
             broadcast_buffer_size: 10000,