Push DRB result to Membership

crates/hotshot/src/traits/election/static_committee.rs

@@ -6,7 +6,9 @@
 
 use std::{cmp::max, collections::BTreeMap, num::NonZeroU64};
 
+use async_trait::async_trait;
 use hotshot_types::{
+    drb::{self, DrbResult, INITIAL_DRB_RESULT},
     traits::{
         election::Membership,
         node_implementation::NodeType,
@@ -15,7 +17,7 @@ use hotshot_types::{
     PeerConfig,
 };
 use primitive_types::U256;
-use utils::anytrace::Result;
+use utils::anytrace::*;
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 /// The static committee election
@@ -38,8 +40,12 @@ pub struct StaticCommittee<T: NodeType> {
     /// The nodes on the committee and their stake, indexed by public key
     indexed_da_stake_table:
         BTreeMap<T::SignatureKey, <T::SignatureKey as SignatureKey>::StakeTableEntry>,
+
+    /// The results of DRB calculations
+    drb_result_table: BTreeMap<T::Epoch, DrbResult>,
 }
 
+#[async_trait]
 impl<TYPES: NodeType> Membership<TYPES> for StaticCommittee<TYPES> {
     type Error = utils::anytrace::Error;
 
@@ -95,6 +101,7 @@ impl<TYPES: NodeType> Membership<TYPES> for StaticCommittee<TYPES> {
             da_stake_table: da_members,
             indexed_stake_table,
             indexed_da_stake_table,
+            drb_result_table: BTreeMap::new(),
         }
     }
 
@@ -196,12 +203,29 @@ impl<TYPES: NodeType> Membership<TYPES> for StaticCommittee<TYPES> {
     fn lookup_leader(
         &self,
         view_number: <TYPES as NodeType>::View,
-        _epoch: Option<<TYPES as NodeType>::Epoch>,
+        epoch: Option<<TYPES as NodeType>::Epoch>,
     ) -> Result<TYPES::SignatureKey> {
-        #[allow(clippy::cast_possible_truncation)]
-        let index = *view_number as usize % self.eligible_leaders.len();
-        let res = self.eligible_leaders[index].clone();
-        Ok(TYPES::SignatureKey::public_key(&res))
+        if false{//let Some(epoch) = epoch {
+            /*let drb_result = if *epoch <= 2 {
+                &INITIAL_DRB_RESULT
+            } else {
+                self.drb_result_table
+                    .get(&epoch)
+                    .ok_or_else(|| panic!("DRB result not available for epoch {:?}", epoch))?
+            };*/
+            let drb_result = &INITIAL_DRB_RESULT;
+
+            Ok(drb::leader::<TYPES>(
+                view_number,
+                &self.eligible_leaders,
+                *drb_result,
+            ))
+        } else {
+            #[allow(clippy::cast_possible_truncation)]
+            let index = ((*view_number as usize) + 1) % self.eligible_leaders.len();
+            let res = self.eligible_leaders[index].clone();
+            Ok(TYPES::SignatureKey::public_key(&res))
+        }
     }
 
     /// Get the total number of nodes in the committee
@@ -234,4 +258,14 @@ impl<TYPES: NodeType> Membership<TYPES> for StaticCommittee<TYPES> {
         let len = self.stake_table.len();
         NonZeroU64::new(max((len as u64 * 9) / 10, ((len as u64 * 2) / 3) + 1)).unwrap()
     }
+
+    async fn add_drb_result(
+        &self,
+        epoch: TYPES::Epoch,
+        drb_result: DrbResult,
+    ) -> Option<Box<dyn FnOnce(&mut Self) + Send>> {
+        Some(Box::new(move |s: &mut Self| {
+            s.drb_result_table.insert(epoch, drb_result);
+        }))
+    }
 }

crates/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() {

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

@@ -46,6 +46,22 @@ 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,
+) {
+    let write_callback = {
+        let membership_reader = membership.read().await;
+        membership_reader.add_drb_result(epoch, drb_result).await
+    };
+
+    if let Some(write_callback) = write_callback {
+        let mut membership_writer = membership.write().await;
+        write_callback(&mut *membership_writer);
+    }
+}
+
 /// Store the DRB result from the computation task to the shared `results` table.
 ///
 /// Returns the result if it exists.
@@ -89,7 +105,12 @@ 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)
         }
         Err(e) => Err(warn!("Error in DRB calculation: {:?}.", e)),
@@ -196,6 +217,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 +310,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 +336,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.");
         }

crates/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

crates/testing/src/test_builder.rs

@@ -15,8 +15,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,
@@ -451,6 +452,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