apollo_consensus_orchestrator: improve inputs to build_proposal::get_proposal_content (#10091)

Author: guy-starkware

crates/apollo_consensus_orchestrator/src/build_proposal.rs

@@ -2,6 +2,7 @@
 #[path = "build_proposal_test.rs"]
 mod build_proposal_test;
 
+use std::borrow::BorrowMut;
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
@@ -11,11 +12,8 @@ use apollo_batcher_types::batcher_types::{
     ProposalId,
     ProposeBlockInput,
 };
-use apollo_batcher_types::communication::{BatcherClient, BatcherClientError};
-use apollo_class_manager_types::transaction_converter::{
-    TransactionConverterError,
-    TransactionConverterTrait,
-};
+use apollo_batcher_types::communication::BatcherClientError;
+use apollo_class_manager_types::transaction_converter::TransactionConverterError;
 use apollo_consensus::types::{ProposalCommitment, Round};
 use apollo_l1_gas_price_types::errors::{EthToStrkOracleClientError, L1GasPriceClientError};
 use apollo_protobuf::consensus::{
@@ -25,7 +23,7 @@ use apollo_protobuf::consensus::{
     ProposalPart,
     TransactionBatch,
 };
-use apollo_time::time::{Clock, DateTime};
+use apollo_time::time::DateTime;
 use starknet_api::block::{BlockNumber, GasPrice};
 use starknet_api::consensus_transaction::InternalConsensusTransaction;
 use starknet_api::core::ContractAddress;
@@ -116,17 +114,7 @@ pub(crate) async fn build_proposal(
         .await
         .expect("Failed to send block info");
 
-    let (proposal_commitment, content) = get_proposal_content(
-        args.proposal_id,
-        args.deps.batcher.as_ref(),
-        args.stream_sender,
-        args.cende_write_success,
-        args.deps.transaction_converter,
-        args.cancel_token,
-        args.deps.clock,
-        batcher_deadline,
-    )
-    .await?;
+    let (proposal_commitment, content) = get_proposal_content(&mut args, batcher_deadline).await?;
 
     // Update valid_proposals before sending fin to avoid a race condition
     // with `repropose` being called before `valid_proposals` is updated.
@@ -187,33 +175,26 @@ async fn initiate_build(args: &ProposalBuildArguments) -> BuildProposalResult<Co
 /// 1. Receive chunks of content from the batcher.
 /// 2. Forward these to the stream handler to be streamed out to the network.
 /// 3. Once finished, receive the commitment from the batcher.
-// TODO(guyn): consider passing a ref to BuildProposalArguments instead of all the fields
-// separately.
-#[allow(clippy::too_many_arguments)]
 async fn get_proposal_content(
-    proposal_id: ProposalId,
-    batcher: &dyn BatcherClient,
-    mut stream_sender: StreamSender,
-    cende_write_success: AbortOnDropHandle<bool>,
-    transaction_converter: Arc<dyn TransactionConverterTrait>,
-    cancel_token: CancellationToken,
-    clock: Arc<dyn Clock>,
+    args: &mut ProposalBuildArguments,
     batcher_deadline: DateTime,
 ) -> BuildProposalResult<(ProposalCommitment, Vec<Vec<InternalConsensusTransaction>>)> {
     let mut content = Vec::new();
     loop {
-        if cancel_token.is_cancelled() {
+        if args.cancel_token.is_cancelled() {
             return Err(BuildProposalError::Interrupted);
         }
         // We currently want one part of the node failing to cause all components to fail. If this
         // changes, we can simply return None and consider this as a failed proposal which consensus
         // should support.
-        let response = batcher
-            .get_proposal_content(GetProposalContentInput { proposal_id })
+        let response = args
+            .deps
+            .batcher
+            .get_proposal_content(GetProposalContentInput { proposal_id: args.proposal_id })
             .await
             .map_err(|err| {
                 BuildProposalError::Batcher(
-                    format!("Failed to get proposal content for proposal_id {proposal_id}."),
+                    format!("Failed to get proposal content for proposal_id {}.", args.proposal_id),
                     err,
                 )
             })?;
@@ -228,14 +209,16 @@ async fn get_proposal_content(
                     txs.len()
                 );
                 let transactions = futures::future::join_all(txs.into_iter().map(|tx| {
-                    transaction_converter.convert_internal_consensus_tx_to_consensus_tx(tx)
+                    args.deps
+                        .transaction_converter
+                        .convert_internal_consensus_tx_to_consensus_tx(tx)
                 }))
                 .await
                 .into_iter()
                 .collect::<Result<Vec<_>, _>>()?;
 
                 trace!(?transactions, "Sending transaction batch with {} txs.", transactions.len());
-                stream_sender
+                args.stream_sender
                     .send(ProposalPart::Transactions(TransactionBatch { transactions }))
                     .await
                     .expect("Failed to broadcast proposal content");
@@ -256,11 +239,11 @@ async fn get_proposal_content(
                 // If the blob writing operation to Aerospike doesn't return a success status, we
                 // can't finish the proposal. Must wait for it at least until batcher_timeout is
                 // reached.
-                let remaining = (batcher_deadline - clock.now())
+                let remaining = (batcher_deadline - args.deps.clock.now())
                     .to_std()
                     .unwrap_or_default()
                     .max(MIN_WAIT_DURATION); // Ensure we wait at least 1 ms to avoid immediate timeout. 
-                match tokio::time::timeout(remaining, cende_write_success).await {
+                match tokio::time::timeout(remaining, args.cende_write_success.borrow_mut()).await {
                     Err(_) => {
                         return Err(BuildProposalError::CendeWriteError(
                             "Writing blob to Aerospike didn't return in time.".to_string(),
@@ -282,13 +265,13 @@ async fn get_proposal_content(
                 let final_n_executed_txs_u64 = final_n_executed_txs
                     .try_into()
                     .expect("Number of executed transactions should fit in u64");
-                stream_sender
+                args.stream_sender
                     .send(ProposalPart::ExecutedTransactionCount(final_n_executed_txs_u64))
                     .await
                     .expect("Failed to broadcast executed transaction count");
                 let fin = ProposalFin { proposal_commitment };
                 info!("Sending fin={fin:?}");
-                stream_sender
+                args.stream_sender
                     .send(ProposalPart::Fin(fin))
                     .await
                     .expect("Failed to broadcast proposal fin");