[qs] send proof v2 support with flag

ibalajiarun · ibalajiarun · commit f39d8044ec3e · 2025-11-07T10:30:45.000-08:00
diff --git a/config/src/config/quorum_store_config.rs b/config/src/config/quorum_store_config.rs
@@ -99,6 +99,8 @@ pub struct QuorumStoreConfig {
     pub enable_opt_quorum_store: bool,
     pub opt_qs_minimum_batch_age_usecs: u64,
     pub enable_payload_v2: bool,
+    /// Boolean flag that controls the usage of `BatchInfoExt::V1`
+    pub enable_proof_v2: bool,
 }
 
 impl Default for QuorumStoreConfig {
@@ -140,6 +142,7 @@ impl Default for QuorumStoreConfig {
             enable_opt_quorum_store: true,
             opt_qs_minimum_batch_age_usecs: Duration::from_millis(50).as_micros() as u64,
             enable_payload_v2: false,
+            enable_proof_v2: false,
         }
     }
 }
diff --git a/consensus/src/network.rs b/consensus/src/network.rs
@@ -29,7 +29,7 @@ use aptos_consensus_types::{
     order_vote_msg::OrderVoteMsg,
     pipeline::{commit_decision::CommitDecision, commit_vote::CommitVote},
     proof_of_store::{
-        BatchInfo, ProofOfStore, ProofOfStoreMsg, SignedBatchInfo, SignedBatchInfoMsg,
+        BatchInfo, BatchInfoExt, ProofOfStore, ProofOfStoreMsg, SignedBatchInfo, SignedBatchInfoMsg,
     },
     proposal_msg::ProposalMsg,
     round_timeout::RoundTimeoutMsg,
@@ -212,9 +212,14 @@ pub trait QuorumStoreSender: Send + Clone {
 
     async fn broadcast_proof_of_store_msg(&mut self, proof_of_stores: Vec<ProofOfStore<BatchInfo>>);
 
+    async fn broadcast_proof_of_store_msg_v2(
+        &mut self,
+        proof_of_stores: Vec<ProofOfStore<BatchInfoExt>>,
+    );
+
     async fn send_proof_of_store_msg_to_self(
         &mut self,
-        proof_of_stores: Vec<ProofOfStore<BatchInfo>>,
+        proof_of_stores: Vec<ProofOfStore<BatchInfoExt>>,
     );
 }
 
@@ -582,9 +587,15 @@ impl QuorumStoreSender for NetworkSender {
         self.broadcast(msg).await
     }
 
-    async fn send_proof_of_store_msg_to_self(&mut self, proofs: Vec<ProofOfStore<BatchInfo>>) {
+    async fn broadcast_proof_of_store_msg_v2(&mut self, proofs: Vec<ProofOfStore<BatchInfoExt>>) {
         fail_point!("consensus::send::proof_of_store", |_| ());
-        let msg = ConsensusMsg::ProofOfStoreMsg(Box::new(ProofOfStoreMsg::new(proofs)));
+        let msg = ConsensusMsg::ProofOfStoreMsgV2(Box::new(ProofOfStoreMsg::new(proofs)));
+        self.broadcast(msg).await
+    }
+
+    async fn send_proof_of_store_msg_to_self(&mut self, proofs: Vec<ProofOfStore<BatchInfoExt>>) {
+        fail_point!("consensus::send::proof_of_store", |_| ());
+        let msg = ConsensusMsg::ProofOfStoreMsgV2(Box::new(ProofOfStoreMsg::new(proofs)));
         self.send(msg, vec![self.author]).await
     }
 }
diff --git a/consensus/src/quorum_store/batch_store.rs b/consensus/src/quorum_store/batch_store.rs
@@ -12,7 +12,7 @@ use crate::{
     },
 };
 use anyhow::bail;
-use aptos_consensus_types::proof_of_store::{BatchInfo, SignedBatchInfo};
+use aptos_consensus_types::proof_of_store::{BatchInfo, BatchInfoExt, SignedBatchInfo};
 use aptos_crypto::{CryptoMaterialError, HashValue};
 use aptos_executor_types::{ExecutorError, ExecutorResult};
 use aptos_infallible::Mutex;
diff --git a/consensus/src/quorum_store/proof_coordinator.rs b/consensus/src/quorum_store/proof_coordinator.rs
@@ -171,6 +171,7 @@ pub(crate) struct ProofCoordinator {
     proof_cache: ProofCache,
     broadcast_proofs: bool,
     batch_expiry_gap_when_init_usecs: u64,
+    enable_proof_v2_msg: bool,
 }
 
 //PoQS builder object - gather signed digest to form PoQS
