From ddaf489f63e3dfa181db80bf308d260e6cba906e Mon Sep 17 00:00:00 2001
From: Brian Cho <brian@aptoslabs.com>
Date: Mon, 4 Nov 2024 23:29:45 -0800
Subject: [PATCH] also constrain windowed blocks to max_block_txns

---
 consensus/src/block_preparer.rs     |  8 ++------
 consensus/src/execution_pipeline.rs | 18 ++++++++++++------
 consensus/src/state_computer.rs     |  1 +
 3 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/consensus/src/block_preparer.rs b/consensus/src/block_preparer.rs
index 306db22e14208..3d9e6f677b28d 100644
--- a/consensus/src/block_preparer.rs
+++ b/consensus/src/block_preparer.rs
@@ -155,12 +155,8 @@ impl BlockPreparer {
         let txn_deduper = self.txn_deduper.clone();
         let block_id = block.id();
         let block_timestamp_usecs = block.timestamp_usecs();
-        let max_prepared_block_txns =
-            if let Some(max_txns_from_block_to_execute) = max_txns_from_block_to_execute {
-                max_txns_from_block_to_execute
-            } else {
-                self.max_block_txns * 2
-            };
+        // Always use max_block_txns * 2 regardless of max_txns_from_block_to_execute for better shuffling
+        let max_prepared_block_txns = self.max_block_txns * 2;
         // Transaction filtering, deduplication and shuffling are CPU intensive tasks, so we run them in a blocking task.
         let result = tokio::task::spawn_blocking(move || {
             // stable sort to ensure batches with same gas are in the same order
diff --git a/consensus/src/execution_pipeline.rs b/consensus/src/execution_pipeline.rs
index c46d6d867f9ae..6103c3dd26462 100644
--- a/consensus/src/execution_pipeline.rs
+++ b/consensus/src/execution_pipeline.rs
@@ -104,6 +104,7 @@ impl ExecutionPipeline {
         &self,
         block: PipelinedBlock,
         block_window: OrderedBlockWindow,
+        max_block_txns: u64,
         metadata: BlockMetadataExt,
         parent_block_id: HashValue,
         txn_generator: BlockPreparer,
@@ -118,6 +119,7 @@ impl ExecutionPipeline {
             .send(PrepareBlockCommand {
                 block,
                 block_window,
+                max_block_txns,
                 metadata,
                 block_executor_onchain_config,
                 parent_block_id,
@@ -153,6 +155,7 @@ impl ExecutionPipeline {
         let PrepareBlockCommand {
             block,
             block_window,
+            max_block_txns,
             metadata,
             block_executor_onchain_config,
             parent_block_id,
@@ -194,6 +197,7 @@ impl ExecutionPipeline {
             execute_block_tx
                 .send(ExecuteBlockCommand {
                     input_txns,
+                    max_block_txns,
                     max_txns_to_execute,
                     pipelined_block: block,
                     block: (block_id, sig_verified_txns).into(),
@@ -231,6 +235,7 @@ impl ExecutionPipeline {
     ) {
         'outer: while let Some(ExecuteBlockCommand {
             input_txns: _,
+            max_block_txns,
             max_txns_to_execute,
             pipelined_block,
             block,
@@ -331,12 +336,11 @@ impl ExecutionPipeline {
             } else {
                 txns.len()
             };
-            let num_txns_to_execute = if let Some(max_user_txns_to_execute) = max_txns_to_execute {
-                txns.len()
-                    .min(num_validator_txns.saturating_add(max_user_txns_to_execute as usize))
-            } else {
-                txns.len()
-            };
+            let mut num_txns_to_execute = txns.len().min(max_block_txns as usize);
+            if let Some(max_user_txns_to_execute) = max_txns_to_execute {
+                num_txns_to_execute = num_txns_to_execute
+                    .min(num_validator_txns.saturating_add(max_user_txns_to_execute as usize));
+            }
             let blocking_txns_provider = Arc::new(BlockingTxnsProvider::new(num_txns_to_execute));
             let blocking_txns_writer = blocking_txns_provider.clone();
             let join_shuffle = tokio::task::spawn_blocking(move || {
@@ -583,6 +587,7 @@ impl ExecutionPipeline {
 struct PrepareBlockCommand {
     block: PipelinedBlock,
     block_window: OrderedBlockWindow,
+    max_block_txns: u64,
     metadata: BlockMetadataExt,
     block_executor_onchain_config: BlockExecutorConfigFromOnchain,
     // The parent block id.
@@ -596,6 +601,7 @@ struct PrepareBlockCommand {
 
 struct ExecuteBlockCommand {
     input_txns: Vec<SignedTransaction>,
+    max_block_txns: u64,
     max_txns_to_execute: Option<u64>,
     pipelined_block: PipelinedBlock,
     block: ExecutableBlock,
diff --git a/consensus/src/state_computer.rs b/consensus/src/state_computer.rs
index 2cc2aed9bd853..888ceaa4c0fc5 100644
--- a/consensus/src/state_computer.rs
+++ b/consensus/src/state_computer.rs
@@ -234,6 +234,7 @@ impl StateComputer for ExecutionProxy {
             .queue(
                 block.clone(),
                 block_window.clone(),
+                self.max_block_txns,
                 metadata.clone(),
                 parent_block_id,
                 transaction_generator,