From bf8011e5028a1b35ce38ad38a1c93072a063a38d Mon Sep 17 00:00:00 2001 From: dafnamatsry <92669167+dafnamatsry@users.noreply.github.com> Date: Mon, 9 Dec 2024 13:47:47 +0200 Subject: [PATCH] refactor(starknet_batcher): add getters to the proposal manger (#2480) --- crates/starknet_batcher/src/batcher_test.rs | 22 ++++++++++ .../starknet_batcher/src/proposal_manager.rs | 40 ++++++++++++++----- 2 files changed, 52 insertions(+), 10 deletions(-) diff --git a/crates/starknet_batcher/src/batcher_test.rs b/crates/starknet_batcher/src/batcher_test.rs index e6ca53d1e1..05ab9c25ae 100644 --- a/crates/starknet_batcher/src/batcher_test.rs +++ b/crates/starknet_batcher/src/batcher_test.rs @@ -580,6 +580,14 @@ trait ProposalManagerTraitWrapper: Send + Sync { proposal_id: ProposalId, ) -> BoxFuture<'_, ProposalResult<ProposalOutput>>; + fn wrap_get_active_proposal(&self) -> BoxFuture<'_, Option<ProposalId>>; + + fn wrap_get_completed_proposals( + &self, + ) -> BoxFuture<'_, Arc<tokio::sync::Mutex<HashMap<ProposalId, ProposalResult<ProposalOutput>>>>>; + + fn wrap_await_active_proposal(&mut self) -> BoxFuture<'_, bool>; + fn wrap_get_proposal_status( &self, proposal_id: ProposalId, @@ -613,6 +621,20 @@ impl<T: ProposalManagerTraitWrapper> ProposalManagerTrait for T { self.wrap_take_proposal_result(proposal_id).await } + async fn get_active_proposal(&self) -> Option<ProposalId> { + self.wrap_get_active_proposal().await + } + + async fn get_completed_proposals( + &self, + ) -> Arc<tokio::sync::Mutex<HashMap<ProposalId, ProposalResult<ProposalOutput>>>> { + self.wrap_get_completed_proposals().await + } + + async fn await_active_proposal(&mut self) -> bool { + self.wrap_await_active_proposal().await + } + async fn get_proposal_status(&self, proposal_id: ProposalId) -> InternalProposalStatus { self.wrap_get_proposal_status(proposal_id).await } diff --git a/crates/starknet_batcher/src/proposal_manager.rs b/crates/starknet_batcher/src/proposal_manager.rs index d588861c10..b910e46280 100644 --- a/crates/starknet_batcher/src/proposal_manager.rs +++ b/crates/starknet_batcher/src/proposal_manager.rs @@ -63,6 +63,16 @@ pub trait ProposalManagerTrait: Send + Sync { proposal_id: ProposalId, ) -> ProposalResult<ProposalOutput>; + #[allow(dead_code)] + async fn get_active_proposal(&self) -> Option<ProposalId>; + + #[allow(dead_code)] + async fn get_completed_proposals( + &self, + ) -> Arc<Mutex<HashMap<ProposalId, ProposalResult<ProposalOutput>>>>; + + async fn await_active_proposal(&mut self) -> bool; + async fn get_proposal_status(&self, proposal_id: ProposalId) -> InternalProposalStatus; async fn await_proposal_commitment( @@ -162,6 +172,26 @@ impl ProposalManagerTrait for ProposalManager { .ok_or(GetProposalResultError::ProposalDoesNotExist { proposal_id })? } + async fn get_active_proposal(&self) -> Option<ProposalId> { + *self.active_proposal.lock().await + } + + async fn get_completed_proposals( + &self, + ) -> Arc<Mutex<HashMap<ProposalId, ProposalResult<ProposalOutput>>>> { + self.executed_proposals.clone() + } + + // Awaits the active proposal. + // Returns true if there was an active proposal, and false otherwise. + async fn await_active_proposal(&mut self) -> bool { + if let Some(proposal_task) = self.active_proposal_task.take() { + proposal_task.join_handle.await.ok(); + return true; + } + false + } + // Returns None if the proposal does not exist, otherwise, returns the status of the proposal. async fn get_proposal_status(&self, proposal_id: ProposalId) -> InternalProposalStatus { match self.executed_proposals.lock().await.get(&proposal_id) { @@ -245,16 +275,6 @@ impl ProposalManager { Ok(()) } - // Awaits the active proposal. - // Returns true if there was an active proposal, and false otherwise. - pub async fn await_active_proposal(&mut self) -> bool { - if let Some(proposal_task) = self.active_proposal_task.take() { - proposal_task.join_handle.await.ok(); - return true; - } - false - } - // Ends the current active proposal. // This call is non-blocking. async fn abort_active_proposal(&mut self) {