[DRB] - Start the DRB calculation two epochs in advance (#3911)

* Implement DRB computation logic

* Fix epoch check, replace initial values

* Use decided qc for computation

crates/hotshot/src/tasks/task_state.rs

@@ -231,6 +231,7 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> CreateTaskState
             vote_dependencies: BTreeMap::new(),
             network: Arc::clone(&handle.hotshot.network),
             membership: (*handle.hotshot.memberships).clone().into(),
+            drb_computations: BTreeMap::new(),
             output_event_stream: handle.hotshot.external_event_stream.0.clone(),
             id: handle.hotshot.id,
             storage: Arc::clone(&handle.storage),

crates/hotshot/src/traits/election.rs

@@ -6,8 +6,6 @@
 
 //! elections used for consensus
 
-/// Dynamic leader election with epochs.
-pub mod dynamic;
 /// leader completely randomized every view
 pub mod randomized_committee;
 /// static (round robin) committee election

crates/task-impls/src/quorum_vote/handlers.rs

@@ -13,18 +13,22 @@ use committable::Committable;
 use hotshot_types::{
     consensus::OuterConsensus,
     data::{Leaf2, QuorumProposal2, VidDisperseShare},
+    drb::compute_drb_result,
     event::{Event, EventType},
     message::{Proposal, UpgradeLock},
     simple_vote::{QuorumData2, QuorumVote2},
     traits::{
+        block_contents::BlockHeader,
         election::Membership,
         node_implementation::{ConsensusTime, NodeImplementation, NodeType},
         signature_key::SignatureKey,
         storage::Storage,
         ValidatedState,
     },
+    utils::epoch_from_block_number,
     vote::HasViewNumber,
 };
+use tokio::spawn;
 use tracing::instrument;
 use utils::anytrace::*;
 use vbs::version::StaticVersionType;
@@ -135,13 +139,13 @@ pub(crate) async fn handle_quorum_proposal_validated<
         // We don't need to hold this while we broadcast
         drop(consensus_writer);
 
-        // First, send an update to everyone saying that we've reached a decide
+        // Send an update to everyone saying that we've reached a decide
         broadcast_event(
             Event {
                 view_number: decided_view_number,
                 event: EventType::Decide {
-                    leaf_chain: Arc::new(leaf_views),
-                    // This is never *not* none if we've reached a new decide, so this is safe to unwrap.
+                    leaf_chain: Arc::new(leaf_views.clone()),
+                    // This is never none if we've reached a new decide, so this is safe to unwrap.
                     qc: Arc::new(new_decide_qc.unwrap()),
                     block_size: included_txns.map(|txns| txns.len().try_into().unwrap()),
                 },
@@ -150,6 +154,65 @@ pub(crate) async fn handle_quorum_proposal_validated<
         )
         .await;
         tracing::debug!("Successfully sent decide event");
+
+        // Start the DRB computation two epochs in advance, if the decided block is the last but
+        // third block in the current epoch and we are in the quorum committee of the next epoch.
+        //
+        // Special cases:
+        // * Epoch 0: No DRB computation since we'll transition to epoch 1 immediately.
+        // * Epoch 1 and 2: Use `[0u8; 32]` as the DRB result since when we first start the
+        // computation in epoch 1, the result is for epoch 3.
+        //
+        // We don't need to handle the special cases explicitly here, because the first proposal
+        // with which we'll start the DRB computation is for epoch 3.
+        if version >= V::Epochs::VERSION {
+            // This is never none if we've reached a new decide, so this is safe to unwrap.
+            let decided_block_number = leaf_views
+                .last()
+                .unwrap()
+                .leaf
+                .block_header()
+                .block_number();
+
+            // Skip if this is not the expected block.
+            if task_state.epoch_height != 0
+                && (decided_block_number + 3) % task_state.epoch_height == 0
+            {
+                // Cancel old DRB computation tasks.
+                let current_epoch_number = TYPES::Epoch::new(epoch_from_block_number(
+                    decided_block_number,
+                    task_state.epoch_height,
+                ));
+                let current_tasks = task_state.drb_computations.split_off(&current_epoch_number);
+                while let Some((_, task)) = task_state.drb_computations.pop_last() {
+                    task.abort();
+                }
+                task_state.drb_computations = current_tasks;
+
+                // Skip if we are not in the committee of the next epoch.
+                if task_state
+                    .membership
+                    .has_stake(&task_state.public_key, current_epoch_number + 1)
+                {
+                    let new_epoch_number = current_epoch_number + 2;
+                    let Ok(drb_seed_input_vec) =
+                        bincode::serialize(&proposal.justify_qc.signatures)
+                    else {
+                        bail!("Failed to serialize the QC signature.");
+                    };
+                    let Ok(drb_seed_input) = drb_seed_input_vec.try_into() else {
+                        bail!(
+                            "Failed to convert the serialized QC signature into a DRB seed input."
+                        );
+                    };
+                    let new_drb_task =
+                        spawn(async move { compute_drb_result::<TYPES>(drb_seed_input) });
+                    task_state
+                        .drb_computations
+                        .insert(new_epoch_number, new_drb_task);
+                }
+            }
+        }
     }
 
     Ok(())

crates/task-impls/src/quorum_vote/mod.rs

@@ -18,6 +18,7 @@ use hotshot_task::{
 use hotshot_types::{
     consensus::OuterConsensus,
     data::{Leaf2, QuorumProposal2},
+    drb::DrbResult,
     event::Event,
     message::{Proposal, UpgradeLock},
     traits::{
@@ -272,6 +273,9 @@ pub struct QuorumVoteTaskState<TYPES: NodeType, I: NodeImplementation<TYPES>, V:
     /// Membership for Quorum certs/votes and DA committee certs/votes.
     pub membership: Arc<TYPES::Membership>,
 
+    /// Table for the in-progress DRB computation tasks.
+    pub drb_computations: BTreeMap<TYPES::Epoch, JoinHandle<DrbResult>>,
+
     /// Output events to application
     pub output_event_stream: async_broadcast::Sender<Event<TYPES>>,
 

crates/testing/src/view_generator.rs

@@ -25,6 +25,7 @@ use hotshot_types::{
         DaProposal, EpochNumber, Leaf, Leaf2, QuorumProposal2, VidDisperse, VidDisperseShare,
         ViewChangeEvidence, ViewNumber,
     },
+    drb::{INITIAL_DRB_RESULT, INITIAL_DRB_SEED_INPUT},
     message::{Proposal, UpgradeLock},
     simple_certificate::{
         DaCertificate, QuorumCertificate, QuorumCertificate2, TimeoutCertificate,
@@ -141,8 +142,8 @@ impl TestView {
             .to_qc2(),
             upgrade_certificate: None,
             view_change_evidence: None,
-            drb_result: [0; 32],
-            drb_seed: [0; 96],
+            drb_result: INITIAL_DRB_RESULT,
+            drb_seed: INITIAL_DRB_SEED_INPUT,
         };
 
         let encoded_transactions = Arc::from(TestTransaction::encode(&transactions));
@@ -368,8 +369,8 @@ impl TestView {
             justify_qc: quorum_certificate.clone(),
             upgrade_certificate: upgrade_certificate.clone(),
             view_change_evidence,
-            drb_result: [0; 32],
-            drb_seed: [0; 96],
+            drb_result: INITIAL_DRB_RESULT,
+            drb_seed: INITIAL_DRB_SEED_INPUT,
         };
 
         let mut leaf = Leaf2::from_quorum_proposal(&proposal);

crates/testing/tests/tests_1/message.rs

@@ -66,8 +66,7 @@ async fn test_certificate2_validity() {
     use hotshot_testing::{helpers::build_system_handle, view_generator::TestViewGenerator};
     use hotshot_types::{
         data::{EpochNumber, Leaf, Leaf2},
-        traits::election::Membership,
-        traits::node_implementation::ConsensusTime,
+        traits::{election::Membership, node_implementation::ConsensusTime},
         vote::Certificate,
     };
 

crates/types/src/data.rs

@@ -30,6 +30,7 @@ use utils::anytrace::*;
 use vec1::Vec1;
 
 use crate::{
+    drb::{DrbResult, DrbSeedInput, INITIAL_DRB_RESULT, INITIAL_DRB_SEED_INPUT},
     message::{Proposal, UpgradeLock},
     simple_certificate::{
         QuorumCertificate, QuorumCertificate2, TimeoutCertificate, UpgradeCertificate,
@@ -391,13 +392,17 @@ pub struct QuorumProposal2<TYPES: NodeType> {
     /// Possible timeout or view sync certificate. If the `justify_qc` is not for a proposal in the immediately preceding view, then either a timeout or view sync certificate must be attached.
     pub view_change_evidence: Option<ViewChangeEvidence<TYPES>>,
 
-    /// the DRB seed currently being calculated
+    /// The DRB seed for the next epoch.
+    ///
+    /// The DRB computation using this seed was started in the previous epoch.
     #[serde(with = "serde_bytes")]
-    pub drb_seed: [u8; 96],
+    pub drb_seed: DrbSeedInput,
 
-    /// the result of the DRB calculation
+    /// The DRB result for the current epoch.
+    ///
+    /// The DRB computation with this result was started two epochs ago.
     #[serde(with = "serde_bytes")]
-    pub drb_result: [u8; 32],
+    pub drb_result: DrbResult,
 }
 
 impl<TYPES: NodeType> From<QuorumProposal<TYPES>> for QuorumProposal2<TYPES> {
@@ -408,8 +413,8 @@ impl<TYPES: NodeType> From<QuorumProposal<TYPES>> for QuorumProposal2<TYPES> {
             justify_qc: quorum_proposal.justify_qc.to_qc2(),
             upgrade_certificate: quorum_proposal.upgrade_certificate,
             view_change_evidence: quorum_proposal.proposal_certificate,
-            drb_seed: [0; 96],
-            drb_result: [0; 32],
+            drb_seed: INITIAL_DRB_SEED_INPUT,
+            drb_result: INITIAL_DRB_RESULT,
         }
     }
 }
@@ -438,8 +443,8 @@ impl<TYPES: NodeType> From<Leaf<TYPES>> for Leaf2<TYPES> {
             upgrade_certificate: leaf.upgrade_certificate,
             block_payload: leaf.block_payload,
             view_change_evidence: None,
-            drb_seed: [0; 96],
-            drb_result: [0; 32],
+            drb_seed: INITIAL_DRB_SEED_INPUT,
+            drb_result: INITIAL_DRB_RESULT,
         }
     }
 }
@@ -562,13 +567,17 @@ pub struct Leaf2<TYPES: NodeType> {
     /// Possible timeout or view sync certificate. If the `justify_qc` is not for a proposal in the immediately preceding view, then either a timeout or view sync certificate must be attached.
     pub view_change_evidence: Option<ViewChangeEvidence<TYPES>>,
 
-    /// the DRB seed currently being calculated
+    /// The DRB seed for the next epoch.
+    ///
+    /// The DRB computation using this seed was started in the previous epoch.
     #[serde(with = "serde_bytes")]
-    pub drb_seed: [u8; 96],
+    pub drb_seed: DrbSeedInput,
 
-    /// the result of the DRB calculation
+    /// The DRB result for the current epoch.
+    ///
+    /// The DRB computation with this result was started two epochs ago.
     #[serde(with = "serde_bytes")]
-    pub drb_result: [u8; 32],
+    pub drb_result: DrbResult,
 }
 
 impl<TYPES: NodeType> Leaf2<TYPES> {
@@ -688,7 +697,7 @@ impl<TYPES: NodeType> Leaf2<TYPES> {
 
 impl<TYPES: NodeType> Committable for Leaf2<TYPES> {
     fn commit(&self) -> committable::Commitment<Self> {
-        if self.drb_seed == [0; 96] && self.drb_result == [0; 32] {
+        if self.drb_seed == [0; 32] && self.drb_result == [0; 32] {
             RawCommitmentBuilder::new("leaf commitment")
                 .u64_field("view number", *self.view_number)
                 .field("parent leaf commitment", self.parent_commitment)

crates/hotshot/src/traits/election/dynamic.rs → crates/types/src/drb.rs

@@ -6,9 +6,10 @@
 
 use std::hash::{DefaultHasher, Hash, Hasher};
 
-use hotshot_types::traits::{node_implementation::NodeType, signature_key::SignatureKey};
 use sha2::{Digest, Sha256};
 
+use crate::traits::{node_implementation::NodeType, signature_key::SignatureKey};
+
 // TODO: Add the following consts once we bench the hash time.
 // <https://github.com/EspressoSystems/HotShot/issues/3880>
 // /// Highest number of hashes that a hardware can complete in a second.
@@ -22,6 +23,17 @@ use sha2::{Digest, Sha256};
 /// Arbitrary number of times the hash function will be repeatedly called.
 const DIFFICULTY_LEVEL: u64 = 10;
 
+/// DRB seed input for epoch 1 and 2.
+pub const INITIAL_DRB_SEED_INPUT: [u8; 32] = [0; 32];
+/// DRB result for epoch 1 and 2.
+pub const INITIAL_DRB_RESULT: [u8; 32] = [0; 32];
+
+/// Alias for DRB seed input for `compute_drb_result`, serialized from the QC signature.
+pub type DrbSeedInput = [u8; 32];
+
+/// Alias for DRB result from `compute_drb_result`.
+pub type DrbResult = [u8; 32];
+
 // TODO: Use `HASHES_PER_SECOND` * `VIEW_TIMEOUT` * `DRB_CALCULATION_NUM_VIEW` to calculate this
 // once we bench the hash time.
 // <https://github.com/EspressoSystems/HotShot/issues/3880>
@@ -40,7 +52,7 @@ pub fn difficulty_level() -> u64 {
 /// # Arguments
 /// * `drb_seed_input` - Serialized QC signature.
 #[must_use]
-pub fn compute_drb_result<TYPES: NodeType>(drb_seed_input: [u8; 32]) -> [u8; 32] {
+pub fn compute_drb_result<TYPES: NodeType>(drb_seed_input: DrbSeedInput) -> DrbResult {
     let mut hash = drb_seed_input.to_vec();
     for _iter in 0..DIFFICULTY_LEVEL {
         // TODO: This may be optimized to avoid memcopies after we bench the hash time.
@@ -61,7 +73,7 @@ pub fn compute_drb_result<TYPES: NodeType>(drb_seed_input: [u8; 32]) -> [u8; 32]
 pub fn leader<TYPES: NodeType>(
     view_number: usize,
     stake_table: &[<TYPES::SignatureKey as SignatureKey>::StakeTableEntry],
-    drb_result: [u8; 32],
+    drb_result: DrbResult,
 ) -> TYPES::SignatureKey {
     let mut hasher = DefaultHasher::new();
     drb_result.hash(&mut hasher);

crates/types/src/lib.rs

@@ -20,6 +20,8 @@ pub mod bundle;
 pub mod consensus;
 pub mod constants;
 pub mod data;
+/// Holds the types and functions for DRB computation.
+pub mod drb;
 pub mod error;
 pub mod event;
 /// Holds the configuration file specification for a HotShot node.