[dag] reliable broadcast storage (aptos-labs#9015)

Author: ibalajiarun

consensus/src/consensusdb/mod.rs

@@ -7,7 +7,7 @@ mod consensusdb_test;
 mod schema;
 
 use crate::{
-    dag::{CertifiedNode, Node},
+    dag::{CertifiedNode, Node, NodeDigestSignature, NodeId},
     error::DbError,
 };
 use anyhow::Result;
@@ -17,7 +17,7 @@ use aptos_logger::prelude::*;
 use aptos_schemadb::{Options, ReadOptions, SchemaBatch, DB, DEFAULT_COLUMN_FAMILY_NAME};
 use schema::{
     block::BlockSchema,
-    dag::{CertifiedNodeSchema, NodeSchema},
+    dag::{CertifiedNodeSchema, NodeDigestSignatureSchema, NodeSchema},
     quorum_certificate::QCSchema,
     single_entry::{SingleEntryKey, SingleEntrySchema},
     BLOCK_CF_NAME, CERTIFIED_NODE_CF_NAME, NODE_CF_NAME, QC_CF_NAME, SINGLE_ENTRY_CF_NAME,
@@ -199,6 +199,38 @@ impl ConsensusDB {
         Ok(())
     }
 
+    pub fn delete_node(&self, digest: HashValue) -> Result<(), DbError> {
+        let batch = SchemaBatch::new();
+        batch.delete::<NodeSchema>(&digest)?;
+        self.commit(batch)
+    }
+
+    pub fn save_node_signature(
+        &self,
+        node_id: &NodeId,
+        node_digest_signature: &NodeDigestSignature,
+    ) -> Result<(), DbError> {
+        let batch = SchemaBatch::new();
+        batch.put::<NodeDigestSignatureSchema>(node_id, node_digest_signature)?;
+        self.commit(batch)
+    }
+
+    pub fn get_node_signatures(&self) -> Result<HashMap<NodeId, NodeDigestSignature>, DbError> {
+        let mut iter = self
+            .db
+            .iter::<NodeDigestSignatureSchema>(ReadOptions::default())?;
+        iter.seek_to_first();
+        Ok(iter.collect::<Result<HashMap<NodeId, NodeDigestSignature>>>()?)
+    }
+
+    pub fn delete_node_signatures(&self, node_ids: Vec<NodeId>) -> Result<(), DbError> {
+        let batch = SchemaBatch::new();
+        node_ids
+            .iter()
+            .try_for_each(|node_id| batch.delete::<NodeDigestSignatureSchema>(node_id))?;
+        self.commit(batch)
+    }
+
     pub fn save_certified_node(&self, node: &CertifiedNode) -> Result<(), DbError> {
         let batch = SchemaBatch::new();
         batch.put::<CertifiedNodeSchema>(&node.digest(), node)?;

consensus/src/consensusdb/schema/dag/mod.rs

@@ -9,7 +9,7 @@
 //! |   digest   |   node/certified node    |
 //! ```
 
-use crate::dag::{CertifiedNode, Node};
+use crate::dag::{CertifiedNode, Node, NodeDigestSignature, NodeId};
 use anyhow::Result;
 use aptos_crypto::HashValue;
 use aptos_schemadb::{
@@ -42,6 +42,35 @@ impl ValueCodec<NodeSchema> for Node {
     }
 }
 
+pub const NODE_DIGEST_SIGNATURE_CF_NAME: ColumnFamilyName = "node_digest_signature";
+
+define_schema!(
+    NodeDigestSignatureSchema,
+    NodeId,
+    NodeDigestSignature,
+    NODE_DIGEST_SIGNATURE_CF_NAME
+);
+
+impl KeyCodec<NodeDigestSignatureSchema> for NodeId {
+    fn encode_key(&self) -> Result<Vec<u8>> {
+        Ok(bcs::to_bytes(&self)?)
+    }
+
+    fn decode_key(data: &[u8]) -> Result<Self> {
+        Ok(bcs::from_bytes(data)?)
+    }
+}
+
+impl ValueCodec<NodeDigestSignatureSchema> for NodeDigestSignature {
+    fn encode_value(&self) -> Result<Vec<u8>> {
+        Ok(bcs::to_bytes(&self)?)
+    }
+
+    fn decode_value(data: &[u8]) -> Result<Self> {
+        Ok(bcs::from_bytes(data)?)
+    }
+}
+
 pub const CERTIFIED_NODE_CF_NAME: ColumnFamilyName = "certified_node";
 
 define_schema!(

consensus/src/dag/dag_driver.rs

@@ -1,6 +1,7 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
+use super::storage::DAGStorage;
 use crate::{
     dag::{
         dag_store::Dag,
@@ -28,6 +29,7 @@ pub(crate) struct DagDriver {
     current_round: Round,
     time_service: Arc<dyn TimeService>,
     rb_abort_handle: Option<AbortHandle>,
+    storage: Arc<dyn DAGStorage>,
 }
 
 impl DagDriver {
@@ -39,7 +41,9 @@ impl DagDriver {
         reliable_broadcast: Arc<ReliableBroadcast>,
         current_round: Round,
         time_service: Arc<dyn TimeService>,
+        storage: Arc<dyn DAGStorage>,
     ) -> Self {
+        // TODO: rebroadcast nodes after recovery
         Self {
             author,
             epoch_state,
@@ -49,6 +53,7 @@ impl DagDriver {
             current_round,
             time_service,
             rb_abort_handle: None,
+            storage,
         }
     }
 
@@ -84,6 +89,9 @@ impl DagDriver {
             payload,
             strong_links,
         );
+        self.storage
+            .save_node(&new_node)
+            .expect("node must be saved");
         self.broadcast_node(new_node);
     }
 

consensus/src/dag/dag_handler.rs

@@ -1,6 +1,6 @@
 // Copyright © Aptos Foundation
 
-use super::{reliable_broadcast::CertifiedNodeHandler, types::TDAGMessage};
+use super::{reliable_broadcast::CertifiedNodeHandler, storage::DAGStorage, types::TDAGMessage};
 use crate::{
     dag::{
         dag_network::RpcHandler, dag_store::Dag, reliable_broadcast::NodeBroadcastHandler,
@@ -32,20 +32,23 @@ impl NetworkHandler {
         dag_rpc_rx: aptos_channel::Receiver<Author, IncomingDAGRequest>,
         signer: ValidatorSigner,
         epoch_state: Arc<EpochState>,
+        storage: Arc<dyn DAGStorage>,
     ) -> Self {
         Self {
             dag_rpc_rx,
             node_receiver: NodeBroadcastHandler::new(
                 dag.clone(),
                 signer,
-                epoch_state.verifier.clone(),
+                epoch_state.clone(),
+                storage,
             ),
             certified_node_receiver: CertifiedNodeHandler::new(dag),
             epoch_state,
         }
     }
 
     async fn start(mut self) {
+        // TODO(ibalajiarun): clean up Reliable Broadcast storage periodically.
         while let Some(msg) = self.dag_rpc_rx.next().await {
             if let Err(e) = self.process_rpc(msg).await {
                 warn!(error = ?e, "error processing rpc");

consensus/src/dag/mod.rs

@@ -16,4 +16,4 @@ mod tests;
 mod types;
 
 pub use dag_network::RpcHandler;
-pub use types::{CertifiedNode, DAGNetworkMessage, Node};
+pub use types::{CertifiedNode, DAGNetworkMessage, Node, NodeDigestSignature, NodeId};

consensus/src/dag/reliable_broadcast.rs

@@ -1,7 +1,11 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-use super::types::{CertifiedAck, CertifiedNode};
+use super::{
+    storage::DAGStorage,
+    types::{CertifiedAck, CertifiedNode},
+    NodeId,
+};
 use crate::{
     dag::{
         dag_network::{DAGNetworkSender, RpcHandler},
@@ -13,9 +17,10 @@ use crate::{
 use anyhow::{bail, ensure};
 use aptos_consensus_types::common::{Author, Round};
 use aptos_infallible::RwLock;
-use aptos_types::{validator_signer::ValidatorSigner, validator_verifier::ValidatorVerifier};
+use aptos_logger::error;
+use aptos_types::{epoch_state::EpochState, validator_signer::ValidatorSigner};
 use futures::{stream::FuturesUnordered, StreamExt};
-use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration};
+use std::{collections::BTreeMap, future::Future, mem, sync::Arc, time::Duration};
 use thiserror::Error as ThisError;
 
 pub trait BroadcastStatus {
@@ -96,25 +101,42 @@ pub struct NodeBroadcastHandler {
     dag: Arc<RwLock<Dag>>,
     signatures_by_round_peer: BTreeMap<Round, BTreeMap<Author, NodeDigestSignature>>,
     signer: ValidatorSigner,
-    verifier: ValidatorVerifier,
+    epoch_state: Arc<EpochState>,
+    storage: Arc<dyn DAGStorage>,
 }
 
 impl NodeBroadcastHandler {
     pub fn new(
         dag: Arc<RwLock<Dag>>,
         signer: ValidatorSigner,
-        verifier: ValidatorVerifier,
+        epoch_state: Arc<EpochState>,
+        storage: Arc<dyn DAGStorage>,
     ) -> Self {
+        let epoch = epoch_state.epoch;
+        let signatures_by_round_peer = read_signatures_from_storage(&storage, epoch);
+
         Self {
             dag,
-            signatures_by_round_peer: BTreeMap::new(),
+            signatures_by_round_peer,
             signer,
-            verifier,
+            epoch_state,
+            storage,
         }
     }
 
-    pub fn gc_before_round(&mut self, min_round: Round) {
-        self.signatures_by_round_peer.retain(|r, _| r >= &min_round);
+    pub fn gc_before_round(&mut self, min_round: Round) -> anyhow::Result<()> {
+        let to_retain = self.signatures_by_round_peer.split_off(&min_round);
+        let to_delete = mem::replace(&mut self.signatures_by_round_peer, to_retain);
+
+        let to_delete = to_delete
+            .iter()
+            .flat_map(|(r, peer_and_digest)| {
+                peer_and_digest
+                    .iter()
+                    .map(|(author, _)| NodeId::new(self.epoch_state.epoch, *r, *author))
+            })
+            .collect();
+        self.storage.delete_node_signatures(to_delete)
     }
 
     fn validate(&self, node: &Node) -> anyhow::Result<()> {
@@ -146,7 +168,8 @@ impl NodeBroadcastHandler {
             ensure!(
                 missing_parents.iter().all(|parent| {
                     let node_digest = NodeDigest::new(*parent.metadata().digest());
-                    self.verifier
+                    self.epoch_state
+                        .verifier
                         .verify_multi_signatures(&node_digest, parent.signatures())
                         .is_ok()
                 }),
@@ -160,6 +183,31 @@ impl NodeBroadcastHandler {
     }
 }
 
+fn read_signatures_from_storage(
+    storage: &Arc<dyn DAGStorage>,
+    epoch: u64,
+) -> BTreeMap<u64, BTreeMap<Author, NodeDigestSignature>> {
+    let mut signatures_by_round_peer = BTreeMap::new();
+
+    let all_node_signatures = storage.get_node_signatures().unwrap_or_default();
+    let mut to_delete = vec![];
+    for (node_id, node_sig) in all_node_signatures {
+        if node_id.epoch() == epoch {
+            signatures_by_round_peer
+                .entry(node_id.round())
+                .or_insert_with(BTreeMap::new)
+                .insert(node_id.author(), node_sig);
+        } else {
+            to_delete.push(node_id);
+        }
+    }
+    if let Err(err) = storage.delete_node_signatures(to_delete) {
+        error!("unable to clear old signatures: {}", err);
+    }
+
+    signatures_by_round_peer
+}
+
 impl RpcHandler for NodeBroadcastHandler {
     type Request = Node;
     type Response = NodeDigestSignature;
@@ -171,13 +219,16 @@ impl RpcHandler for NodeBroadcastHandler {
             .signatures_by_round_peer
             .entry(node.metadata().round())
             .or_insert(BTreeMap::new());
-        // TODO(ibalajiarun): persist node before voting
         match signatures_by_peer.get(node.metadata().author()) {
             None => {
                 let signature = node.sign(&self.signer)?;
                 let digest_signature =
                     NodeDigestSignature::new(node.metadata().epoch(), node.digest(), signature);
+
+                self.storage
+                    .save_node_signature(&NodeId::from(&node), &digest_signature)?;
                 signatures_by_peer.insert(*node.metadata().author(), digest_signature.clone());
+
                 Ok(digest_signature)
             },
             Some(ack) => Ok(ack.clone()),

consensus/src/dag/storage.rs

@@ -1,16 +1,30 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
+use super::{types::NodeDigestSignature, NodeId};
 use crate::{
     consensusdb::ConsensusDB,
     dag::{CertifiedNode, Node},
 };
+use anyhow::Ok;
 use aptos_crypto::HashValue;
 use std::collections::HashMap;
 
 pub trait DAGStorage {
     fn save_node(&self, node: &Node) -> anyhow::Result<()>;
 
+    fn delete_node(&self, digest: HashValue) -> anyhow::Result<()>;
+
+    fn save_node_signature(
+        &self,
+        node_id: &NodeId,
+        node_digest_signature: &NodeDigestSignature,
+    ) -> anyhow::Result<()>;
+
+    fn get_node_signatures(&self) -> anyhow::Result<HashMap<NodeId, NodeDigestSignature>>;
+
+    fn delete_node_signatures(&self, node_ids: Vec<NodeId>) -> anyhow::Result<()>;
+
     fn save_certified_node(&self, node: &CertifiedNode) -> anyhow::Result<()>;
 
     fn get_certified_nodes(&self) -> anyhow::Result<HashMap<HashValue, CertifiedNode>>;
@@ -23,6 +37,26 @@ impl DAGStorage for ConsensusDB {
         Ok(self.save_node(node)?)
     }
 
+    fn delete_node(&self, digest: HashValue) -> anyhow::Result<()> {
+        Ok(self.delete_node(digest)?)
+    }
+
+    fn save_node_signature(
+        &self,
+        node_id: &NodeId,
+        node_digest_signature: &NodeDigestSignature,
+    ) -> anyhow::Result<()> {
+        Ok(self.save_node_signature(node_id, node_digest_signature)?)
+    }
+
+    fn get_node_signatures(&self) -> anyhow::Result<HashMap<NodeId, NodeDigestSignature>> {
+        Ok(self.get_node_signatures()?)
+    }
+
+    fn delete_node_signatures(&self, node_ids: Vec<NodeId>) -> anyhow::Result<()> {
+        Ok(self.delete_node_signatures(node_ids)?)
+    }
+
     fn save_certified_node(&self, node: &CertifiedNode) -> anyhow::Result<()> {
         Ok(self.save_certified_node(node)?)
     }