diff --git a/crates/starknet_batcher/src/batcher.rs b/crates/starknet_batcher/src/batcher.rs index 0ffcc0fb62..b42a92c3e0 100644 --- a/crates/starknet_batcher/src/batcher.rs +++ b/crates/starknet_batcher/src/batcher.rs @@ -41,8 +41,8 @@ use crate::block_builder::{ use crate::config::BatcherConfig; use crate::proposal_manager::{ GenerateProposalError, - GetProposalResultError, InternalProposalStatus, + ProposalError, ProposalManager, ProposalManagerTrait, ProposalOutput, @@ -280,13 +280,13 @@ impl Batcher { self.close_input_transaction_stream(proposal_id)?; let response = match self.proposal_manager.await_proposal_commitment(proposal_id).await { - Ok(proposal_commitment) => ProposalStatus::Finished(proposal_commitment), - Err(GetProposalResultError::BlockBuilderError(err)) => match err.as_ref() { + Some(Ok(proposal_commitment)) => ProposalStatus::Finished(proposal_commitment), + Some(Err(ProposalError::BlockBuilderError(err))) => match err.as_ref() { BlockBuilderError::FailOnError(_) => ProposalStatus::InvalidProposal, _ => return Err(BatcherError::InternalError), }, - Err(GetProposalResultError::ProposalDoesNotExist { proposal_id: _ }) - | Err(GetProposalResultError::Aborted) => { + Some(Err(ProposalError::Aborted)) => return Err(BatcherError::ProposalAborted), + None => { panic!("Proposal {} should exist in the proposal manager.", proposal_id) } }; @@ -327,8 +327,11 @@ impl Batcher { // TODO: Consider removing the proposal from the proposal manager and keep it in the batcher // for decision reached. self.propose_tx_streams.remove(&proposal_id); - let proposal_commitment = - self.proposal_manager.await_proposal_commitment(proposal_id).await?; + let proposal_commitment = self + .proposal_manager + .await_proposal_commitment(proposal_id) + .await + .ok_or(BatcherError::ProposalNotFound { proposal_id })??; Ok(GetProposalContentResponse { content: GetProposalContent::Finished(proposal_commitment), }) @@ -337,7 +340,11 @@ impl Batcher { #[instrument(skip(self), err)] pub async fn decision_reached(&mut self, input: DecisionReachedInput) -> BatcherResult<()> { let proposal_id = input.proposal_id; - let proposal_output = self.proposal_manager.take_proposal_result(proposal_id).await?; + let proposal_output = self + .proposal_manager + .take_proposal_result(proposal_id) + .await + .ok_or(BatcherError::ExecutedProposalNotFound { proposal_id })??; let ProposalOutput { state_diff, nonces: address_to_nonce, tx_hashes, .. } = proposal_output; // TODO: Keep the height from start_height or get it from the input. @@ -437,14 +444,11 @@ impl From for BatcherError { } } -impl From for BatcherError { - fn from(err: GetProposalResultError) -> Self { +impl From for BatcherError { + fn from(err: ProposalError) -> Self { match err { - GetProposalResultError::BlockBuilderError(..) => BatcherError::InternalError, - GetProposalResultError::ProposalDoesNotExist { proposal_id } => { - BatcherError::ExecutedProposalNotFound { proposal_id } - } - GetProposalResultError::Aborted => BatcherError::ProposalAborted, + ProposalError::BlockBuilderError(..) => BatcherError::InternalError, + ProposalError::Aborted => BatcherError::ProposalAborted, } } } diff --git a/crates/starknet_batcher/src/batcher_test.rs b/crates/starknet_batcher/src/batcher_test.rs index 05ab9c25ae..c4134ea930 100644 --- a/crates/starknet_batcher/src/batcher_test.rs +++ b/crates/starknet_batcher/src/batcher_test.rs @@ -49,8 +49,8 @@ use crate::block_builder::{ use crate::config::BatcherConfig; use crate::proposal_manager::{ GenerateProposalError, - GetProposalResultError, InternalProposalStatus, + ProposalError, ProposalManagerTrait, ProposalOutput, ProposalResult, @@ -156,7 +156,7 @@ fn mock_proposal_manager_common_expectations( .expect_wrap_await_proposal_commitment() .times(1) .with(eq(PROPOSAL_ID)) - .return_once(move |_| { async move { Ok(proposal_commitment()) } }.boxed()); + .return_once(move |_| { async move { Some(Ok(proposal_commitment())) } }.boxed()); } fn mock_proposal_manager_validate_flow() -> MockProposalManagerTraitWrapper { @@ -372,14 +372,14 @@ async fn send_finish_to_an_invalid_proposal() { .with(eq(PROPOSAL_ID), always(), always()) .return_once(|_, _, _| { async move { Ok(()) } }.boxed()); - let proposal_error = GetProposalResultError::BlockBuilderError(Arc::new( + let proposal_error = ProposalError::BlockBuilderError(Arc::new( BlockBuilderError::FailOnError(FailOnErrorCause::BlockFull), )); proposal_manager .expect_wrap_await_proposal_commitment() .times(1) .with(eq(PROPOSAL_ID)) - .return_once(move |_| { async move { Err(proposal_error) } }.boxed()); + .return_once(move |_| { async move { Some(Err(proposal_error)) } }.boxed()); let mut batcher = create_batcher(MockDependencies { proposal_manager, @@ -517,12 +517,12 @@ async fn decision_reached() { .with(eq(PROPOSAL_ID)) .return_once(move |_| { async move { - Ok(ProposalOutput { + Some(Ok(ProposalOutput { state_diff: ThinStateDiff::default(), commitment: ProposalCommitment::default(), tx_hashes: test_tx_hashes(), nonces: test_contract_nonces(), - }) + })) } .boxed() }); @@ -553,11 +553,11 @@ async fn decision_reached_no_executed_proposal() { let expected_error = BatcherError::ExecutedProposalNotFound { proposal_id: PROPOSAL_ID }; let mut proposal_manager = MockProposalManagerTraitWrapper::new(); - proposal_manager.expect_wrap_take_proposal_result().times(1).with(eq(PROPOSAL_ID)).return_once( - |proposal_id| { - async move { Err(GetProposalResultError::ProposalDoesNotExist { proposal_id }) }.boxed() - }, - ); + proposal_manager + .expect_wrap_take_proposal_result() + .times(1) + .with(eq(PROPOSAL_ID)) + .return_once(|_| async move { None }.boxed()); let mut batcher = create_batcher(MockDependencies { proposal_manager, ..Default::default() }); let decision_reached_result = @@ -578,7 +578,7 @@ trait ProposalManagerTraitWrapper: Send + Sync { fn wrap_take_proposal_result( &mut self, proposal_id: ProposalId, - ) -> BoxFuture<'_, ProposalResult>; + ) -> BoxFuture<'_, Option>>; fn wrap_get_active_proposal(&self) -> BoxFuture<'_, Option>; @@ -596,7 +596,7 @@ trait ProposalManagerTraitWrapper: Send + Sync { fn wrap_await_proposal_commitment( &self, proposal_id: ProposalId, - ) -> BoxFuture<'_, ProposalResult>; + ) -> BoxFuture<'_, Option>>; fn wrap_abort_proposal(&mut self, proposal_id: ProposalId) -> BoxFuture<'_, ()>; @@ -617,7 +617,7 @@ impl ProposalManagerTrait for T { async fn take_proposal_result( &mut self, proposal_id: ProposalId, - ) -> ProposalResult { + ) -> Option> { self.wrap_take_proposal_result(proposal_id).await } @@ -642,7 +642,7 @@ impl ProposalManagerTrait for T { async fn await_proposal_commitment( &mut self, proposal_id: ProposalId, - ) -> ProposalResult { + ) -> Option> { self.wrap_await_proposal_commitment(proposal_id).await } diff --git a/crates/starknet_batcher/src/proposal_manager.rs b/crates/starknet_batcher/src/proposal_manager.rs index f4dabdbbc2..7f27dffdc0 100644 --- a/crates/starknet_batcher/src/proposal_manager.rs +++ b/crates/starknet_batcher/src/proposal_manager.rs @@ -33,11 +33,9 @@ pub enum GenerateProposalError { } #[derive(Clone, Debug, Error)] -pub enum GetProposalResultError { +pub enum ProposalError { #[error(transparent)] BlockBuilderError(Arc), - #[error("Proposal with id {proposal_id} does not exist.")] - ProposalDoesNotExist { proposal_id: ProposalId }, #[error("Proposal was aborted")] Aborted, } @@ -61,7 +59,7 @@ pub trait ProposalManagerTrait: Send + Sync { async fn take_proposal_result( &mut self, proposal_id: ProposalId, - ) -> ProposalResult; + ) -> Option>; #[allow(dead_code)] async fn get_active_proposal(&self) -> Option; @@ -78,7 +76,7 @@ pub trait ProposalManagerTrait: Send + Sync { async fn await_proposal_commitment( &mut self, proposal_id: ProposalId, - ) -> ProposalResult; + ) -> Option>; async fn abort_proposal(&mut self, proposal_id: ProposalId); @@ -109,7 +107,7 @@ pub(crate) struct ProposalManager { executed_proposals: Arc>>>, } -pub type ProposalResult = Result; +pub type ProposalResult = Result; #[derive(Debug, PartialEq)] pub struct ProposalOutput { @@ -143,7 +141,7 @@ impl ProposalManagerTrait for ProposalManager { .build_block() .await .map(ProposalOutput::from) - .map_err(|e| GetProposalResultError::BlockBuilderError(Arc::new(e))); + .map_err(|e| ProposalError::BlockBuilderError(Arc::new(e))); // The proposal is done, clear the active proposal. // Keep the proposal result only if it is the same as the active proposal. @@ -164,12 +162,8 @@ impl ProposalManagerTrait for ProposalManager { async fn take_proposal_result( &mut self, proposal_id: ProposalId, - ) -> ProposalResult { - self.executed_proposals - .lock() - .await - .remove(&proposal_id) - .ok_or(GetProposalResultError::ProposalDoesNotExist { proposal_id })? + ) -> Option> { + self.executed_proposals.lock().await.remove(&proposal_id) } async fn get_active_proposal(&self) -> Option { @@ -210,17 +204,16 @@ impl ProposalManagerTrait for ProposalManager { async fn await_proposal_commitment( &mut self, proposal_id: ProposalId, - ) -> ProposalResult { + ) -> Option> { if *self.active_proposal.lock().await == Some(proposal_id) { self.await_active_proposal().await; } let proposals = self.executed_proposals.lock().await; - let output = proposals - .get(&proposal_id) - .ok_or(GetProposalResultError::ProposalDoesNotExist { proposal_id })?; + let output = proposals.get(&proposal_id); match output { - Ok(output) => Ok(output.commitment), - Err(e) => Err(e.clone()), + Some(Ok(output)) => Some(Ok(output.commitment)), + Some(Err(e)) => Some(Err(e.clone())), + None => None, } } @@ -229,10 +222,7 @@ impl ProposalManagerTrait for ProposalManager { async fn abort_proposal(&mut self, proposal_id: ProposalId) { if *self.active_proposal.lock().await == Some(proposal_id) { self.abort_active_proposal().await; - self.executed_proposals - .lock() - .await - .insert(proposal_id, Err(GetProposalResultError::Aborted)); + self.executed_proposals.lock().await.insert(proposal_id, Err(ProposalError::Aborted)); } } diff --git a/crates/starknet_batcher/src/proposal_manager_test.rs b/crates/starknet_batcher/src/proposal_manager_test.rs index 318ddb619e..2f92a8492b 100644 --- a/crates/starknet_batcher/src/proposal_manager_test.rs +++ b/crates/starknet_batcher/src/proposal_manager_test.rs @@ -6,7 +6,7 @@ use starknet_batcher_types::batcher_types::ProposalId; use crate::block_builder::{BlockBuilderTrait, BlockExecutionArtifacts, MockBlockBuilderTrait}; use crate::proposal_manager::{ GenerateProposalError, - GetProposalResultError, + ProposalError, ProposalManager, ProposalManagerTrait, ProposalOutput, @@ -70,7 +70,7 @@ async fn spawn_proposal( async fn spawn_proposal_success(mut proposal_manager: ProposalManager) { spawn_proposal(&mut proposal_manager, ProposalId(0), mock_build_block()).await; - proposal_manager.take_proposal_result(ProposalId(0)).await.unwrap(); + proposal_manager.take_proposal_result(ProposalId(0)).await.unwrap().unwrap(); } #[rstest] @@ -117,13 +117,10 @@ async fn take_proposal_result_no_active_proposal(mut proposal_manager: ProposalM let expected_proposal_output = ProposalOutput::from(BlockExecutionArtifacts::create_for_testing()); assert_eq!( - proposal_manager.take_proposal_result(ProposalId(0)).await.unwrap(), + proposal_manager.take_proposal_result(ProposalId(0)).await.unwrap().unwrap(), expected_proposal_output ); - assert_matches!( - proposal_manager.take_proposal_result(ProposalId(0)).await, - Err(GetProposalResultError::ProposalDoesNotExist { .. }) - ); + assert_matches!(proposal_manager.take_proposal_result(ProposalId(0)).await, None); } #[rstest] @@ -137,7 +134,7 @@ async fn abort_active_proposal(mut proposal_manager: ProposalManager) { assert_matches!( proposal_manager.take_proposal_result(ProposalId(0)).await, - Err(GetProposalResultError::Aborted) + Some(Err(ProposalError::Aborted)) ); // Make sure there is no active proposal. @@ -156,10 +153,7 @@ async fn reset(mut proposal_manager: ProposalManager) { proposal_manager.reset().await; // Make sure executed proposals are deleted. - assert_matches!( - proposal_manager.take_proposal_result(ProposalId(0)).await, - Err(GetProposalResultError::ProposalDoesNotExist { .. }) - ); + assert_matches!(proposal_manager.take_proposal_result(ProposalId(0)).await, None); // Make sure there is no active proposal. assert!(!proposal_manager.await_active_proposal().await);