randomness #4: RandManager update from randomnet (#12224)

* RandManager update from randomnet

* lint

* lint

consensus/Cargo.toml

@@ -23,6 +23,7 @@ aptos-consensus-notifications = { workspace = true }
 aptos-consensus-types = { workspace = true }
 aptos-crypto = { workspace = true }
 aptos-crypto-derive = { workspace = true }
+aptos-dkg = { workspace = true }
 aptos-enum-conversion-derive = { workspace = true }
 aptos-event-notifications = { workspace = true }
 aptos-executor = { workspace = true }
@@ -75,6 +76,7 @@ serde = { workspace = true }
 serde_bytes = { workspace = true }
 serde_json = { workspace = true }
 serde_yaml = { workspace = true }
+sha3 = { workspace = true }
 strum_macros = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
@@ -90,6 +92,7 @@ aptos-keygen = { workspace = true }
 aptos-mempool = { workspace = true, features = ["fuzzing"] }
 aptos-network = { workspace = true, features = ["fuzzing"] }
 aptos-safety-rules = { workspace = true, features = ["testing"] }
+aptos-vm = { workspace = true, features = ["fuzzing"] }
 aptos-vm-validator = { workspace = true }
 claims = { workspace = true }
 move-core-types = { workspace = true }

consensus/src/block_storage/tracing.rs

