From f7d38ef11118f0ea21931ff304ef08ee7fb5a995 Mon Sep 17 00:00:00 2001 From: pls148 <184445976+pls148@users.noreply.github.com> Date: Fri, 14 Feb 2025 15:23:18 -0800 Subject: [PATCH 1/2] Push DRB result to Membership --- .../src/quorum_proposal/handlers.rs | 2 +- .../src/quorum_vote/handlers.rs | 34 ++++++++++- hotshot-testing/src/helpers.rs | 17 +++++- hotshot-testing/src/test_builder.rs | 15 ++++- hotshot-testing/src/view_generator.rs | 58 +++++++++++++++++-- .../tests/tests_1/block_builder.rs | 4 +- hotshot-testing/tests/tests_1/da_task.rs | 16 ++--- hotshot-testing/tests/tests_1/message.rs | 8 +-- hotshot-testing/tests/tests_1/network_task.rs | 14 ++--- .../tests_1/quorum_proposal_recv_task.rs | 14 ++--- .../tests/tests_1/quorum_proposal_task.rs | 47 ++++++++------- .../tests/tests_1/quorum_vote_task.rs | 21 +++---- .../tests_1/upgrade_task_with_proposal.rs | 8 +-- .../tests/tests_1/upgrade_task_with_vote.rs | 7 +-- .../tests/tests_1/vote_dependency_handle.rs | 7 +-- hotshot-types/src/drb.rs | 31 +--------- hotshot-types/src/traits/election.rs | 6 +- .../traits/election/randomized_committee.rs | 6 +- .../election/randomized_committee_members.rs | 4 +- .../src/traits/election/static_committee.rs | 5 +- .../static_committee_leader_two_views.rs | 6 +- .../traits/election/two_static_committees.rs | 6 +- types/src/v0/impls/stake_table.rs | 12 +++- 23 files changed, 216 insertions(+), 132 deletions(-) diff --git a/hotshot-task-impls/src/quorum_proposal/handlers.rs b/hotshot-task-impls/src/quorum_proposal/handlers.rs index 31a9d5ccc6..6e4e73d671 100644 --- a/hotshot-task-impls/src/quorum_proposal/handlers.rs +++ b/hotshot-task-impls/src/quorum_proposal/handlers.rs @@ -202,7 +202,7 @@ impl ProposalDependencyHandle { ) -> Option> { 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() { diff --git a/hotshot-task-impls/src/quorum_vote/handlers.rs b/hotshot-task-impls/src/quorum_vote/handlers.rs index 2869595088..4e17dd32a6 100644 --- a/hotshot-task-impls/src/quorum_vote/handlers.rs +++ b/hotshot-task-impls/src/quorum_vote/handlers.rs @@ -46,6 +46,14 @@ use crate::{ quorum_vote::Versions, }; +async fn notify_membership_of_drb_result( + membership: &Arc>, + epoch: ::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::(&task_state.membership, epoch_number, result) + .await; task_state.drb_computation = None; Ok(result) } @@ -196,6 +208,12 @@ async fn start_drb_task, V: Versio .drb_seeds_and_results .results .insert(*task_epoch, result); + notify_membership_of_drb_result::( + &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 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. + // + 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 .drb_seeds_and_results .results .insert(current_epoch_number + 1, result); + notify_membership_of_drb_result::( + &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."); } diff --git a/hotshot-testing/src/helpers.rs b/hotshot-testing/src/helpers.rs index 1bbb099c93..21a9451d3e 100644 --- a/hotshot-testing/src/helpers.rs +++ b/hotshot-testing/src/helpers.rs @@ -5,7 +5,7 @@ // along with the HotShot repository. If not, see . #![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< + ::SignatureKey, + <::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, Sender>>, Receiver>>, + Arc, ) { let builder: TestDescription = TestDescription::default_multiple_rounds(); @@ -91,6 +97,7 @@ pub async fn build_system_handle_from_launcher< SystemContextHandle, Sender>>, Receiver>>, + Arc, ) { 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 diff --git a/hotshot-testing/src/test_builder.rs b/hotshot-testing/src/test_builder.rs index 0e3ad9c487..452d90473a 100644 --- a/hotshot-testing/src/test_builder.rs +++ b/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, V: Versions> TestDescription ..self } } + + pub fn build_node_key_map(&self) -> Arc { + 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::(i as u64); + node_key_map.insert(public_key, private_key); + } + + Arc::new(node_key_map) + } } impl, V: Versions> Default diff --git a/hotshot-testing/src/view_generator.rs b/hotshot-testing/src/view_generator.rs index abe4277590..7c30f4ab16 100644 --- a/hotshot-testing/src/view_generator.rs +++ b/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, pub membership: Arc::Membership>>, + pub node_key_map: Arc, pub vid_disperse: Proposal>, pub vid_proposal: ( Vec>>, @@ -73,8 +75,31 @@ pub struct TestView { } impl TestView { + async fn find_leader_key_pair( + membership: &Arc::Membership>>, + node_key_map: &Arc, + view_number: ::View, + epoch: Option<::Epoch>, + ) -> ( + <::SignatureKey as SignatureKey>::PrivateKey, + ::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( membership: &Arc::Membership>>, + node_key_map: Arc, ) -> Self { let genesis_view = ViewNumber::new(1); let genesis_epoch = genesis_epoch_from_version::(); @@ -96,7 +121,10 @@ impl TestView { &metadata, ); - let (private_key, public_key) = key_pair_for_id::(*genesis_view); + //let (private_key, public_key) = key_pair_for_id::(*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::(*old_view); + //let (old_private_key, old_public_key) = key_pair_for_id::(*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::(*next_view); + //let (private_key, public_key) = key_pair_for_id::(*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 { pub current_view: Option, pub membership: Arc::Membership>>, + pub node_key_map: Arc, pub _pd: PhantomData, } impl TestViewGenerator { - pub fn generate(membership: Arc::Membership>>) -> Self { + pub fn generate( + membership: Arc::Membership>>, + node_key_map: Arc, + ) -> Self { TestViewGenerator { current_view: None, membership, + node_key_map, _pd: PhantomData, } } @@ -603,12 +648,13 @@ impl Stream for TestViewGenerator { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 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::(&mem).await }.boxed() + async move { TestView::genesis::(&mem, nkm).await }.boxed() }; match fut.as_mut().poll(cx) { diff --git a/hotshot-testing/tests/tests_1/block_builder.rs b/hotshot-testing/tests/tests_1/block_builder.rs index 00c111dfe5..58db015f20 100644 --- a/hotshot-testing/tests/tests_1/block_builder.rs +++ b/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; diff --git a/hotshot-testing/tests/tests_1/da_task.rs b/hotshot-testing/tests/tests_1/da_task.rs index 5006cab51d..c942dd8b76 100644 --- a/hotshot-testing/tests/tests_1/da_task.rs +++ b/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::(2) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(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::::generate(membership.clone()); + let mut generator = + TestViewGenerator::::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::(2) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(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::::generate(Arc::clone(&membership)); + let mut generator = + TestViewGenerator::::generate(Arc::clone(&membership), node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); diff --git a/hotshot-testing/tests/tests_1/message.rs b/hotshot-testing/tests/tests_1/message.rs index c2970befd5..568b2906f4 100644 --- a/hotshot-testing/tests/tests_1/message.rs +++ b/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::(node_id) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(node_id).await; let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(Arc::clone(&membership)); + let mut generator = + TestViewGenerator::::generate(Arc::clone(&membership), node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); diff --git a/hotshot-testing/tests/tests_1/network_task.rs b/hotshot-testing/tests/tests_1/network_task.rs index f634c60113..ef3584a55e 100644 --- a/hotshot-testing/tests/tests_1/network_task.rs +++ b/hotshot-testing/tests/tests_1/network_task.rs @@ -43,9 +43,8 @@ async fn test_network_task() { TestDescription::default_multiple_rounds(); let upgrade_lock = UpgradeLock::::new(); let node_id = 1; - let handle = build_system_handle::(node_id) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(node_id).await; let launcher = builder.gen_launcher(); let network = (launcher.resource_generators.channel_generator)(node_id).await; @@ -80,7 +79,7 @@ async fn test_network_task() { let task = Task::new(network_state, tx.clone(), rx); task_reg.run_task(task); - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); let view = generator.next().await.unwrap(); let (out_tx_internal, mut out_rx_internal) = async_broadcast::broadcast(10); @@ -215,9 +214,8 @@ async fn test_network_storage_fail() { let builder: TestDescription = TestDescription::default_multiple_rounds(); let node_id = 1; - let handle = build_system_handle::(node_id) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(node_id).await; let launcher = builder.gen_launcher(); let network = (launcher.resource_generators.channel_generator)(node_id).await; @@ -253,7 +251,7 @@ async fn test_network_storage_fail() { let task = Task::new(network_state, tx.clone(), rx); task_reg.run_task(task); - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); let view = generator.next().await.unwrap(); let (out_tx_internal, mut out_rx_internal): (Sender>>, _) = diff --git a/hotshot-testing/tests/tests_1/quorum_proposal_recv_task.rs b/hotshot-testing/tests/tests_1/quorum_proposal_recv_task.rs index 29d2389f2d..d7cd877ddc 100644 --- a/hotshot-testing/tests/tests_1/quorum_proposal_recv_task.rs +++ b/hotshot-testing/tests/tests_1/quorum_proposal_recv_task.rs @@ -51,14 +51,13 @@ async fn test_quorum_proposal_recv_task() { hotshot::helpers::initialize_logging(); - let handle = build_system_handle::(2) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(2).await; let membership = Arc::clone(&handle.hotshot.memberships); let consensus = handle.hotshot.consensus(); let mut consensus_writer = consensus.write().await; - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); let mut votes = Vec::new(); @@ -126,14 +125,13 @@ async fn test_quorum_proposal_recv_task_liveness_check() { hotshot::helpers::initialize_logging(); - let handle = build_system_handle::(4) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(4).await; let membership = Arc::clone(&handle.hotshot.memberships); let consensus = handle.hotshot.consensus(); let mut consensus_writer = consensus.write().await; - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); let mut votes = Vec::new(); diff --git a/hotshot-testing/tests/tests_1/quorum_proposal_task.rs b/hotshot-testing/tests/tests_1/quorum_proposal_task.rs index 9545f79cfc..65e8e4fa14 100644 --- a/hotshot-testing/tests/tests_1/quorum_proposal_task.rs +++ b/hotshot-testing/tests/tests_1/quorum_proposal_task.rs @@ -47,9 +47,8 @@ async fn test_quorum_proposal_task_quorum_proposal_view_1() { hotshot::helpers::initialize_logging(); let node_id = 1; - let handle = build_system_handle::(node_id) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(node_id).await; let membership = Arc::clone(&handle.hotshot.memberships); let version = handle @@ -66,7 +65,8 @@ async fn test_quorum_proposal_task_quorum_proposal_view_1() { ) .await; - let mut generator = TestViewGenerator::::generate(Arc::clone(&membership)); + let mut generator = + TestViewGenerator::::generate(Arc::clone(&membership), node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); @@ -151,13 +151,13 @@ async fn test_quorum_proposal_task_quorum_proposal_view_gt_1() { hotshot::helpers::initialize_logging(); let node_id = 3; - let handle = build_system_handle::(node_id) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(node_id).await; let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(membership.clone()); + let mut generator = + TestViewGenerator::::generate(membership.clone(), node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); @@ -335,9 +335,8 @@ async fn test_quorum_proposal_task_qc_timeout() { hotshot::helpers::initialize_logging(); let node_id = 3; - let handle = build_system_handle::(node_id) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(node_id).await; let membership = Arc::clone(&handle.hotshot.memberships); let version = handle .hotshot @@ -354,7 +353,8 @@ async fn test_quorum_proposal_task_qc_timeout() { .await; let builder_commitment = BuilderCommitment::from_raw_digest(sha2::Sha256::new().finalize()); - let mut generator = TestViewGenerator::::generate(Arc::clone(&membership)); + let mut generator = + TestViewGenerator::::generate(Arc::clone(&membership), node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); @@ -433,9 +433,8 @@ async fn test_quorum_proposal_task_view_sync() { hotshot::helpers::initialize_logging(); let node_id = 2; - let handle = build_system_handle::(node_id) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(node_id).await; let membership = Arc::clone(&handle.hotshot.memberships); let version = handle @@ -453,7 +452,8 @@ async fn test_quorum_proposal_task_view_sync() { .await; let builder_commitment = BuilderCommitment::from_raw_digest(sha2::Sha256::new().finalize()); - let mut generator = TestViewGenerator::::generate(Arc::clone(&membership)); + let mut generator = + TestViewGenerator::::generate(Arc::clone(&membership), node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); @@ -532,13 +532,13 @@ async fn test_quorum_proposal_task_liveness_check() { hotshot::helpers::initialize_logging(); let node_id = 3; - let handle = build_system_handle::(node_id) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(node_id).await; let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(Arc::clone(&membership)); + let mut generator = + TestViewGenerator::::generate(Arc::clone(&membership), node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); @@ -711,12 +711,11 @@ async fn test_quorum_proposal_task_liveness_check() { async fn test_quorum_proposal_task_with_incomplete_events() { hotshot::helpers::initialize_logging(); - let handle = build_system_handle::(2) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(2).await; let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); diff --git a/hotshot-testing/tests/tests_1/quorum_vote_task.rs b/hotshot-testing/tests/tests_1/quorum_vote_task.rs index 4ea37ebe00..ae8c2bc9ad 100644 --- a/hotshot-testing/tests/tests_1/quorum_vote_task.rs +++ b/hotshot-testing/tests/tests_1/quorum_vote_task.rs @@ -41,13 +41,12 @@ async fn test_quorum_vote_task_success() { hotshot::helpers::initialize_logging(); - let handle = build_system_handle::(2) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(2).await; let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); let mut proposals = Vec::new(); let mut leaves = Vec::new(); @@ -108,13 +107,12 @@ async fn test_quorum_vote_task_miss_dependency() { hotshot::helpers::initialize_logging(); - let handle = build_system_handle::(2) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(2).await; let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); let mut proposals = Vec::new(); let mut leaders = Vec::new(); @@ -192,13 +190,12 @@ async fn test_quorum_vote_task_incorrect_dependency() { hotshot::helpers::initialize_logging(); - let handle = build_system_handle::(2) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(2).await; let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); let mut proposals = Vec::new(); let mut leaves = Vec::new(); diff --git a/hotshot-testing/tests/tests_1/upgrade_task_with_proposal.rs b/hotshot-testing/tests/tests_1/upgrade_task_with_proposal.rs index 820fb3a562..c80d0419d5 100644 --- a/hotshot-testing/tests/tests_1/upgrade_task_with_proposal.rs +++ b/hotshot-testing/tests/tests_1/upgrade_task_with_proposal.rs @@ -55,9 +55,8 @@ async fn test_upgrade_task_with_proposal() { hotshot::helpers::initialize_logging(); - let handle = build_system_handle::(3) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(3).await; let other_handles = futures::future::join_all((0..=9).map(build_system_handle)).await; @@ -85,7 +84,8 @@ async fn test_upgrade_task_with_proposal() { let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(Arc::clone(&membership)); + let mut generator = + TestViewGenerator::::generate(Arc::clone(&membership), node_key_map); for view in (&mut generator).take(1).collect::>().await { proposals.push(view.quorum_proposal.clone()); diff --git a/hotshot-testing/tests/tests_1/upgrade_task_with_vote.rs b/hotshot-testing/tests/tests_1/upgrade_task_with_vote.rs index 841fd86a86..4047f95607 100644 --- a/hotshot-testing/tests/tests_1/upgrade_task_with_vote.rs +++ b/hotshot-testing/tests/tests_1/upgrade_task_with_vote.rs @@ -45,9 +45,8 @@ async fn test_upgrade_task_with_vote() { hotshot::helpers::initialize_logging(); - let handle = build_system_handle::(2) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(2).await; let old_version = Version { major: 0, minor: 1 }; let new_version = Version { major: 0, minor: 2 }; @@ -71,7 +70,7 @@ async fn test_upgrade_task_with_vote() { let mut consensus_writer = consensus.write().await; let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); for view in (&mut generator).take(2).collect::>().await { proposals.push(view.quorum_proposal.clone()); diff --git a/hotshot-testing/tests/tests_1/vote_dependency_handle.rs b/hotshot-testing/tests/tests_1/vote_dependency_handle.rs index a6e36c800f..e6e40f2b8c 100644 --- a/hotshot-testing/tests/tests_1/vote_dependency_handle.rs +++ b/hotshot-testing/tests/tests_1/vote_dependency_handle.rs @@ -33,12 +33,11 @@ async fn test_vote_dependency_handle() { // We use a node ID of 2 here arbitrarily. We just need it to build the system handle. let node_id = 2; // Construct the system handle for the node ID to build all of the state objects. - let handle = build_system_handle::(node_id) - .await - .0; + let (handle, _, _, node_key_map) = + build_system_handle::(node_id).await; let membership = Arc::clone(&handle.hotshot.memberships); - let mut generator = TestViewGenerator::::generate(membership); + let mut generator = TestViewGenerator::::generate(membership, node_key_map); // Generate our state for the test let mut proposals = Vec::new(); diff --git a/hotshot-types/src/drb.rs b/hotshot-types/src/drb.rs index e52fb89605..94727b7bc2 100644 --- a/hotshot-types/src/drb.rs +++ b/hotshot-types/src/drb.rs @@ -4,18 +4,12 @@ // You should have received a copy of the MIT License // along with the HotShot repository. If not, see . -use std::{ - collections::BTreeMap, - hash::{DefaultHasher, Hash, Hasher}, -}; +use std::collections::BTreeMap; use sha2::{Digest, Sha256}; use tokio::task::JoinHandle; -use crate::traits::{ - node_implementation::{ConsensusTime, NodeType}, - signature_key::SignatureKey, -}; +use crate::traits::node_implementation::{ConsensusTime, NodeType}; // TODO: Add the following consts once we bench the hash time. // @@ -76,27 +70,6 @@ pub fn compute_drb_result(drb_seed_input: DrbSeedInput) -> DrbR drb_result } -/// Use the DRB result to get the leader. -/// -/// The DRB result is the output of a spawned `compute_drb_result` call. -#[must_use] -pub fn leader( - view_number: usize, - stake_table: &[::StakeTableEntry], - drb_result: DrbResult, -) -> TYPES::SignatureKey { - let mut hasher = DefaultHasher::new(); - drb_result.hash(&mut hasher); - view_number.hash(&mut hasher); - #[allow(clippy::cast_possible_truncation)] - // TODO: Use the total stake rather than `len()` and update the indexing after switching to - // a weighted stake table. - // - let index = (hasher.finish() as usize) % stake_table.len(); - let entry = stake_table[index].clone(); - TYPES::SignatureKey::public_key(&entry) -} - /// Alias for in-progress DRB computation task, if there's any. pub type DrbComputation = Option<(::Epoch, JoinHandle)>; diff --git a/hotshot-types/src/traits/election.rs b/hotshot-types/src/traits/election.rs index f442182b42..2473f6b357 100644 --- a/hotshot-types/src/traits/election.rs +++ b/hotshot-types/src/traits/election.rs @@ -11,7 +11,7 @@ use async_trait::async_trait; use hotshot_utils::anytrace::Result; use super::node_implementation::NodeType; -use crate::{traits::signature_key::SignatureKey, PeerConfig}; +use crate::{drb::DrbResult, traits::signature_key::SignatureKey, PeerConfig}; #[async_trait] /// A protocol for determining membership in and participating in a committee. @@ -152,4 +152,8 @@ pub trait Membership: Debug + Send + Sync { async fn sync_l1(&self) -> Option> { None } + + /// Called to notify the Membership when a new DRB result has been calculated. + /// Observes the same semantics as add_epoch_root + fn add_drb_result(&mut self, _epoch: TYPES::Epoch, _drb_result: DrbResult); } diff --git a/hotshot/src/traits/election/randomized_committee.rs b/hotshot/src/traits/election/randomized_committee.rs index 92eec78764..72a1711b27 100644 --- a/hotshot/src/traits/election/randomized_committee.rs +++ b/hotshot/src/traits/election/randomized_committee.rs @@ -4,9 +4,8 @@ // You should have received a copy of the MIT License // along with the HotShot repository. If not, see . -use std::{cmp::max, collections::BTreeMap, num::NonZeroU64}; - use hotshot_types::{ + drb::DrbResult, traits::{ election::Membership, node_implementation::NodeType, @@ -17,6 +16,7 @@ use hotshot_types::{ use hotshot_utils::anytrace::Result; use primitive_types::U256; use rand::{rngs::StdRng, Rng}; +use std::{cmp::max, collections::BTreeMap, num::NonZeroU64}; #[derive(Clone, Debug, Eq, PartialEq, Hash)] @@ -247,4 +247,6 @@ impl Membership for RandomizedCommittee { )) .unwrap() } + + fn add_drb_result(&mut self, _epoch: ::Epoch, _drb_result: DrbResult) {} } diff --git a/hotshot/src/traits/election/randomized_committee_members.rs b/hotshot/src/traits/election/randomized_committee_members.rs index 408a4f3a8e..4a3c8d508a 100644 --- a/hotshot/src/traits/election/randomized_committee_members.rs +++ b/hotshot/src/traits/election/randomized_committee_members.rs @@ -3,7 +3,6 @@ // You should have received a copy of the MIT License // along with the HotShot repository. If not, see . - use std::{ cmp::max, collections::{BTreeMap, BTreeSet}, @@ -12,6 +11,7 @@ use std::{ }; use hotshot_types::{ + drb::DrbResult, traits::{ election::Membership, node_implementation::{ConsensusTime, NodeType}, @@ -447,4 +447,6 @@ impl Membership let len = self.total_nodes(epoch); NonZeroU64::new(max((len as u64 * 9) / 10, ((len as u64 * 2) / 3) + 1)).unwrap() } + + fn add_drb_result(&mut self, _epoch: ::Epoch, _drb_result: DrbResult) {} } diff --git a/hotshot/src/traits/election/static_committee.rs b/hotshot/src/traits/election/static_committee.rs index 1e30126a42..7bcf0fc431 100644 --- a/hotshot/src/traits/election/static_committee.rs +++ b/hotshot/src/traits/election/static_committee.rs @@ -7,6 +7,7 @@ use std::{cmp::max, collections::BTreeMap, num::NonZeroU64}; use hotshot_types::{ + drb::DrbResult, traits::{ election::Membership, node_implementation::NodeType, @@ -14,7 +15,7 @@ use hotshot_types::{ }, PeerConfig, }; -use hotshot_utils::anytrace::Result; +use hotshot_utils::anytrace::*; use primitive_types::U256; #[derive(Clone, Debug, Eq, PartialEq, Hash)] @@ -234,4 +235,6 @@ impl Membership for StaticCommittee { let len = self.stake_table.len(); NonZeroU64::new(max((len as u64 * 9) / 10, ((len as u64 * 2) / 3) + 1)).unwrap() } + + fn add_drb_result(&mut self, _epoch: ::Epoch, _drb_result: DrbResult) {} } diff --git a/hotshot/src/traits/election/static_committee_leader_two_views.rs b/hotshot/src/traits/election/static_committee_leader_two_views.rs index 567c24f742..52412c26d5 100644 --- a/hotshot/src/traits/election/static_committee_leader_two_views.rs +++ b/hotshot/src/traits/election/static_committee_leader_two_views.rs @@ -4,9 +4,8 @@ // You should have received a copy of the MIT License // along with the HotShot repository. If not, see . -use std::{collections::BTreeMap, num::NonZeroU64}; - use hotshot_types::{ + drb::DrbResult, traits::{ election::Membership, node_implementation::NodeType, @@ -16,6 +15,7 @@ use hotshot_types::{ }; use hotshot_utils::anytrace::Result; use primitive_types::U256; +use std::{collections::BTreeMap, num::NonZeroU64}; #[derive(Clone, Debug, Eq, PartialEq, Hash)] @@ -235,4 +235,6 @@ impl Membership for StaticCommitteeLeaderForTwoViews::Epoch>) -> NonZeroU64 { NonZeroU64::new(((self.stake_table.len() as u64 * 9) / 10) + 1).unwrap() } + + fn add_drb_result(&mut self, _epoch: ::Epoch, _drb_result: DrbResult) {} } diff --git a/hotshot/src/traits/election/two_static_committees.rs b/hotshot/src/traits/election/two_static_committees.rs index 2f43c9be30..1c7c583cbc 100644 --- a/hotshot/src/traits/election/two_static_committees.rs +++ b/hotshot/src/traits/election/two_static_committees.rs @@ -4,9 +4,8 @@ // You should have received a copy of the MIT License // along with the HotShot repository. If not, see . -use std::{cmp::max, collections::BTreeMap, num::NonZeroU64}; - use hotshot_types::{ + drb::DrbResult, traits::{ election::Membership, node_implementation::NodeType, @@ -16,6 +15,7 @@ use hotshot_types::{ }; use hotshot_utils::anytrace::Result; use primitive_types::U256; +use std::{cmp::max, collections::BTreeMap, num::NonZeroU64}; /// Tuple type for eligible leaders type EligibleLeaders = ( @@ -424,4 +424,6 @@ impl Membership for TwoStaticCommittees { .unwrap() } } + + fn add_drb_result(&mut self, _epoch: ::Epoch, _drb_result: DrbResult) {} } diff --git a/types/src/v0/impls/stake_table.rs b/types/src/v0/impls/stake_table.rs index f5206fe065..a1096df6a9 100644 --- a/types/src/v0/impls/stake_table.rs +++ b/types/src/v0/impls/stake_table.rs @@ -11,6 +11,7 @@ use hotshot::types::{BLSPubKey, SignatureKey as _}; use hotshot_contract_adapter::stake_table::{bls_alloy_to_jf, NodeInfoJf}; use hotshot_types::{ data::EpochNumber, + drb::DrbResult, stake_table::StakeTableEntry, traits::{ election::Membership, @@ -22,7 +23,7 @@ use hotshot_types::{ use itertools::Itertools; use std::{ cmp::max, - collections::{BTreeSet, HashMap}, + collections::{BTreeMap, BTreeSet, HashMap}, num::NonZeroU64, str::FromStr, }; @@ -105,6 +106,9 @@ pub struct EpochCommittees { /// Address of Stake Table Contract contract_address: Option
, + + /// The results of DRB calculations + drb_result_table: BTreeMap, } #[derive(Debug, Clone, PartialEq)] @@ -252,6 +256,7 @@ impl EpochCommittees { _epoch_size: epoch_size, l1_client: instance_state.l1_client.clone(), contract_address: instance_state.chain_config.stake_table_contract, + drb_result_table: BTreeMap::new(), } } @@ -332,6 +337,7 @@ impl Membership for EpochCommittees { l1_client: L1Client::new(vec![Url::from_str("http:://ab.b").unwrap()]) .expect("Failed to create L1 client"), contract_address: None, + drb_result_table: BTreeMap::new(), } } @@ -497,6 +503,10 @@ impl Membership for EpochCommittees { }) }) } + + fn add_drb_result(&mut self, epoch: Epoch, drb_result: DrbResult) { + self.drb_result_table.insert(epoch, drb_result); + } } #[cfg(test)] From 5f82edc8a75e2544e2ce58cf25b067f81a4cb81a Mon Sep 17 00:00:00 2001 From: pls148 <184445976+pls148@users.noreply.github.com> Date: Tue, 18 Feb 2025 18:51:49 -0800 Subject: [PATCH 2/2] hotshot: add tests for drb results --- hotshot-example-types/src/node_types.rs | 33 +++ hotshot-testing/tests/tests_6/test_epochs.rs | 15 +- hotshot-types/src/drb.rs | 35 ++- hotshot/src/traits/election/mod.rs | 3 + .../election/static_committee_with_drb.rs | 263 ++++++++++++++++++ 5 files changed, 340 insertions(+), 9 deletions(-) create mode 100644 hotshot/src/traits/election/static_committee_with_drb.rs diff --git a/hotshot-example-types/src/node_types.rs b/hotshot-example-types/src/node_types.rs index 033feb514b..54c2111b91 100644 --- a/hotshot-example-types/src/node_types.rs +++ b/hotshot-example-types/src/node_types.rs @@ -15,6 +15,7 @@ use hotshot::traits::{ randomized_committee_members::RandomizedCommitteeMembers, static_committee::StaticCommittee, static_committee_leader_two_views::StaticCommitteeLeaderForTwoViews, + static_committee_with_drb::StaticCommitteeWithDrb, two_static_committees::TwoStaticCommittees, }, implementations::{CombinedNetworks, Libp2pNetwork, MemoryNetwork, PushCdnNetwork}, @@ -69,6 +70,38 @@ impl NodeType for TestTypes { type BuilderSignatureKey = BuilderKey; } +#[derive( + Copy, + Clone, + Debug, + Default, + Hash, + PartialEq, + Eq, + PartialOrd, + Ord, + serde::Serialize, + serde::Deserialize, +)] +/// filler struct to implement node type and allow us +/// to select our traits +pub struct TestTypesWithDrb; +impl NodeType for TestTypesWithDrb { + const UPGRADE_CONSTANTS: UpgradeConstants = TEST_UPGRADE_CONSTANTS; + + type AuctionResult = TestAuctionResult; + type View = ViewNumber; + type Epoch = EpochNumber; + type BlockHeader = TestBlockHeader; + type BlockPayload = TestBlockPayload; + type SignatureKey = BLSPubKey; + type Transaction = TestTransaction; + type ValidatedState = TestValidatedState; + type InstanceState = TestInstanceState; + type Membership = StaticCommitteeWithDrb; + type BuilderSignatureKey = BuilderKey; +} + #[derive( Copy, Clone, diff --git a/hotshot-testing/tests/tests_6/test_epochs.rs b/hotshot-testing/tests/tests_6/test_epochs.rs index 67cace8ce3..407432f326 100644 --- a/hotshot-testing/tests/tests_6/test_epochs.rs +++ b/hotshot-testing/tests/tests_6/test_epochs.rs @@ -10,7 +10,7 @@ use hotshot_example_types::{ node_types::{ CombinedImpl, EpochUpgradeTestVersions, EpochsTestVersions, Libp2pImpl, MemoryImpl, PushCdnImpl, RandomOverlapQuorumFilterConfig, StableQuorumFilterConfig, - TestConsecutiveLeaderTypes, TestTwoStakeTablesTypes, TestTypes, + TestConsecutiveLeaderTypes, TestTwoStakeTablesTypes, TestTypes,TestTypesWithDrb, TestTypesRandomizedCommitteeMembers, TestTypesRandomizedLeader, }, testable_delay::{DelayConfig, DelayOptions, DelaySettings, SupportedTraitTypesForAsyncDelay}, @@ -60,6 +60,7 @@ cross_tests!( TestTypesRandomizedCommitteeMembers>, // Overlap = 2F+1 TestTypesRandomizedCommitteeMembers>, // Overlap = 3F TestTypesRandomizedCommitteeMembers>, // Overlap = Dynamic + TestTypesWithDrb, ], Versions: [EpochsTestVersions], Ignore: false, @@ -79,7 +80,7 @@ cross_tests!( cross_tests!( TestName: test_success_with_async_delay_with_epochs, Impls: [Libp2pImpl, PushCdnImpl, CombinedImpl], - Types: [TestTypes, TestTwoStakeTablesTypes], + Types: [TestTypes, TestTwoStakeTablesTypes, TestTypesWithDrb], Versions: [EpochsTestVersions], Ignore: false, Metadata: { @@ -111,7 +112,7 @@ cross_tests!( cross_tests!( TestName: test_success_with_async_delay_2_with_epochs, Impls: [Libp2pImpl, PushCdnImpl, CombinedImpl], - Types: [TestTypes, TestTwoStakeTablesTypes], + Types: [TestTypes, TestTwoStakeTablesTypes, TestTypesWithDrb], Versions: [EpochsTestVersions], Ignore: false, Metadata: { @@ -168,7 +169,7 @@ cross_tests!( cross_tests!( TestName: test_epoch_end, Impls: [CombinedImpl, Libp2pImpl, PushCdnImpl], - Types: [TestTypes, TestTwoStakeTablesTypes], + Types: [TestTypes, TestTwoStakeTablesTypes, TestTypesWithDrb], Versions: [EpochsTestVersions], Ignore: false, Metadata: { @@ -192,7 +193,7 @@ cross_tests!( cross_tests!( TestName: test_shorter_decide, Impls: [Libp2pImpl, PushCdnImpl, CombinedImpl], - Types: [TestTypes, TestTwoStakeTablesTypes], + Types: [TestTypes, TestTwoStakeTablesTypes, TestTypesWithDrb], Versions: [EpochsTestVersions], Ignore: false, Metadata: { @@ -495,7 +496,7 @@ cross_tests!( cross_tests!( TestName: test_all_restart_epochs, Impls: [CombinedImpl, PushCdnImpl], - Types: [TestTypes, TestTypesRandomizedLeader, TestTwoStakeTablesTypes], + Types: [TestTypes, TestTypesRandomizedLeader, TestTwoStakeTablesTypes, TestTypesWithDrb], Versions: [EpochsTestVersions], Ignore: false, Metadata: { @@ -544,7 +545,7 @@ cross_tests!( cross_tests!( TestName: test_all_restart_one_da_with_epochs, Impls: [CombinedImpl], - Types: [TestTypes, TestTwoStakeTablesTypes], + Types: [TestTypes, TestTwoStakeTablesTypes, TestTypesWithDrb], Versions: [EpochsTestVersions], Ignore: false, Metadata: { diff --git a/hotshot-types/src/drb.rs b/hotshot-types/src/drb.rs index 94727b7bc2..4fae3904c1 100644 --- a/hotshot-types/src/drb.rs +++ b/hotshot-types/src/drb.rs @@ -4,12 +4,18 @@ // You should have received a copy of the MIT License // along with the HotShot repository. If not, see . -use std::collections::BTreeMap; +use std::{ + collections::BTreeMap, + hash::{DefaultHasher, Hash, Hasher}, +}; use sha2::{Digest, Sha256}; use tokio::task::JoinHandle; -use crate::traits::node_implementation::{ConsensusTime, NodeType}; +use crate::traits::{ + node_implementation::{ConsensusTime, NodeType}, + signature_key::SignatureKey, +}; // TODO: Add the following consts once we bench the hash time. // @@ -70,6 +76,31 @@ pub fn compute_drb_result(drb_seed_input: DrbSeedInput) -> DrbR drb_result } +/// Use the DRB result to get the leader. +/// +/// The DRB result is the output of a spawned `compute_drb_result` call. +#[must_use] +pub fn leader( + view_number: TYPES::View, + stake_table: &[::StakeTableEntry], + drb_result: DrbResult, +) -> TYPES::SignatureKey { + let mut hasher = DefaultHasher::new(); + + drb_result.hash(&mut hasher); + view_number.hash(&mut hasher); + + #[allow(clippy::cast_possible_truncation)] + let hash = hasher.finish() as usize; + + // TODO: Use the total stake rather than `len()` and update the indexing after switching to + // a weighted stake table. + // + let index = hash % stake_table.len(); + let entry = stake_table[index].clone(); + TYPES::SignatureKey::public_key(&entry) +} + /// Alias for in-progress DRB computation task, if there's any. pub type DrbComputation = Option<(::Epoch, JoinHandle)>; diff --git a/hotshot/src/traits/election/mod.rs b/hotshot/src/traits/election/mod.rs index 5cc908a6d6..3dba37bda8 100644 --- a/hotshot/src/traits/election/mod.rs +++ b/hotshot/src/traits/election/mod.rs @@ -15,6 +15,9 @@ pub mod randomized_committee_members; /// static (round robin) committee election pub mod static_committee; +/// static (round robin) committee election with DRB +pub mod static_committee_with_drb; + /// static (round robin leader for 2 consecutive views) committee election pub mod static_committee_leader_two_views; /// two static (round robin) committees for even and odd epochs diff --git a/hotshot/src/traits/election/static_committee_with_drb.rs b/hotshot/src/traits/election/static_committee_with_drb.rs new file mode 100644 index 0000000000..99182104d7 --- /dev/null +++ b/hotshot/src/traits/election/static_committee_with_drb.rs @@ -0,0 +1,263 @@ +// Copyright (c) 2021-2024 Espresso Systems (espressosys.com) +// This file is part of the HotShot repository. + +// You should have received a copy of the MIT License +// along with the HotShot repository. If not, see . + +use std::{cmp::max, collections::BTreeMap, num::NonZeroU64}; + +use hotshot_types::{ + drb::{self, DrbResult, INITIAL_DRB_RESULT}, + traits::{ + election::Membership, + node_implementation::NodeType, + signature_key::{SignatureKey, StakeTableEntryType}, + }, + PeerConfig, +}; +use hotshot_utils::anytrace::*; +use primitive_types::U256; + +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +/// The static committee election +pub struct StaticCommitteeWithDrb { + /// The nodes eligible for leadership. + /// NOTE: This is currently a hack because the DA leader needs to be the quorum + /// leader but without voting rights. + eligible_leaders: Vec<::StakeTableEntry>, + + /// The nodes on the committee and their stake + stake_table: Vec<::StakeTableEntry>, + + /// The nodes on the committee and their stake + da_stake_table: Vec<::StakeTableEntry>, + + /// The nodes on the committee and their stake, indexed by public key + indexed_stake_table: + BTreeMap::StakeTableEntry>, + + /// The nodes on the committee and their stake, indexed by public key + indexed_da_stake_table: + BTreeMap::StakeTableEntry>, + + /// The results of DRB calculations + drb_result_table: BTreeMap, +} + +impl Membership for StaticCommitteeWithDrb { + type Error = hotshot_utils::anytrace::Error; + + /// Create a new election + fn new( + committee_members: Vec::SignatureKey>>, + da_members: Vec::SignatureKey>>, + ) -> Self { + // For each eligible leader, get the stake table entry + let eligible_leaders: Vec<::StakeTableEntry> = + committee_members + .iter() + .map(|member| member.stake_table_entry.clone()) + .filter(|entry| entry.stake() > U256::zero()) + .collect(); + + // For each member, get the stake table entry + let members: Vec<::StakeTableEntry> = + committee_members + .iter() + .map(|member| member.stake_table_entry.clone()) + .filter(|entry| entry.stake() > U256::zero()) + .collect(); + + // For each member, get the stake table entry + let da_members: Vec<::StakeTableEntry> = da_members + .iter() + .map(|member| member.stake_table_entry.clone()) + .filter(|entry| entry.stake() > U256::zero()) + .collect(); + + // Index the stake table by public key + let indexed_stake_table: BTreeMap< + TYPES::SignatureKey, + ::StakeTableEntry, + > = members + .iter() + .map(|entry| (TYPES::SignatureKey::public_key(entry), entry.clone())) + .collect(); + + // Index the stake table by public key + let indexed_da_stake_table: BTreeMap< + TYPES::SignatureKey, + ::StakeTableEntry, + > = da_members + .iter() + .map(|entry| (TYPES::SignatureKey::public_key(entry), entry.clone())) + .collect(); + + Self { + eligible_leaders, + stake_table: members, + da_stake_table: da_members, + indexed_stake_table, + indexed_da_stake_table, + drb_result_table: BTreeMap::new(), + } + } + + /// Get the stake table for the current view + fn stake_table( + &self, + _epoch: Option<::Epoch>, + ) -> Vec<<::SignatureKey as SignatureKey>::StakeTableEntry> { + self.stake_table.clone() + } + + /// Get the stake table for the current view + fn da_stake_table( + &self, + _epoch: Option<::Epoch>, + ) -> Vec<<::SignatureKey as SignatureKey>::StakeTableEntry> { + self.da_stake_table.clone() + } + + /// Get all members of the committee for the current view + fn committee_members( + &self, + _view_number: ::View, + _epoch: Option<::Epoch>, + ) -> std::collections::BTreeSet<::SignatureKey> { + self.stake_table + .iter() + .map(TYPES::SignatureKey::public_key) + .collect() + } + + /// Get all members of the committee for the current view + fn da_committee_members( + &self, + _view_number: ::View, + _epoch: Option<::Epoch>, + ) -> std::collections::BTreeSet<::SignatureKey> { + self.da_stake_table + .iter() + .map(TYPES::SignatureKey::public_key) + .collect() + } + + /// Get all eligible leaders of the committee for the current view + fn committee_leaders( + &self, + _view_number: ::View, + _epoch: Option<::Epoch>, + ) -> std::collections::BTreeSet<::SignatureKey> { + self.eligible_leaders + .iter() + .map(TYPES::SignatureKey::public_key) + .collect() + } + + /// Get the stake table entry for a public key + fn stake( + &self, + pub_key: &::SignatureKey, + _epoch: Option<::Epoch>, + ) -> Option<::StakeTableEntry> { + // Only return the stake if it is above zero + self.indexed_stake_table.get(pub_key).cloned() + } + + /// Get the DA stake table entry for a public key + fn da_stake( + &self, + pub_key: &::SignatureKey, + _epoch: Option<::Epoch>, + ) -> Option<::StakeTableEntry> { + // Only return the stake if it is above zero + self.indexed_da_stake_table.get(pub_key).cloned() + } + + /// Check if a node has stake in the committee + fn has_stake( + &self, + pub_key: &::SignatureKey, + _epoch: Option<::Epoch>, + ) -> bool { + self.indexed_stake_table + .get(pub_key) + .is_some_and(|x| x.stake() > U256::zero()) + } + + /// Check if a node has stake in the committee + fn has_da_stake( + &self, + pub_key: &::SignatureKey, + _epoch: Option<::Epoch>, + ) -> bool { + self.indexed_da_stake_table + .get(pub_key) + .is_some_and(|x| x.stake() > U256::zero()) + } + + /// Index the vector of public keys with the current view number + fn lookup_leader( + &self, + view_number: ::View, + epoch: Option<::Epoch>, + ) -> Result { + if 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::( + view_number, + &self.eligible_leaders, + *drb_result, + )) + } else { + #[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)) + } + } + + /// Get the total number of nodes in the committee + fn total_nodes(&self, _epoch: Option<::Epoch>) -> usize { + self.stake_table.len() + } + + /// Get the total number of DA nodes in the committee + fn da_total_nodes(&self, _epoch: Option<::Epoch>) -> usize { + self.da_stake_table.len() + } + + /// Get the voting success threshold for the committee + fn success_threshold(&self, _epoch: Option<::Epoch>) -> NonZeroU64 { + NonZeroU64::new(((self.stake_table.len() as u64 * 2) / 3) + 1).unwrap() + } + + /// Get the voting success threshold for the committee + fn da_success_threshold(&self, _epoch: Option<::Epoch>) -> NonZeroU64 { + NonZeroU64::new(((self.da_stake_table.len() as u64 * 2) / 3) + 1).unwrap() + } + + /// Get the voting failure threshold for the committee + fn failure_threshold(&self, _epoch: Option<::Epoch>) -> NonZeroU64 { + NonZeroU64::new(((self.stake_table.len() as u64) / 3) + 1).unwrap() + } + + /// Get the voting upgrade threshold for the committee + fn upgrade_threshold(&self, _epoch: Option<::Epoch>) -> NonZeroU64 { + let len = self.stake_table.len(); + NonZeroU64::new(max((len as u64 * 9) / 10, ((len as u64 * 2) / 3) + 1)).unwrap() + } + + fn add_drb_result(&mut self, epoch: ::Epoch, drb_result: DrbResult) { + self.drb_result_table.insert(epoch, drb_result); + } +}