@@ -183,6 +184,7 @@ impl ProofCoordinator {
         proof_cache: ProofCache,
         broadcast_proofs: bool,
         batch_expiry_gap_when_init_usecs: u64,
+        enable_proof_v2_msg: bool,
     ) -> Self {
         Self {
             peer_id,
@@ -195,6 +197,7 @@ impl ProofCoordinator {
             proof_cache,
             broadcast_proofs,
             batch_expiry_gap_when_init_usecs,
+            enable_proof_v2_msg,
         }
     }
 
@@ -374,45 +377,56 @@ impl ProofCoordinator {
                             let approx_created_ts_usecs = signed_batch_info
                                 .expiration()
                                 .saturating_sub(self.batch_expiry_gap_when_init_usecs);
+                            let self_peer_id = self.peer_id;
+                            let enable_broadcast_proofs = self.broadcast_proofs;
+                            let enable_proof_v2_msg = self.enable_proof_v2_msg;
 
-                            let mut proofs = vec![];
-                            for signed_batch_info in signed_batch_infos.into_iter() {
+                            let mut proofs_iter = signed_batch_infos.into_iter().filter_map(|signed_batch_info| {
                                 let peer_id = signed_batch_info.signer();
                                 let digest = *signed_batch_info.digest();
                                 let batch_id = signed_batch_info.batch_id();
                                 match self.add_signature(signed_batch_info, &validator_verifier) {
-                                    Ok(result) => {
-                                        if let Some(proof) = result {
-                                            debug!(
-                                                LogSchema::new(LogEvent::ProofOfStoreReady),
-                                                digest = digest,
-                                                batch_id = batch_id.id,
-                                            );
-                                            let (info, sig) = proof.unpack();
-                                            proofs.push(ProofOfStore::new(info.info().clone(), sig));
-                                        }
+                                    Ok(Some(proof)) => {
+                                        debug!(
+                                            LogSchema::new(LogEvent::ProofOfStoreReady),
+                                            digest = digest,
+                                            batch_id = batch_id.id,
+                                        );
+                                        Some(proof)
                                     },
+                                    Ok(None) => None,
                                     Err(e) => {
                                         // Can happen if we already garbage collected, the commit notification is late, or the peer is misbehaving.
                                         if peer_id == self.peer_id {
                                             info!("QS: could not add signature from self, digest = {}, batch_id = {}, err = {:?}", digest, batch_id, e);
                                         } else {
                                             debug!("QS: could not add signature from peer {}, digest = {}, batch_id = {}, err = {:?}", peer_id, digest, batch_id, e);
                                         }
+                                        None
                                     },
                                 }
-                            }
-                            if let Some(value) = self.batch_info_to_proof.get_mut(&info) {
-                                value.observe_voting_pct(approx_created_ts_usecs, &validator_verifier);
-                            }
-                            if !proofs.is_empty() {
-                                observe_batch(approx_created_ts_usecs, self.peer_id, BatchStage::POS_FORMED);
-                                if self.broadcast_proofs {
-                                    network_sender.broadcast_proof_of_store_msg(proofs).await;
+                            }).peekable();
+                            if proofs_iter.peek().is_some() {
+                                observe_batch(approx_created_ts_usecs, self_peer_id, BatchStage::POS_FORMED);
+                                if enable_broadcast_proofs {
+                                    if enable_proof_v2_msg {
+                                        let proofs: Vec<_> = proofs_iter.collect();
+                                        network_sender.broadcast_proof_of_store_msg_v2(proofs).await;
+                                    } else {
+                                        let proofs: Vec<_> = proofs_iter.map(|proof| {
+                                            let (info, sig) = proof.unpack();
+                                            ProofOfStore::new(info.info().clone(), sig)
+                                        }).collect();
+                                        network_sender.broadcast_proof_of_store_msg(proofs).await;
+                                    }
                                 } else {
+                                    let proofs: Vec<_> = proofs_iter.collect();
                                     network_sender.send_proof_of_store_msg_to_self(proofs).await;
                                 }
                             }
+                            if let Some(value) = self.batch_info_to_proof.get_mut(&info) {
+                                value.observe_voting_pct(approx_created_ts_usecs, &validator_verifier);
+                            }
                         },
                     }
                 }),
diff --git a/consensus/src/quorum_store/quorum_store_builder.rs b/consensus/src/quorum_store/quorum_store_builder.rs
@@ -351,6 +351,7 @@ impl InnerBuilder {
             self.proof_cache,
             self.broadcast_proofs,
             self.config.batch_expiry_gap_when_init_usecs,
+            self.config.enable_proof_v2,
         );
         spawn_named!(
             "proof_coordinator",
diff --git a/consensus/src/quorum_store/tests/batch_requester_test.rs b/consensus/src/quorum_store/tests/batch_requester_test.rs
@@ -11,7 +11,7 @@ use crate::{
 };
 use aptos_consensus_types::{
     common::Author,
-    proof_of_store::{BatchInfo, ProofOfStore, SignedBatchInfo},
+    proof_of_store::{BatchInfo, BatchInfoExt, ProofOfStore, SignedBatchInfo},
 };
 use aptos_crypto::HashValue;
 use aptos_infallible::Mutex;
@@ -73,9 +73,16 @@ impl QuorumStoreSender for MockBatchRequester {
         unimplemented!()
     }
 
+    async fn broadcast_proof_of_store_msg_v2(
+        &mut self,
+        _proof_of_stores: Vec<ProofOfStore<BatchInfoExt>>,
+    ) {
+        unimplemented!()
+    }
+
     async fn send_proof_of_store_msg_to_self(
         &mut self,
-        _proof_of_stores: Vec<ProofOfStore<BatchInfo>>,
+        _proof_of_stores: Vec<ProofOfStore<BatchInfoExt>>,
     ) {
         unimplemented!()
     }
diff --git a/consensus/src/quorum_store/tests/proof_coordinator_test.rs b/consensus/src/quorum_store/tests/proof_coordinator_test.rs
@@ -60,6 +60,7 @@ async fn test_proof_coordinator_basic() {
         proof_cache.clone(),
         true,
         10,
+        false,
     );
     let (proof_coordinator_tx, proof_coordinator_rx) = channel(100);
     let (tx, mut rx) = channel(100);
@@ -113,6 +114,7 @@ async fn test_proof_coordinator_with_unverified_signatures() {
         proof_cache.clone(),
         true,
         10,
+        false,
     );
     let (proof_coordinator_tx, proof_coordinator_rx) = channel(100);
     let (tx, mut rx) = channel(100);
diff --git a/consensus/src/test_utils/mock_quorum_store_sender.rs b/consensus/src/test_utils/mock_quorum_store_sender.rs
@@ -9,7 +9,7 @@ use crate::{
 use aptos_consensus_types::{
     common::Author,
     proof_of_store::{
-        BatchInfo, ProofOfStore, ProofOfStoreMsg, SignedBatchInfo, SignedBatchInfoMsg,
+        BatchInfo, BatchInfoExt, ProofOfStore, ProofOfStoreMsg, SignedBatchInfo, SignedBatchInfoMsg,
     },
 };
 use std::time::Duration;
@@ -72,8 +72,21 @@ impl QuorumStoreSender for MockQuorumStoreSender {
 
     async fn send_proof_of_store_msg_to_self(
         &mut self,
-        _proof_of_stores: Vec<ProofOfStore<BatchInfo>>,
+        _proof_of_stores: Vec<ProofOfStore<BatchInfoExt>>,
     ) {
         unimplemented!()
     }
+
+    async fn broadcast_proof_of_store_msg_v2(
+        &mut self,
+        proof_of_stores: Vec<ProofOfStore<BatchInfoExt>>,
+    ) {
+        self.tx
+            .send((
+                ConsensusMsg::ProofOfStoreMsgV2(Box::new(ProofOfStoreMsg::new(proof_of_stores))),
+                vec![],
+            ))
+            .await
+            .expect("We should be able to send the proof of store message");
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -99,6 +99,8 @@ pub struct QuorumStoreConfig {`
`99`	`99`	`pub enable_opt_quorum_store: bool,`
`100`	`100`	`pub opt_qs_minimum_batch_age_usecs: u64,`
`101`	`101`	`pub enable_payload_v2: bool,`
	`102`	+ /// Boolean flag that controls the usage of `BatchInfoExt::V1`
	`103`	`+ pub enable_proof_v2: bool,`
`102`	`104`	`}`
`103`	`105`
`104`	`106`	`impl Default for QuorumStoreConfig {`
`@@ -140,6 +142,7 @@ impl Default for QuorumStoreConfig {`
`140`	`142`	`enable_opt_quorum_store: true,`
`141`	`143`	`opt_qs_minimum_batch_age_usecs: Duration::from_millis(50).as_micros() as u64,`
`142`	`144`	`enable_payload_v2: false,`
	`145`	`+ enable_proof_v2: false,`
`143`	`146`	`}`
`144`	`147`	`}`
`145`	`148`	`}`