@@ -19,8 +19,7 @@ impl BlockStage {
     pub const QC_ADDED: &'static str = "qc_added";
     pub const QC_AGGREGATED: &'static str = "qc_aggregated";
     pub const RAND_ADD_DECISION: &'static str = "rand_add_decision";
-    pub const RAND_ADD_SHARE: &'static str = "rand_add_share";
-    pub const RAND_AGG_DECISION: &'static str = "rand_agg_decision";
+    pub const RAND_ADD_ENOUGH_SHARE: &'static str = "rand_add_enough_share";
     pub const RAND_ENTER: &'static str = "rand_enter";
     pub const RAND_READY: &'static str = "rand_ready";
     pub const ROUND_MANAGER_RECEIVED: &'static str = "round_manager_received";

consensus/src/counters.rs

@@ -868,6 +868,15 @@ pub static MAX_TXNS_FROM_BLOCK_TO_EXECUTE: Lazy<Histogram> = Lazy::new(|| {
     .unwrap()
 });
 
+/// Count of the number of `DKG` validator transactions received while the feature is disabled.
+pub static UNEXPECTED_DKG_VTXN_COUNT: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "aptos_consensus_unexpected_dkg_vtxn_count",
+        "Count of the number of `DKG` validator transactions received while the feature is disabled."
+    )
+        .unwrap()
+});
+
 /// Update various counters for committed blocks
 pub fn update_counters_for_committed_blocks(blocks_to_commit: &[Arc<PipelinedBlock>]) {
     for block in blocks_to_commit {
@@ -938,3 +947,11 @@ pub static PROPOSED_VTXN_BYTES: Lazy<IntCounterVec> = Lazy::new(|| {
     )
     .unwrap()
 });
+
+pub static RAND_QUEUE_SIZE: Lazy<IntGauge> = Lazy::new(|| {
+    register_int_gauge!(
+        "aptos_consensus_rand_queue_size",
+        "Number of randomness-pending blocks."
+    )
+    .unwrap()
+});

consensus/src/logging.rs

@@ -3,16 +3,19 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use aptos_consensus_types::common::Author;
+use aptos_crypto::HashValue;
 use aptos_logger::Schema;
 use aptos_types::block_info::Round;
 use serde::Serialize;
 
 #[derive(Schema)]
 pub struct LogSchema {
     event: LogEvent,
+    author: Option<Author>,
     remote_peer: Option<Author>,
     epoch: Option<u64>,
     round: Option<Round>,
+    id: Option<HashValue>,
 }
 
 #[derive(Serialize)]
@@ -40,15 +43,25 @@ pub enum LogEvent {
     Timeout,
     Vote,
     VoteNIL,
+    // log events related to randomness generation
+    BroadcastRandShare,
+    ReceiveProactiveRandShare,
+    ReceiveReactiveRandShare,
+    BroadcastAugData,
+    ReceiveAugData,
+    BroadcastCertifiedAugData,
+    ReceiveCertifiedAugData,
 }
 
 impl LogSchema {
     pub fn new(event: LogEvent) -> Self {
         Self {
             event,
+            author: None,
             remote_peer: None,
             epoch: None,
             round: None,
+            id: None,
         }
     }
 }

consensus/src/network.rs

@@ -14,7 +14,7 @@ use crate::{
     network_interface::{ConsensusMsg, ConsensusNetworkClient, RPC},
     pipeline::commit_reliable_broadcast::CommitMessage,
     quorum_store::types::{Batch, BatchMsg, BatchRequest, BatchResponse},
-    rand::rand_gen::RandGenMessage,
+    rand::rand_gen::network_messages::RandGenMessage,
 };
 use anyhow::{anyhow, bail, ensure};
 use aptos_channels::{self, aptos_channel, message_queues::QueueStyle};

consensus/src/network_interface.rs

@@ -8,7 +8,7 @@ use crate::{
     dag::DAGNetworkMessage,
     pipeline,
     quorum_store::types::{Batch, BatchMsg, BatchRequest, BatchResponse},
-    rand::rand_gen::RandGenMessage,
+    rand::rand_gen::network_messages::RandGenMessage,
 };
 use aptos_config::network_id::{NetworkId, PeerNetworkId};
 use aptos_consensus_types::{

consensus/src/rand/rand_gen/aug_data_store.rs

@@ -2,10 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::rand::rand_gen::{
-    storage::interface::AugDataStorage,
+    storage::interface::RandStorage,
     types::{
-        AugData, AugDataId, AugDataSignature, AugmentedData, CertifiedAugData, CertifiedAugDataAck,
-        RandConfig,
+        AugData, AugDataId, AugDataSignature, CertifiedAugData, CertifiedAugDataAck, RandConfig,
+        TAugmentedData,
     },
 };
 use anyhow::ensure;
@@ -14,16 +14,16 @@ use aptos_logger::error;
 use aptos_types::validator_signer::ValidatorSigner;
 use std::{collections::HashMap, sync::Arc};
 
-pub struct AugDataStore<D, Storage> {
+pub struct AugDataStore<D> {
     epoch: u64,
     signer: Arc<ValidatorSigner>,
     config: RandConfig,
     data: HashMap<Author, AugData<D>>,
     certified_data: HashMap<Author, CertifiedAugData<D>>,
-    db: Arc<Storage>,
+    db: Arc<dyn RandStorage<D>>,
 }
 
-impl<D: AugmentedData, Storage: AugDataStorage<D>> AugDataStore<D, Storage> {
+impl<D: TAugmentedData> AugDataStore<D> {
     fn filter_by_epoch<T>(
         epoch: u64,
         all_data: impl Iterator<Item = (AugDataId, T)>,
@@ -44,24 +44,30 @@ impl<D: AugmentedData, Storage: AugDataStorage<D>> AugDataStore<D, Storage> {
         epoch: u64,
         signer: Arc<ValidatorSigner>,
         config: RandConfig,
-        db: Arc<Storage>,
+        db: Arc<dyn RandStorage<D>>,
     ) -> Self {
         let all_data = db.get_all_aug_data().unwrap_or_default();
         let (to_remove, aug_data) = Self::filter_by_epoch(epoch, all_data.into_iter());
-        if let Err(e) = db.remove_aug_data(to_remove.into_iter()) {
+        if let Err(e) = db.remove_aug_data(to_remove) {
             error!("[AugDataStore] failed to remove aug data: {:?}", e);
         }
 
         let all_certified_data = db.get_all_certified_aug_data().unwrap_or_default();
         let (to_remove, certified_data) =
             Self::filter_by_epoch(epoch, all_certified_data.into_iter());
-        if let Err(e) = db.remove_certified_aug_data(to_remove.into_iter()) {
+        if let Err(e) = db.remove_certified_aug_data(to_remove) {
             error!(
                 "[AugDataStore] failed to remove certified aug data: {:?}",
                 e
             );
         }
 
+        for (_, certified_data) in &certified_data {
+            certified_data
+                .data()
+                .augment(&config, certified_data.author());
+        }
+
         Self {
             epoch,
             signer,
@@ -79,11 +85,11 @@ impl<D: AugmentedData, Storage: AugDataStorage<D>> AugDataStore<D, Storage> {
     }
 
     pub fn get_my_aug_data(&self) -> Option<AugData<D>> {
-        self.data.get(self.config.author()).cloned()
+        self.data.get(&self.config.author()).cloned()
     }
 
     pub fn get_my_certified_aug_data(&self) -> Option<CertifiedAugData<D>> {
-        self.certified_data.get(self.config.author()).cloned()
+        self.certified_data.get(&self.config.author()).cloned()
     }
 
     pub fn add_aug_data(&mut self, data: AugData<D>) -> anyhow::Result<AugDataSignature> {

consensus/src/rand/rand_gen/block_queue.rs

@@ -94,6 +94,10 @@ impl BlockQueue {
         }
     }
 
+    pub fn queue(&self) -> &BTreeMap<Round, QueueItem> {
+        &self.queue
+    }
+
     pub fn push_back(&mut self, item: QueueItem) {
         for block in item.blocks() {
             observe_block(block.timestamp_usecs(), BlockStage::RAND_ENTER);

consensus/src/rand/rand_gen/mod.rs

@@ -4,14 +4,12 @@
 #[cfg(test)]
 mod test_utils;
 
-mod block_queue;
-mod network_messages;
-mod rand_store;
-mod types;
-
-mod aug_data_store;
-mod rand_manager;
-mod reliable_broadcast_state;
-mod storage;
-
-pub use network_messages::RandGenMessage;
+pub mod block_queue;
+pub mod network_messages;
+pub mod rand_store;
+pub mod types;
+
+pub mod aug_data_store;
+pub mod rand_manager;
+pub mod reliable_broadcast_state;
+pub mod storage;

consensus/src/rand/rand_gen/network_messages.rs

@@ -5,8 +5,8 @@ use crate::{
     network::TConsensusMsg,
     network_interface::ConsensusMsg,
     rand::rand_gen::types::{
-        AugData, AugDataSignature, AugmentedData, CertifiedAugData, CertifiedAugDataAck,
-        RandConfig, RandShare, RequestShare, Share,
+        AugData, AugDataSignature, CertifiedAugData, CertifiedAugDataAck, RandConfig, RandShare,
+        RequestShare, TAugmentedData, TShare,
     },
 };
 use anyhow::bail;
@@ -30,7 +30,7 @@ pub enum RandMessage<S, D> {
     CertifiedAugDataAck(CertifiedAugDataAck),
 }
 
-impl<S: Share, D: AugmentedData> RandMessage<S, D> {
+impl<S: TShare, D: TAugmentedData> RandMessage<S, D> {
     pub fn verify(
         &self,
         epoch_state: &EpochState,
@@ -49,9 +49,9 @@ impl<S: Share, D: AugmentedData> RandMessage<S, D> {
     }
 }
 
-impl<S: Share, D: AugmentedData> RBMessage for RandMessage<S, D> {}
+impl<S: TShare, D: TAugmentedData> RBMessage for RandMessage<S, D> {}
 
-impl<S: Share, D: AugmentedData> TConsensusMsg for RandMessage<S, D> {
+impl<S: TShare, D: TAugmentedData> TConsensusMsg for RandMessage<S, D> {
     fn epoch(&self) -> u64 {
         match self {
             RandMessage::RequestShare(request) => request.epoch(),