Push DRB result to Membership

hotshot-task-impls/src/quorum_proposal/handlers.rs

@@ -202,7 +202,7 @@ impl<TYPES: NodeType, V: Versions> ProposalDependencyHandle<TYPES, V> {
     ) -> Option<NextEpochQuorumCertificate2<TYPES>> {
         tracing::debug!("getting the next epoch QC");
         // If we haven't upgraded to Epochs just return None right away
-        if self.upgrade_lock.version_infallible(self.view_number).await < V::Epochs::VERSION {
+        if !self.upgrade_lock.epochs_enabled(self.view_number).await {
             return None;
         }
         if let Some(next_epoch_qc) = self.consensus.read().await.next_epoch_high_qc() {

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

@@ -46,6 +46,14 @@ use crate::{
     quorum_vote::Versions,
 };
 
+async fn notify_membership_of_drb_result<TYPES: NodeType>(
+    membership: &Arc<RwLock<TYPES::Membership>>,
+    epoch: <TYPES as NodeType>::Epoch,
+    drb_result: DrbResult,
+) {
+    membership.write().await.add_drb_result(epoch, drb_result);
+}
+
 /// Store the DRB result from the computation task to the shared `results` table.
 ///
 /// Returns the result if it exists.
@@ -89,6 +97,10 @@ async fn store_and_get_computed_drb_result<
                 .drb_seeds_and_results
                 .results
                 .insert(epoch_number, result);
+            drop(consensus_writer);
+
+            notify_membership_of_drb_result::<TYPES>(&task_state.membership, epoch_number, result)
+                .await;
             task_state.drb_computation = None;
             Ok(result)
         }
@@ -196,6 +208,12 @@ async fn start_drb_task<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versio
                             .drb_seeds_and_results
                             .results
                             .insert(*task_epoch, result);
+                        notify_membership_of_drb_result::<TYPES>(
+                            &task_state.membership,
+                            *task_epoch,
+                            result,
+                        )
+                        .await;
                         task_state.drb_computation = None;
                     }
                     Err(e) => {
@@ -283,9 +301,13 @@ async fn store_drb_seed_and_result<TYPES: NodeType, I: NodeImplementation<TYPES>
         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.");
-        };
+
+        // TODO: Replace the leader election with a weighted version.
+        // <https://github.com/EspressoSystems/HotShot/issues/3898>
+        let mut drb_seed_input = [0u8; 32];
+        let len = drb_seed_input_vec.len().min(32);
+        drb_seed_input[..len].copy_from_slice(&drb_seed_input_vec[..len]);
+
         task_state
             .consensus
             .write()
@@ -305,6 +327,12 @@ async fn store_drb_seed_and_result<TYPES: NodeType, I: NodeImplementation<TYPES>
                 .drb_seeds_and_results
                 .results
                 .insert(current_epoch_number + 1, result);
+            notify_membership_of_drb_result::<TYPES>(
+                &task_state.membership,
+                current_epoch_number + 1,
+                result,
+            )
+            .await;
         } else {
             bail!("The last block of the epoch is decided but doesn't contain a DRB result.");
         }

hotshot-testing/src/helpers.rs

@@ -5,7 +5,7 @@
 // along with the HotShot repository. If not, see <https://mit-license.org/>.
 
 #![allow(clippy::panic)]
-use std::{fmt::Debug, hash::Hash, marker::PhantomData, sync::Arc};
+use std::{collections::BTreeMap, fmt::Debug, hash::Hash, marker::PhantomData, sync::Arc};
 
 use async_broadcast::{Receiver, Sender};
 use async_lock::RwLock;
@@ -47,6 +47,11 @@ use vbs::version::Version;
 
 use crate::{test_builder::TestDescription, test_launcher::TestLauncher};
 
+pub type TestNodeKeyMap = BTreeMap<
+    <TestTypes as NodeType>::SignatureKey,
+    <<TestTypes as NodeType>::SignatureKey as SignatureKey>::PrivateKey,
+>;
+
 /// create the [`SystemContextHandle`] from a node id, with no epochs
 /// # Panics
 /// if cannot create a [`HotShotInitializer`]
