wip

aptos-labs · Oct 1, 2024 · 87c18e4 · 87c18e4
1 parent 7bdd146
commit 87c18e4
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/consensus/src/quorum_store/batch_proof_queue.rs b/consensus/src/quorum_store/batch_proof_queue.rs
@@ -7,7 +7,7 @@ use super::{
 };
 use crate::quorum_store::counters;
 use aptos_consensus_types::{
-    common::TxnSummaryWithExpiration,
+    common::{Author, TxnSummaryWithExpiration},
     payload::TDataInfo,
     proof_of_store::{BatchInfo, ProofOfStore},
     utils::PayloadTxnsSize,
@@ -389,6 +389,7 @@ impl BatchProofQueue {
         let (result, all_txns, unique_txns, is_full) = self.pull_internal(
             false,
             excluded_batches,
+            &HashSet::new(),
             max_txns,
             max_txns_after_filtering,
             soft_max_txns_after_filtering,
@@ -429,6 +430,7 @@ impl BatchProofQueue {
     pub fn pull_batches(
         &mut self,
         excluded_batches: &HashSet<BatchInfo>,
+        exclude_authors: &HashSet<Author>,
         max_txns: PayloadTxnsSize,
         max_txns_after_filtering: u64,
         soft_max_txns_after_filtering: u64,
@@ -438,6 +440,7 @@ impl BatchProofQueue {
         let (result, all_txns, unique_txns, _) = self.pull_internal(
             true,
             excluded_batches,
+            exclude_authors,
             max_txns,
             max_txns_after_filtering,
             soft_max_txns_after_filtering,
@@ -463,6 +466,7 @@ impl BatchProofQueue {
     ) {
         let (batches, all_txns, unique_txns) = self.pull_batches(
             excluded_batches,
+            &HashSet::new(),
             max_txns,
             max_txns_after_filtering,
             soft_max_txns_after_filtering,
@@ -489,6 +493,7 @@ impl BatchProofQueue {
         &mut self,
         batches_without_proofs: bool,
         excluded_batches: &HashSet<BatchInfo>,
+        exclude_authors: &HashSet<Author>,
         max_txns: PayloadTxnsSize,
         max_txns_after_filtering: u64,
         soft_max_txns_after_filtering: u64,
@@ -516,7 +521,11 @@ impl BatchProofQueue {
         }
 
         let mut iters = vec![];
-        for (_, batches) in self.author_to_batches.iter() {
+        for (_, batches) in self
+            .author_to_batches
+            .iter()
+            .filter(|(author, _)| !exclude_authors.contains(author))
+        {
             let batch_iter = batches.iter().rev().filter_map(|(sort_key, info)| {
                 if let Some(item) = self.items.get(&sort_key.batch_key) {
                     if item.is_committed() {

diff --git a/consensus/src/quorum_store/proof_manager.rs b/consensus/src/quorum_store/proof_manager.rs
@@ -136,13 +136,18 @@ impl ProofManager {
 
                 let (opt_batches, opt_batch_txns_size) = if self.enable_opt_quorum_store {
                     // TODO(ibalajiarun): Support unique txn calculation
+                    let exclude_authors = request
+                        .maybe_optqs_payload_pull_params
+                        .map(|p| p.exclude_authors)
+                        .unwrap_or_default();
                     let max_opt_batch_txns_size = request.max_txns - txns_with_proof_size;
                     let (opt_batches, opt_payload_size, _) = self.batch_proof_queue.pull_batches(
                         &excluded_batches
                             .iter()
                             .cloned()
                             .chain(proof_block.iter().map(|proof| proof.info().clone()))
                             .collect(),
+                        exclude_authors,
                         max_opt_batch_txns_size,
                         request.max_txns_after_filtering,
                         request.soft_max_txns_after_filtering,