[qs] support signing BatchInfoExt behind flag

ibalajiarun · ibalajiarun · commit 75caee4095c3 · 2025-11-07T12:36:20.000-08:00
diff --git a/consensus/src/network.rs b/consensus/src/network.rs
@@ -208,6 +208,12 @@ pub trait QuorumStoreSender: Send + Clone {
         recipients: Vec<Author>,
     );
 
+    async fn send_signed_batch_info_msg_v2(
+        &self,
+        signed_batch_infos: Vec<SignedBatchInfo<BatchInfoExt>>,
+        recipients: Vec<Author>,
+    );
+
     async fn broadcast_batch_msg(&mut self, batches: Vec<Batch>);
 
     async fn broadcast_proof_of_store_msg(&mut self, proof_of_stores: Vec<ProofOfStore<BatchInfo>>);
@@ -575,6 +581,18 @@ impl QuorumStoreSender for NetworkSender {
         self.send(msg, recipients).await
     }
 
+    async fn send_signed_batch_info_msg_v2(
+        &self,
+        signed_batch_infos: Vec<SignedBatchInfo<BatchInfoExt>>,
+        recipients: Vec<Author>,
+    ) {
+        fail_point!("consensus::send::signed_batch_info", |_| ());
+        let msg = ConsensusMsg::SignedBatchInfoMsgV2(Box::new(SignedBatchInfoMsg::new(
+            signed_batch_infos,
+        )));
+        self.send(msg, recipients).await
+    }
+
     async fn broadcast_batch_msg(&mut self, batches: Vec<Batch>) {
         fail_point!("consensus::send::broadcast_batch", |_| ());
         let msg = ConsensusMsg::BatchMsg(Box::new(BatchMsg::new(batches)));
diff --git a/consensus/src/quorum_store/batch_coordinator.rs b/consensus/src/quorum_store/batch_coordinator.rs
@@ -44,6 +44,7 @@ pub struct BatchCoordinator {
     max_total_bytes: u64,
     batch_expiry_gap_when_init_usecs: u64,
     transaction_filter_config: BatchTransactionFilterConfig,
+    enable_proof_v2: bool,
 }
 
 impl BatchCoordinator {
@@ -59,6 +60,7 @@ impl BatchCoordinator {
         max_total_bytes: u64,
         batch_expiry_gap_when_init_usecs: u64,
         transaction_filter_config: BatchTransactionFilterConfig,
+        enable_proof_v2: bool,
     ) -> Self {
         Self {
             my_peer_id,
@@ -72,6 +74,7 @@ impl BatchCoordinator {
             max_total_bytes,
             batch_expiry_gap_when_init_usecs,
             transaction_filter_config,
+            enable_proof_v2,
         }
     }
 
@@ -87,6 +90,7 @@ impl BatchCoordinator {
         let batch_store = self.batch_store.clone();
         let network_sender = self.network_sender.clone();
         let sender_to_proof_manager = self.sender_to_proof_manager.clone();
+        let enable_proof_v2 = self.enable_proof_v2;
         tokio::spawn(async move {
             let peer_id = persist_requests[0].author();
             let batches = persist_requests
@@ -98,14 +102,27 @@ impl BatchCoordinator {
                     )
                 })
                 .collect();
-            let signed_batch_infos = batch_store.persist(persist_requests);
-            if !signed_batch_infos.is_empty() {
-                if approx_created_ts_usecs > 0 {
-                    observe_batch(approx_created_ts_usecs, peer_id, BatchStage::SIGNED);
+
+            if enable_proof_v2 {
+                let signed_batch_infos = batch_store.persist_v2(persist_requests);
+                if !signed_batch_infos.is_empty() {
+                    if approx_created_ts_usecs > 0 {
+                        observe_batch(approx_created_ts_usecs, peer_id, BatchStage::SIGNED);
+                    }
+                    network_sender
+                        .send_signed_batch_info_msg_v2(signed_batch_infos, vec![peer_id])
+                        .await;
+                }
+            } else {
+                let signed_batch_infos = batch_store.persist(persist_requests);
+                if !signed_batch_infos.is_empty() {
+                    if approx_created_ts_usecs > 0 {
+                        observe_batch(approx_created_ts_usecs, peer_id, BatchStage::SIGNED);
+                    }
+                    network_sender
+                        .send_signed_batch_info_msg(signed_batch_infos, vec![peer_id])
+                        .await;
                 }
-                network_sender
-                    .send_signed_batch_info_msg(signed_batch_infos, vec![peer_id])
-                    .await;
             }
             let _ = sender_to_proof_manager
                 .send(ProofManagerCommand::ReceiveBatches(batches))
diff --git a/consensus/src/quorum_store/batch_store.rs b/consensus/src/quorum_store/batch_store.rs
@@ -12,7 +12,9 @@ use crate::{
     },
 };
 use anyhow::bail;
-use aptos_consensus_types::proof_of_store::{BatchInfo, BatchInfoExt, SignedBatchInfo};
+use aptos_consensus_types::proof_of_store::{
+    BatchInfo, BatchInfoExt, BatchKind, ExtraBatchInfo, SignedBatchInfo, TBatchInfo,
+};
 use aptos_crypto::{CryptoMaterialError, HashValue};
 use aptos_executor_types::{ExecutorError, ExecutorResult};
 use aptos_infallible::Mutex;
@@ -26,6 +28,7 @@ use fail::fail_point;
 use futures::{future::Shared, FutureExt};
 use once_cell::sync::OnceCell;
 use std::{
+    any::TypeId,
     collections::{BTreeSet, HashMap},
     future::Future,
     pin::Pin,
@@ -378,10 +381,10 @@ impl BatchStore {
         ret
     }
 
-    fn generate_signed_batch_info(
+    fn generate_signed_batch_info<T: TBatchInfo>(
         &self,
-        batch_info: BatchInfo,
-    ) -> Result<SignedBatchInfo<BatchInfo>, CryptoMaterialError> {
+        batch_info: T,
+    ) -> Result<SignedBatchInfo<T>, CryptoMaterialError> {
         fail_point!("quorum_store::create_invalid_signed_batch_info", |_| {
             Ok(SignedBatchInfo::new_with_signature(
                 batch_info.clone(),
@@ -392,10 +395,17 @@ impl BatchStore {
         SignedBatchInfo::new(batch_info, &self.validator_signer)
     }
 
-    fn persist_inner(&self, persist_request: PersistedValue) -> Option<SignedBatchInfo<BatchInfo>> {
+    fn persist_inner<T: TBatchInfo>(
+        &self,
+        batch_info: T,
+        persist_request: PersistedValue,
+    ) -> Option<SignedBatchInfo<T>> {
+        assert!(
+            batch_info.as_batch_info() == persist_request.batch_info(),
+            "Provided batch info doesn't match persist request batch info"
+        );
         match self.save(&persist_request) {
             Ok(needs_db) => {
-                let batch_info = persist_request.batch_info().clone();
                 trace!("QS: sign digest {}", persist_request.digest());
                 if needs_db {
                     #[allow(clippy::unwrap_in_result)]
@@ -405,7 +415,6 @@ impl BatchStore {
                 }
                 self.generate_signed_batch_info(batch_info).ok()
             },
-
             Err(e) => {
                 debug!("QS: failed to store to cache {:?}", e);
                 None
@@ -486,7 +495,37 @@ impl BatchWriter for BatchStore {
     fn persist(&self, persist_requests: Vec<PersistedValue>) -> Vec<SignedBatchInfo<BatchInfo>> {
         let mut signed_infos = vec![];
         for persist_request in persist_requests.into_iter() {
-            if let Some(signed_info) = self.persist_inner(persist_request.clone()) {
+            let batch_info = persist_request.batch_info().clone();
+            if let Some(signed_info) = self.persist_inner(batch_info, persist_request.clone()) {
+                self.notify_subscribers(persist_request);
+                signed_infos.push(signed_info);
+            }
+        }
+        signed_infos
+    }
+
+    fn persist_v2(
+        &self,
+        persist_requests: Vec<PersistedValue>,
+    ) -> Vec<SignedBatchInfo<BatchInfoExt>> {
+        let mut signed_infos = vec![];
+        for persist_request in persist_requests.into_iter() {
+            let is_encrypted_batch = persist_request
+                .payload()
+                .as_ref()
+                .expect("Payload must be available for persistence")
+                .iter()
+                .any(|txn| txn.is_encrypted());
+            let batch_kind = if is_encrypted_batch {
+                BatchKind::Encrypted
+            } else {
+                BatchKind::Normal
+            };
+            let batch_info = BatchInfoExt::V2 {
+                info: persist_request.batch_info().clone(),
+                extra: ExtraBatchInfo { batch_kind },
+            };
+            if let Some(signed_info) = self.persist_inner(batch_info, persist_request.clone()) {
                 self.notify_subscribers(persist_request);
                 signed_infos.push(signed_info);
             }
@@ -611,4 +650,9 @@ impl<T: QuorumStoreSender + Clone + Send + Sync + 'static> BatchReader for Batch
 
 pub trait BatchWriter: Send + Sync {
     fn persist(&self, persist_requests: Vec<PersistedValue>) -> Vec<SignedBatchInfo<BatchInfo>>;
+
+    fn persist_v2(
+        &self,
+        persist_requests: Vec<PersistedValue>,
+    ) -> Vec<SignedBatchInfo<BatchInfoExt>>;
 }
diff --git a/consensus/src/quorum_store/quorum_store_builder.rs b/consensus/src/quorum_store/quorum_store_builder.rs
@@ -333,6 +333,7 @@ impl InnerBuilder {
                 self.config.receiver_max_total_bytes as u64,
                 self.config.batch_expiry_gap_when_init_usecs,
                 self.transaction_filter_config.clone(),
+                self.config.enable_proof_v2,
             );
             #[allow(unused_variables)]
             let name = format!("batch_coordinator-{}", i);
diff --git a/consensus/src/quorum_store/tests/batch_coordinator_test.rs b/consensus/src/quorum_store/tests/batch_coordinator_test.rs
@@ -142,6 +142,7 @@ fn create_batch_coordinator(
         10_000,
         10_000,
         transaction_filter_config,
+        false,
     )
 }
 
diff --git a/consensus/src/quorum_store/tests/batch_generator_test.rs b/consensus/src/quorum_store/tests/batch_generator_test.rs
@@ -38,6 +38,13 @@ impl BatchWriter for MockBatchWriter {
     fn persist(&self, _persist_requests: Vec<PersistedValue>) -> Vec<SignedBatchInfo<BatchInfo>> {
         vec![]
     }
+
+    fn persist_v2(
+        &self,
+        persist_requests: Vec<PersistedValue>,
+    ) -> Vec<SignedBatchInfo<aptos_consensus_types::proof_of_store::BatchInfoExt>> {
+        vec![]
+    }
 }
 
 #[allow(clippy::needless_collect)]
diff --git a/consensus/src/quorum_store/tests/batch_requester_test.rs b/consensus/src/quorum_store/tests/batch_requester_test.rs
@@ -62,6 +62,14 @@ impl QuorumStoreSender for MockBatchRequester {
         unimplemented!()
     }
 
+    async fn send_signed_batch_info_msg_v2(
+        &self,
+        _signed_batch_infos: Vec<SignedBatchInfo<BatchInfoExt>>,
+        _recipients: Vec<Author>,
+    ) {
+        unimplemented!()
+    }
+
     async fn broadcast_batch_msg(&mut self, _batches: Vec<Batch>) {
         unimplemented!()
     }
diff --git a/consensus/src/test_utils/mock_quorum_store_sender.rs b/consensus/src/test_utils/mock_quorum_store_sender.rs
@@ -53,6 +53,22 @@ impl QuorumStoreSender for MockQuorumStoreSender {
             .expect("could not send");
     }
 
+    async fn send_signed_batch_info_msg_v2(
+        &self,
+        signed_batch_infos: Vec<SignedBatchInfo<BatchInfoExt>>,
+        recipients: Vec<Author>,
+    ) {
+        self.tx
+            .send((
+                ConsensusMsg::SignedBatchInfoMsgV2(Box::new(SignedBatchInfoMsg::new(
+                    signed_batch_infos,
+                ))),
+                recipients,
+            ))
+            .await
+            .expect("could not send");
+    }
+
     async fn broadcast_batch_msg(&mut self, _batches: Vec<Batch>) {
         unimplemented!()
     }
diff --git a/types/src/transaction/mod.rs b/types/src/transaction/mod.rs
@@ -940,6 +940,10 @@ impl TransactionPayload {
             extra_config,
         })
     }
+
+    pub fn is_encrypted(&self) -> bool {
+        matches!(self, Self::EncryptedPayload(_))
+    }
 }
 
 impl TransactionExtraConfig {
@@ -1296,6 +1300,10 @@ impl SignedTransaction {
     pub fn replay_protector(&self) -> ReplayProtector {
         self.raw_txn.replay_protector()
     }
+
+    pub fn is_encrypted(&self) -> bool {
+        self.payload().is_encrypted()
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]

Original file line number	Diff line number	Diff line change
`@@ -142,6 +142,7 @@ fn create_batch_coordinator(`
`142`	`142`	`10_000,`
`143`	`143`	`10_000,`
`144`	`144`	`transaction_filter_config,`
	`145`	`+ false,`
`145`	`146`	`)`
`146`	`147`	`}`
`147`	`148`
Original file line number	Diff line number	Diff line change
`@@ -38,6 +38,13 @@ impl BatchWriter for MockBatchWriter {`
`38`	`38`	`fn persist(&self, _persist_requests: Vec<PersistedValue>) -> Vec<SignedBatchInfo<BatchInfo>> {`
`39`	`39`	`vec![]`
`40`	`40`	`}`
	`41`	`+`
	`42`	`+ fn persist_v2(`
	`43`	`+ &self,`
	`44`	`+ persist_requests: Vec<PersistedValue>,`
	`45`	`+ ) -> Vec<SignedBatchInfo<aptos_consensus_types::proof_of_store::BatchInfoExt>> {`
	`46`	`+ vec![]`
	`47`	`+ }`
`41`	`48`	`}`
`42`	`49`
`43`	`50`	`#[allow(clippy::needless_collect)]`
Original file line number	Diff line number	Diff line change
`@@ -62,6 +62,14 @@ impl QuorumStoreSender for MockBatchRequester {`
`62`	`62`	`unimplemented!()`
`63`	`63`	`}`
`64`	`64`
	`65`	`+ async fn send_signed_batch_info_msg_v2(`
	`66`	`+ &self,`
	`67`	`+ _signed_batch_infos: Vec<SignedBatchInfo<BatchInfoExt>>,`
	`68`	`+ _recipients: Vec<Author>,`
	`69`	`+ ) {`
	`70`	`+ unimplemented!()`
	`71`	`+ }`
	`72`	`+`
`65`	`73`	`async fn broadcast_batch_msg(&mut self, _batches: Vec<Batch>) {`
`66`	`74`	`unimplemented!()`
`67`	`75`	`}`
Original file line number	Diff line number	Diff line change
`@@ -940,6 +940,10 @@ impl TransactionPayload {`
`940`	`940`	`extra_config,`
`941`	`941`	`})`
`942`	`942`	`}`
	`943`	`+`
	`944`	`+ pub fn is_encrypted(&self) -> bool {`
	`945`	`+ matches!(self, Self::EncryptedPayload(_))`
	`946`	`+ }`
`943`	`947`	`}`
`944`	`948`
`945`	`949`	`impl TransactionExtraConfig {`
`@@ -1296,6 +1300,10 @@ impl SignedTransaction {`
`1296`	`1300`	`pub fn replay_protector(&self) -> ReplayProtector {`
`1297`	`1301`	`self.raw_txn.replay_protector()`
`1298`	`1302`	`}`
	`1303`	`+`
	`1304`	`+ pub fn is_encrypted(&self) -> bool {`
	`1305`	`+ self.payload().is_encrypted()`
	`1306`	`+ }`
`1299`	`1307`	`}`
`1300`	`1308`
`1301`	`1309`	`#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]`