apollo_batcher: add revert task

Author: yoavGrs

crates/apollo_batcher/src/commitment_manager/commitment_manager_impl.rs

@@ -3,7 +3,11 @@
 use std::sync::Arc;
 
 use apollo_batcher_config::config::{BatcherConfig, CommitmentManagerConfig};
-use apollo_committer_types::committer_types::{CommitBlockRequest, CommitBlockResponse};
+use apollo_committer_types::committer_types::{
+    CommitBlockRequest,
+    CommitBlockResponse,
+    RevertBlockRequest,
+};
 use apollo_committer_types::communication::SharedCommitterClient;
 use starknet_api::block::BlockNumber;
 use starknet_api::block_hash::block_hash_calculator::{
@@ -73,57 +77,63 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
     }
 
     /// Adds a commitment task to the state committer. If the task height does not match the
-    /// task offset, an error is returned. If the tasks channel is full, the behavior depends on
-    /// the config: if `wait_for_tasks_channel` is true, it will wait until there is space in the
-    /// channel; otherwise, it will panic. Any other error when sending the task will also cause a
-    /// panic.
+    /// task offset, an error is returned.
     pub(crate) async fn add_commitment_task(
         &mut self,
         height: BlockNumber,
         state_diff: ThinStateDiff,
         state_diff_commitment: Option<StateDiffCommitment>,
     ) -> CommitmentManagerResult<()> {
         if height != self.commitment_task_offset {
-            return Err(CommitmentManagerError::WrongTaskHeight {
+            return Err(CommitmentManagerError::WrongCommitmentTaskHeight {
                 expected: self.commitment_task_offset,
                 actual: height,
                 state_diff_commitment,
             });
         }
-        let commitment_task_input = CommitterTaskInput::Commit(CommitBlockRequest {
-            height,
-            state_diff,
-            state_diff_commitment,
-        });
-        let error_message = format!(
-            "Failed to send commitment task to state committer. Block: {height}, state diff \
-             commitment: {state_diff_commitment:?}",
-        );
+        let commitment_task_input =
+            CommitterTaskInput::Commit(CommitBlockRequest {
+                height,
+                state_diff,
+                state_diff_commitment,
+            });
+        let task_details =
+            format!("Commit block {height}, state diff commitment: {state_diff_commitment:?}",);
+        self.add_task(commitment_task_input, &task_details).await?;
+        self.successfully_added_commitment_task(height, state_diff_commitment);
+        Ok(())
+    }
+
+    /// If the tasks channel is full, the behavior depends on the config: if
+    /// `wait_for_tasks_channel` is true, it will wait until there is space in the channel;
+    /// otherwise, it will panic. Any other error when sending the task will also cause a panic.
+    async fn add_task(
+        &self,
+        task_input: CommitterTaskInput,
+        task_details: &str,
+    ) -> CommitmentManagerResult<()> {
+        let error_message = format!("Failed to send task to state committer: {task_details}.",);
 
         if self.config.wait_for_tasks_channel {
-            info!(
-                "Waiting to send commitment task for block {height} and state diff \
-                 {state_diff_commitment:?} to state committer."
-            );
-            match self.tasks_sender.send(commitment_task_input).await {
-                Ok(_) => self.successfully_added_commitment_task(height, state_diff_commitment),
-                Err(err) => panic!("{error_message}. error: {err}"),
-            }
+            info!("Waiting to send task for {task_details} to state committer.");
+            self.tasks_sender
+                .send(task_input)
+                .await
+                .unwrap_or_else(|err| panic!("{error_message}. error: {err}"));
         } else {
-            match self.tasks_sender.try_send(commitment_task_input) {
-                Ok(_) => self.successfully_added_commitment_task(height, state_diff_commitment),
+            match self.tasks_sender.try_send(task_input) {
+                Ok(_) => (),
                 Err(TrySendError::Full(_)) => {
                     let channel_size = self.tasks_sender.max_capacity();
                     panic!(
-                        "Failed to send commitment task to state committer because the channel is \
-                         full. Block: {height}, state diff commitment: {state_diff_commitment:?}, \
-                         channel size: {channel_size}. Consider increasing the channel size or \
-                         enabling waiting in the config.",
+                        "{error_message}. The channel is full. channel size: {channel_size}. \
+                         Consider increasing the channel size or enabling waiting in the config.",
                     );
                 }
                 Err(err) => panic!("{error_message}. error: {err}"),
             }
         }
+        Ok(())
     }
 
     /// Fetches all ready commitment results from the state committer. Panics if any task is a
@@ -168,13 +178,12 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
         &mut self,
         height: BlockNumber,
         state_diff_commitment: Option<StateDiffCommitment>,
-    ) -> CommitmentManagerResult<()> {
+    ) {
         info!(
             "Sent commitment task for block {height} and state diff {state_diff_commitment:?} to \
              state committer."
         );
         self.increase_commitment_task_offset();
-        Ok(())
     }
 
     /// Initializes the CommitmentManager. This includes starting the state committer task.
@@ -256,8 +265,28 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
 
     // Associated functions.
 
-    pub(crate) async fn revert_block(height: BlockNumber, reversed_state_diff: ThinStateDiff) {
-        unimplemented!()
+    pub(crate) async fn add_revert_task(
+        &mut self,
+        height: BlockNumber,
+        reversed_state_diff: ThinStateDiff,
+    ) -> CommitmentManagerResult<()> {
+        let expected_height =
+            self.commitment_task_offset.prev().expect("Can't revert before the genesis block.");
+        if height != expected_height {
+            return Err(CommitmentManagerError::WrongRevertTaskHeight {
+                expected: expected_height,
+                actual: height,
+            });
+        }
+        let revert_task_input =
+            CommitterTaskInput::Revert(RevertBlockRequest {
+                height,
+                reversed_state_diff,
+            });
+        let task_details = format!("Revert block {height}",);
+        self.add_task(revert_task_input, &task_details).await?;
+        info!("Sent revert task for block {height}.");
+        Ok(())
     }
 
     /// Returns the final commitment output for a given commitment task output.

crates/apollo_batcher/src/commitment_manager/commitment_manager_test.rs

@@ -167,7 +167,7 @@ async fn test_add_commitment_task(mut mock_dependencies: MockDependencies) {
         .await;
     assert_matches!(
         result,
-        Err(CommitmentManagerError::WrongTaskHeight { expected, actual, .. })
+        Err(CommitmentManagerError::WrongCommitmentTaskHeight { expected, actual, .. })
         if expected == INITIAL_HEIGHT && actual == incorrect_height
     );
 
@@ -212,8 +212,7 @@ async fn test_add_commitment_task(mut mock_dependencies: MockDependencies) {
 
 #[rstest]
 #[tokio::test]
-#[should_panic(expected = "Failed to send commitment task to state committer because the channel \
-                           is full. Block: 4")]
+#[should_panic(expected = "The channel is full. channel size: 1.")]
 async fn test_add_commitment_task_full(mut mock_dependencies: MockDependencies) {
     add_initial_heights(&mut mock_dependencies);
     let state_diff = test_state_diff();

crates/apollo_batcher/src/commitment_manager/errors.rs

@@ -17,9 +17,11 @@ pub enum CommitmentManagerError {
         "Wrong commitment task height. Expected: {expected}, Actual: {actual}. State diff \
          commitment: {state_diff_commitment:?}"
     )]
-    WrongTaskHeight {
+    WrongCommitmentTaskHeight {
         expected: BlockNumber,
         actual: BlockNumber,
         state_diff_commitment: Option<StateDiffCommitment>,
     },
+    #[error("Wrong revert task height. Expected: {expected}, Actual: {actual}.")]
+    WrongRevertTaskHeight { expected: BlockNumber, actual: BlockNumber },
 }