From f5917a336fa3088a80045b3ca2ffcb5872a48d89 Mon Sep 17 00:00:00 2001 From: Yael Doweck Date: Sun, 10 Nov 2024 18:00:00 +0200 Subject: [PATCH] refactor(batcher): move build_block parameter into the block_builder object --- crates/batcher/src/block_builder.rs | 63 ++++++++++++++------- crates/batcher/src/block_builder_test.rs | 11 +++- crates/batcher/src/proposal_manager.rs | 13 ++++- crates/batcher/src/proposal_manager_test.rs | 8 +-- 4 files changed, 65 insertions(+), 30 deletions(-) diff --git a/crates/batcher/src/block_builder.rs b/crates/batcher/src/block_builder.rs index 66c67ea2c4..dc0590e21e 100644 --- a/crates/batcher/src/block_builder.rs +++ b/crates/batcher/src/block_builder.rs @@ -77,41 +77,47 @@ pub struct BlockExecutionArtifacts { #[cfg_attr(test, automock)] #[async_trait] pub trait BlockBuilderTrait: Send { - async fn build_block( - &self, - deadline: tokio::time::Instant, - tx_provider: Box, - output_content_sender: Option>, - fail_on_err: bool, - ) -> BlockBuilderResult; + async fn build_block(&mut self) -> BlockBuilderResult; } pub struct BlockBuilder { // TODO(Yael 14/10/2024): make the executor thread safe and delete this mutex. executor: Mutex>, tx_chunk_size: usize, + deadline: tokio::time::Instant, + tx_provider: Box, + output_content_sender: Option>, + fail_on_err: bool, } impl BlockBuilder { - pub fn new(executor: Box, tx_chunk_size: usize) -> Self { - Self { executor: Mutex::new(executor), tx_chunk_size } + pub fn new( + executor: Box, + tx_chunk_size: usize, + deadline: tokio::time::Instant, + tx_provider: Box, + output_content_sender: Option>, + fail_on_err: bool, + ) -> Self { + Self { + executor: Mutex::new(executor), + tx_chunk_size, + deadline, + tx_provider, + output_content_sender, + fail_on_err, + } } } #[async_trait] impl BlockBuilderTrait for BlockBuilder { - async fn build_block( - &self, - deadline: tokio::time::Instant, - mut tx_provider: Box, - output_content_sender: Option>, - fail_on_err: bool, - ) -> BlockBuilderResult { + async fn build_block(&mut self) -> BlockBuilderResult { let mut block_is_full = false; let mut execution_infos = IndexMap::new(); // TODO(yael 6/10/2024): delete the timeout condition once the executor has a timeout - while !block_is_full && tokio::time::Instant::now() < deadline { - let next_txs = tx_provider.get_txs(self.tx_chunk_size).await?; + while !block_is_full && tokio::time::Instant::now() < self.deadline { + let next_txs = self.tx_provider.get_txs(self.tx_chunk_size).await?; let next_tx_chunk = match next_txs { NextTxs::Txs(txs) => txs, NextTxs::End => break, @@ -134,8 +140,8 @@ impl BlockBuilderTrait for BlockBuilder { next_tx_chunk, results, &mut execution_infos, - &output_content_sender, - fail_on_err, + &self.output_content_sender, + self.fail_on_err, ) .await?; } @@ -190,6 +196,10 @@ pub trait BlockBuilderFactoryTrait { &self, height: BlockNumber, retrospective_block_hash: Option, + deadline: tokio::time::Instant, + tx_provider: Box, + output_content_sender: Option>, + fail_on_err: bool, ) -> BlockBuilderResult>; } @@ -313,9 +323,20 @@ impl BlockBuilderFactoryTrait for BlockBuilderFactory { &self, height: BlockNumber, retrospective_block_hash: Option, + deadline: tokio::time::Instant, + tx_provider: Box, + output_content_sender: Option>, + fail_on_err: bool, ) -> BlockBuilderResult> { let executor = self.preprocess_and_create_transaction_executor(height, retrospective_block_hash)?; - Ok(Box::new(BlockBuilder::new(Box::new(executor), self.block_builder_config.tx_chunk_size))) + Ok(Box::new(BlockBuilder::new( + Box::new(executor), + self.block_builder_config.tx_chunk_size, + deadline, + tx_provider, + output_content_sender, + fail_on_err, + ))) } } diff --git a/crates/batcher/src/block_builder_test.rs b/crates/batcher/src/block_builder_test.rs index 72f91e987b..8bf1d8f892 100644 --- a/crates/batcher/src/block_builder_test.rs +++ b/crates/batcher/src/block_builder_test.rs @@ -293,11 +293,18 @@ async fn run_build_block( output_sender: Option>, fail_on_err: bool, ) -> BlockBuilderResult { - let block_builder = BlockBuilder::new(Box::new(mock_transaction_executor), TX_CHUNK_SIZE); let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(BLOCK_GENERATION_DEADLINE_SECS); + let mut block_builder = BlockBuilder::new( + Box::new(mock_transaction_executor), + TX_CHUNK_SIZE, + deadline, + Box::new(tx_provider), + output_sender, + fail_on_err, + ); - block_builder.build_block(deadline, Box::new(tx_provider), output_sender, fail_on_err).await + block_builder.build_block().await } // TODO: Add test case for failed transaction. diff --git a/crates/batcher/src/proposal_manager.rs b/crates/batcher/src/proposal_manager.rs index 970fc26a38..cadcc0b2c1 100644 --- a/crates/batcher/src/proposal_manager.rs +++ b/crates/batcher/src/proposal_manager.rs @@ -170,18 +170,25 @@ impl ProposalManagerTrait for ProposalManager { info!("Starting generation of a new proposal with id {}.", proposal_id); let height = self.active_height.expect("No active height."); - let block_builder = - self.block_builder_factory.create_block_builder(height, retrospective_block_hash)?; let tx_provider = ProposeTransactionProvider { mempool_client: self.mempool_client.clone() }; + let mut block_builder = self.block_builder_factory.create_block_builder( + height, + retrospective_block_hash, + deadline, + Box::new(tx_provider), + Some(tx_sender.clone()), + false, + )?; + let active_proposal = self.active_proposal.clone(); let executed_proposals = self.executed_proposals.clone(); self.active_proposal_handle = Some(tokio::spawn( async move { let result = block_builder - .build_block(deadline, Box::new(tx_provider), Some(tx_sender.clone()), false) + .build_block() .await .map(ProposalOutput::from) .map_err(|e| GetProposalResultError::BlockBuilderError(Arc::new(e))); diff --git a/crates/batcher/src/proposal_manager_test.rs b/crates/batcher/src/proposal_manager_test.rs index d9b5c97efe..a30c27f2f8 100644 --- a/crates/batcher/src/proposal_manager_test.rs +++ b/crates/batcher/src/proposal_manager_test.rs @@ -48,14 +48,14 @@ impl MockDependencies { let mut mock_block_builder = MockBlockBuilderTrait::new(); mock_block_builder .expect_build_block() - .return_once(move |_, _, _, _| Ok(BlockExecutionArtifacts::create_for_testing())); + .return_once(move || Ok(BlockExecutionArtifacts::create_for_testing())); Ok(Box::new(mock_block_builder)) }; self.block_builder_factory .expect_create_block_builder() .times(times) - .returning(move |_, _| simulate_build_block()); + .returning(move |_, _, _, _, _, _| simulate_build_block()); } // This function simulates a long build block operation. This is required for a test that @@ -63,7 +63,7 @@ impl MockDependencies { fn expect_long_build_block(&mut self, times: usize) { let simulate_long_build_block = || -> BlockBuilderResult> { let mut mock_block_builder = MockBlockBuilderTrait::new(); - mock_block_builder.expect_build_block().return_once(move |_, _, _, _| { + mock_block_builder.expect_build_block().return_once(move || { std::thread::sleep(BLOCK_GENERATION_TIMEOUT * 10); Ok(BlockExecutionArtifacts::create_for_testing()) }); @@ -73,7 +73,7 @@ impl MockDependencies { self.block_builder_factory .expect_create_block_builder() .times(times) - .returning(move |_, _| simulate_long_build_block()); + .returning(move |_, _, _, _, _, _| simulate_long_build_block()); } }