@@ -64,6 +69,7 @@ pub async fn build_system_handle<
     SystemContextHandle<TYPES, I, V>,
     Sender<Arc<HotShotEvent<TYPES>>>,
     Receiver<Arc<HotShotEvent<TYPES>>>,
+    Arc<TestNodeKeyMap>,
 ) {
     let builder: TestDescription<TYPES, I, V> = TestDescription::default_multiple_rounds();
 
@@ -91,6 +97,7 @@ pub async fn build_system_handle_from_launcher<
     SystemContextHandle<TYPES, I, V>,
     Sender<Arc<HotShotEvent<TYPES>>>,
     Receiver<Arc<HotShotEvent<TYPES>>>,
+    Arc<TestNodeKeyMap>,
 ) {
     let network = (launcher.resource_generators.channel_generator)(node_id).await;
     let storage = (launcher.resource_generators.storage)(node_id);
@@ -118,7 +125,9 @@ pub async fn build_system_handle_from_launcher<
         hotshot_config.known_da_nodes.clone(),
     )));
 
-    SystemContext::init(
+    let node_key_map = launcher.metadata.build_node_key_map();
+
+    let (c, s, r) = SystemContext::init(
         public_key,
         private_key,
         node_id,
@@ -131,7 +140,9 @@ pub async fn build_system_handle_from_launcher<
         marketplace_config,
     )
     .await
-    .expect("Could not init hotshot")
+    .expect("Could not init hotshot");
+
+    (c, s, r, node_key_map)
 }
 
 /// create certificate

hotshot-testing/src/test_builder.rs

@@ -14,8 +14,8 @@ use hotshot::{
     HotShotInitializer, MarketplaceConfig, SystemContext, TwinsHandlerState,
 };
 use hotshot_example_types::{
-    auction_results_provider_types::TestAuctionResultsProvider, state_types::TestInstanceState,
-    storage_types::TestStorage, testable_delay::DelayConfig,
+    auction_results_provider_types::TestAuctionResultsProvider, node_types::TestTypes,
+    state_types::TestInstanceState, storage_types::TestStorage, testable_delay::DelayConfig,
 };
 use hotshot_types::{
     consensus::ConsensusMetricsValue,
@@ -32,6 +32,7 @@ use super::{
     txn_task::TxnTaskDescription,
 };
 use crate::{
+    helpers::{key_pair_for_id, TestNodeKeyMap},
     spinning_task::SpinningTaskDescription,
     test_launcher::{Network, ResourceGenerators, TestLauncher},
     test_task::TestTaskStateSeed,
@@ -441,6 +442,16 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> TestDescription
             ..self
         }
     }
+
+    pub fn build_node_key_map(&self) -> Arc<TestNodeKeyMap> {
+        let mut node_key_map = TestNodeKeyMap::new();
+        for i in 0..self.test_config.num_nodes_with_stake.into() {
+            let (private_key, public_key) = key_pair_for_id::<TestTypes>(i as u64);
+            node_key_map.insert(public_key, private_key);
+        }
+
+        Arc::new(node_key_map)
+    }
 }
 
 impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> Default

hotshot-testing/src/view_generator.rs

@@ -37,6 +37,7 @@ use hotshot_types::{
     },
     traits::{
         consensus_api::ConsensusApi,
+        election::Membership,
         node_implementation::{ConsensusTime, NodeType, Versions},
         BlockPayload,
     },
@@ -46,7 +47,7 @@ use rand::{thread_rng, Rng};
 use sha2::{Digest, Sha256};
 
 use crate::helpers::{
-    build_cert, build_da_certificate, build_vid_proposal, da_payload_commitment, key_pair_for_id,
+    build_cert, build_da_certificate, build_vid_proposal, da_payload_commitment, TestNodeKeyMap,
 };
 
 #[derive(Clone)]
