apollo_batcher: define custom input task type

Author: yoavGrs

crates/apollo_batcher/src/commitment_manager/commitment_manager_impl.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use apollo_batcher_config::config::{BatcherConfig, CommitmentManagerConfig};
 use apollo_committer_types::committer_types::{CommitBlockRequest, CommitBlockResponse};
-use apollo_committer_types::communication::{CommitterRequest, SharedCommitterClient};
+use apollo_committer_types::communication::SharedCommitterClient;
 use starknet_api::block::{BlockHash, BlockNumber};
 use starknet_api::block_hash::block_hash_calculator::{
     calculate_block_hash,
@@ -90,12 +90,11 @@ impl<S: StateCommitterTrait> CommitmentManager<S> {
                 state_diff_commitment,
             });
         }
-        let commitment_task_input =
-            CommitterTaskInput(CommitterRequest::CommitBlock(CommitBlockRequest {
-                height,
-                state_diff,
-                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:?}",

crates/apollo_batcher/src/commitment_manager/state_committer.rs

@@ -1,7 +1,7 @@
 #![allow(dead_code, unused_variables, unused_mut)]
 
 use apollo_committer_types::committer_types::{CommitBlockRequest, RevertBlockRequest};
-use apollo_committer_types::communication::{CommitterRequest, SharedCommitterClient};
+use apollo_committer_types::communication::SharedCommitterClient;
 use apollo_committer_types::errors::{CommitterClientError, CommitterClientResult};
 use tokio::sync::mpsc::error::TrySendError;
 use tokio::sync::mpsc::{Receiver, Sender};
@@ -55,7 +55,7 @@ impl StateCommitter {
         committer_client: SharedCommitterClient,
     ) {
         // TODO(Yoav): Test this function.
-        while let Some(CommitterTaskInput(request)) = tasks_receiver.recv().await {
+        while let Some(request) = tasks_receiver.recv().await {
             let output = perform_task(request, &committer_client).await;
             let height = output.height();
             match results_sender.try_send(output.clone()) {
@@ -77,15 +77,15 @@ impl StateCommitter {
 /// Performs a commitment task by calling the committer.
 /// Retries at recoverable errors.
 async fn perform_task(
-    request: CommitterRequest,
+    request: CommitterTaskInput,
     committer_client: &SharedCommitterClient,
 ) -> CommitterTaskOutput {
     loop {
         let result = match &request {
-            CommitterRequest::CommitBlock(commit_block_request) => {
+            CommitterTaskInput::Commit(commit_block_request) => {
                 perform_commit_block_task(commit_block_request.clone(), committer_client).await
             }
-            CommitterRequest::RevertBlock(revert_block_request) => {
+            CommitterTaskInput::Revert(revert_block_request) => {
                 perform_revert_block_task(revert_block_request.clone(), committer_client).await
             }
         };

crates/apollo_batcher/src/commitment_manager/types.rs

@@ -1,12 +1,43 @@
 #![allow(dead_code)]
 
-use apollo_committer_types::committer_types::{CommitBlockResponse, RevertBlockResponse};
-use apollo_committer_types::communication::CommitterRequest;
+use std::fmt::Display;
+
+use apollo_committer_types::committer_types::{
+    CommitBlockRequest,
+    CommitBlockResponse,
+    RevertBlockRequest,
+    RevertBlockResponse,
+};
 use starknet_api::block::{BlockHash, BlockNumber};
 use starknet_api::core::GlobalRoot;
 
 /// Input for commitment tasks.
-pub(crate) struct CommitterTaskInput(pub(crate) CommitterRequest);
+pub(crate) enum CommitterTaskInput {
+    Commit(CommitBlockRequest),
+    Revert(RevertBlockRequest),
+}
+
+impl CommitterTaskInput {
+    pub(crate) fn height(&self) -> BlockNumber {
+        match self {
+            Self::Commit(request) => request.height,
+            Self::Revert(request) => request.height,
+        }
+    }
+}
+
+impl Display for CommitterTaskInput {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Commit(request) => write!(
+                f,
+                "Commit(height={}, state_diff_commitment={:?})",
+                request.height, request.state_diff_commitment
+            ),
+            Self::Revert(request) => write!(f, "Revert(height={})", request.height),
+        }
+    }
+}
 
 #[derive(Clone, Debug)]
 pub(crate) struct CommitmentTaskOutput {

crates/apollo_batcher/src/test_utils.rs

@@ -347,7 +347,7 @@ impl MockStateCommitter {
             let task = tasks_receiver.try_recv().unwrap();
             let result = CommitterTaskOutput::Commit(CommitmentTaskOutput {
                 response: CommitBlockResponse { state_root: GlobalRoot::default() },
-                height: task.0.height(),
+                height: task.height(),
             });
             results_sender.try_send(result).unwrap();
         }

crates/apollo_committer_types/src/communication.rs

@@ -10,7 +10,6 @@ use async_trait::async_trait;
 #[cfg(any(feature = "testing", test))]
 use mockall::automock;
 use serde::{Deserialize, Serialize};
-use starknet_api::block::BlockNumber;
 use strum::{EnumVariantNames, VariantNames};
 use strum_macros::{AsRefStr, EnumDiscriminants, EnumIter, IntoStaticStr};
 
@@ -55,15 +54,6 @@ pub enum CommitterRequest {
     RevertBlock(RevertBlockRequest),
 }
 
-impl CommitterRequest {
-    pub fn height(&self) -> BlockNumber {
-        match self {
-            CommitterRequest::CommitBlock(request) => request.height,
-            CommitterRequest::RevertBlock(request) => request.height,
-        }
-    }
-}
-
 impl_debug_for_infra_requests_and_responses!(CommitterRequest);
 impl_labeled_request!(CommitterRequest, CommitterRequestLabelValue);
 impl PrioritizedRequest for CommitterRequest {}