@@ -57,6 +58,7 @@ pub struct TestView {
     pub view_number: ViewNumber,
     pub epoch_number: Option<EpochNumber>,
     pub membership: Arc<RwLock<<TestTypes as NodeType>::Membership>>,
+    pub node_key_map: Arc<TestNodeKeyMap>,
     pub vid_disperse: Proposal<TestTypes, VidDisperse<TestTypes>>,
     pub vid_proposal: (
         Vec<Proposal<TestTypes, VidDisperseShare<TestTypes>>>,
@@ -73,8 +75,31 @@ pub struct TestView {
 }
 
 impl TestView {
+    async fn find_leader_key_pair(
+        membership: &Arc<RwLock<<TestTypes as NodeType>::Membership>>,
+        node_key_map: &Arc<TestNodeKeyMap>,
+        view_number: <TestTypes as NodeType>::View,
+        epoch: Option<<TestTypes as NodeType>::Epoch>,
+    ) -> (
+        <<TestTypes as NodeType>::SignatureKey as SignatureKey>::PrivateKey,
+        <TestTypes as NodeType>::SignatureKey,
+    ) {
+        let membership_reader = membership.read().await;
+        let leader = membership_reader
+            .leader(view_number, epoch)
+            .expect("expected Membership::leader to succeed");
+        drop(membership_reader);
+
+        let sk = node_key_map
+            .get(&leader)
+            .expect("expected Membership::leader public key to be in node_key_map");
+
+        (sk.clone(), leader)
+    }
+
     pub async fn genesis<V: Versions>(
         membership: &Arc<RwLock<<TestTypes as NodeType>::Membership>>,
+        node_key_map: Arc<TestNodeKeyMap>,
     ) -> Self {
         let genesis_view = ViewNumber::new(1);
         let genesis_epoch = genesis_epoch_from_version::<V, TestTypes>();
@@ -96,7 +121,10 @@ impl TestView {
             &metadata,
         );
 
-        let (private_key, public_key) = key_pair_for_id::<TestTypes>(*genesis_view);
+        //let (private_key, public_key) = key_pair_for_id::<TestTypes>(*genesis_view);
+        let (private_key, public_key) =
+            Self::find_leader_key_pair(membership, &node_key_map, genesis_view, genesis_epoch)
+                .await;
 
         let leader_public_key = public_key;
 
@@ -198,6 +226,7 @@ impl TestView {
             view_number: genesis_view,
             epoch_number: genesis_epoch,
             membership: membership.clone(),
+            node_key_map,
             vid_disperse,
             vid_proposal: (vid_proposal, public_key),
             da_certificate,
@@ -235,9 +264,19 @@ impl TestView {
             epoch: old_epoch,
         };
 
-        let (old_private_key, old_public_key) = key_pair_for_id::<TestTypes>(*old_view);
+        //let (old_private_key, old_public_key) = key_pair_for_id::<TestTypes>(*old_view);
+        let (old_private_key, old_public_key) =
+            Self::find_leader_key_pair(&self.membership, &self.node_key_map, old_view, old_epoch)
+                .await;
 
-        let (private_key, public_key) = key_pair_for_id::<TestTypes>(*next_view);
+        //let (private_key, public_key) = key_pair_for_id::<TestTypes>(*next_view);
+        let (private_key, public_key) = Self::find_leader_key_pair(
+            &self.membership,
+            &self.node_key_map,
+            next_view,
+            self.epoch_number,
+        )
+        .await;
 
         let leader_public_key = public_key;
 
@@ -441,6 +480,7 @@ impl TestView {
             view_number: next_view,
             epoch_number: self.epoch_number,
             membership: self.membership.clone(),
+            node_key_map: self.node_key_map.clone(),
             vid_disperse,
             vid_proposal: (vid_proposal, public_key),
             da_certificate,
@@ -517,14 +557,19 @@ impl TestView {
 pub struct TestViewGenerator<V: Versions> {
     pub current_view: Option<TestView>,
     pub membership: Arc<RwLock<<TestTypes as NodeType>::Membership>>,
+    pub node_key_map: Arc<TestNodeKeyMap>,
     pub _pd: PhantomData<fn(V)>,
 }
 
 impl<V: Versions> TestViewGenerator<V> {
-    pub fn generate(membership: Arc<RwLock<<TestTypes as NodeType>::Membership>>) -> Self {
+    pub fn generate(
+        membership: Arc<RwLock<<TestTypes as NodeType>::Membership>>,
+        node_key_map: Arc<TestNodeKeyMap>,
+    ) -> Self {
         TestViewGenerator {
             current_view: None,
             membership,
+            node_key_map,
             _pd: PhantomData,
         }
     }
@@ -603,12 +648,13 @@ impl<V: Versions> Stream for TestViewGenerator<V> {
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let mem = Arc::clone(&self.membership);
+        let nkm = Arc::clone(&self.node_key_map);
         let curr_view = &self.current_view.clone();
 
         let mut fut = if let Some(ref view) = curr_view {
             async move { TestView::next_view(view).await }.boxed()
         } else {
-            async move { TestView::genesis::<V>(&mem).await }.boxed()
+            async move { TestView::genesis::<V>(&mem, nkm).await }.boxed()
         };
 
         match fut.as_mut().poll(cx) {

hotshot-testing/tests/tests_1/block_builder.rs

@@ -21,8 +21,8 @@ use hotshot_testing::block_builder::{
 use hotshot_types::{
     network::RandomBuilderConfig,
     traits::{
-        block_contents::advz_commitment, node_implementation::NodeType, signature_key::SignatureKey,
-        BlockPayload,
+        block_contents::advz_commitment, node_implementation::NodeType,
+        signature_key::SignatureKey, BlockPayload,
     },
 };
 use tide_disco::Url;

hotshot-testing/tests/tests_1/da_task.rs

@@ -35,9 +35,8 @@ use vbs::version::{StaticVersionType, Version};
 async fn test_da_task() {
     hotshot::helpers::initialize_logging();
 
-    let handle = build_system_handle::<TestTypes, MemoryImpl, TestVersions>(2)
-        .await
-        .0;
+    let (handle, _, _, node_key_map) =
+        build_system_handle::<TestTypes, MemoryImpl, TestVersions>(2).await;
 
     let membership = Arc::clone(&handle.hotshot.memberships);
     let default_version = Version { major: 0, minor: 0 };
@@ -52,7 +51,8 @@ async fn test_da_task() {
         default_version,
     );
 
-    let mut generator = TestViewGenerator::<TestVersions>::generate(membership.clone());
+    let mut generator =
+        TestViewGenerator::<TestVersions>::generate(membership.clone(), node_key_map);
 
     let mut proposals = Vec::new();
     let mut leaders = Vec::new();
@@ -142,9 +142,8 @@ async fn test_da_task() {
 async fn test_da_task_storage_failure() {
     hotshot::helpers::initialize_logging();
 
-    let handle = build_system_handle::<TestTypes, MemoryImpl, TestVersions>(2)
-        .await
-        .0;
+    let (handle, _, _, node_key_map) =
+        build_system_handle::<TestTypes, MemoryImpl, TestVersions>(2).await;
 
     // Set the error flag here for the system handle. This causes it to emit an error on append.
     handle.storage().write().await.should_return_err = true;
@@ -161,7 +160,8 @@ async fn test_da_task_storage_failure() {
         default_version,
     );
 
-    let mut generator = TestViewGenerator::<TestVersions>::generate(Arc::clone(&membership));
+    let mut generator =
+        TestViewGenerator::<TestVersions>::generate(Arc::clone(&membership), node_key_map);
 
     let mut proposals = Vec::new();
     let mut leaders = Vec::new();

hotshot-testing/tests/tests_1/message.rs

@@ -76,12 +76,12 @@ async fn test_certificate2_validity() {
     hotshot::helpers::initialize_logging();
 
     let node_id = 1;
-    let handle = build_system_handle::<TestTypes, MemoryImpl, TestVersions>(node_id)
-        .await
-        .0;
+    let (handle, _, _, node_key_map) =
+        build_system_handle::<TestTypes, MemoryImpl, TestVersions>(node_id).await;
     let membership = Arc::clone(&handle.hotshot.memberships);
 
-    let mut generator = TestViewGenerator::<TestVersions>::generate(Arc::clone(&membership));
+    let mut generator =
+        TestViewGenerator::<TestVersions>::generate(Arc::clone(&membership), node_key_map);
 
     let mut proposals = Vec::new();
     let mut leaders = Vec::new();