Skip to content

Commit f7d38ef

Browse files
pls148pls148
committed
Push DRB result to Membership
1 parent bac3637 commit f7d38ef

23 files changed

+216
-132
lines changed

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

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

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

+31-3
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,14 @@ use crate::{
4646
quorum_vote::Versions,
4747
};
4848

49+
async fn notify_membership_of_drb_result<TYPES: NodeType>(
50+
membership: &Arc<RwLock<TYPES::Membership>>,
51+
epoch: <TYPES as NodeType>::Epoch,
52+
drb_result: DrbResult,
53+
) {
54+
membership.write().await.add_drb_result(epoch, drb_result);
55+
}
56+
4957
/// Store the DRB result from the computation task to the shared `results` table.
5058
///
5159
/// Returns the result if it exists.
@@ -89,6 +97,10 @@ async fn store_and_get_computed_drb_result<
8997
.drb_seeds_and_results
9098
.results
9199
.insert(epoch_number, result);
100+
drop(consensus_writer);
101+
102+
notify_membership_of_drb_result::<TYPES>(&task_state.membership, epoch_number, result)
103+
.await;
92104
task_state.drb_computation = None;
93105
Ok(result)
94106
}
@@ -196,6 +208,12 @@ async fn start_drb_task<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versio
196208
.drb_seeds_and_results
197209
.results
198210
.insert(*task_epoch, result);
211+
notify_membership_of_drb_result::<TYPES>(
212+
&task_state.membership,
213+
*task_epoch,
214+
result,
215+
)
216+
.await;
199217
task_state.drb_computation = None;
200218
}
201219
Err(e) => {
@@ -283,9 +301,13 @@ async fn store_drb_seed_and_result<TYPES: NodeType, I: NodeImplementation<TYPES>
283301
else {
284302
bail!("Failed to serialize the QC signature.");
285303
};
286-
let Ok(drb_seed_input) = drb_seed_input_vec.try_into() else {
287-
bail!("Failed to convert the serialized QC signature into a DRB seed input.");
288-
};
304+
305+
// TODO: Replace the leader election with a weighted version.
306+
// <https://github.com/EspressoSystems/HotShot/issues/3898>
307+
let mut drb_seed_input = [0u8; 32];
308+
let len = drb_seed_input_vec.len().min(32);
309+
drb_seed_input[..len].copy_from_slice(&drb_seed_input_vec[..len]);
310+
289311
task_state
290312
.consensus
291313
.write()
@@ -305,6 +327,12 @@ async fn store_drb_seed_and_result<TYPES: NodeType, I: NodeImplementation<TYPES>
305327
.drb_seeds_and_results
306328
.results
307329
.insert(current_epoch_number + 1, result);
330+
notify_membership_of_drb_result::<TYPES>(
331+
&task_state.membership,
332+
current_epoch_number + 1,
333+
result,
334+
)
335+
.await;
308336
} else {
309337
bail!("The last block of the epoch is decided but doesn't contain a DRB result.");
310338
}

‎hotshot-testing/src/helpers.rs

+14-3
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
// along with the HotShot repository. If not, see <https://mit-license.org/>.
66

77
#![allow(clippy::panic)]
8-
use std::{fmt::Debug, hash::Hash, marker::PhantomData, sync::Arc};
8+
use std::{collections::BTreeMap, fmt::Debug, hash::Hash, marker::PhantomData, sync::Arc};
99

1010
use async_broadcast::{Receiver, Sender};
1111
use async_lock::RwLock;
@@ -47,6 +47,11 @@ use vbs::version::Version;
4747

4848
use crate::{test_builder::TestDescription, test_launcher::TestLauncher};
4949

50+
pub type TestNodeKeyMap = BTreeMap<
51+
<TestTypes as NodeType>::SignatureKey,
52+
<<TestTypes as NodeType>::SignatureKey as SignatureKey>::PrivateKey,
53+
>;
54+
5055
/// create the [`SystemContextHandle`] from a node id, with no epochs
5156
/// # Panics
5257
/// if cannot create a [`HotShotInitializer`]
@@ -64,6 +69,7 @@ pub async fn build_system_handle<
6469
SystemContextHandle<TYPES, I, V>,
6570
Sender<Arc<HotShotEvent<TYPES>>>,
6671
Receiver<Arc<HotShotEvent<TYPES>>>,
72+
Arc<TestNodeKeyMap>,
6773
) {
6874
let builder: TestDescription<TYPES, I, V> = TestDescription::default_multiple_rounds();
6975

@@ -91,6 +97,7 @@ pub async fn build_system_handle_from_launcher<
9197
SystemContextHandle<TYPES, I, V>,
9298
Sender<Arc<HotShotEvent<TYPES>>>,
9399
Receiver<Arc<HotShotEvent<TYPES>>>,
100+
Arc<TestNodeKeyMap>,
94101
) {
95102
let network = (launcher.resource_generators.channel_generator)(node_id).await;
96103
let storage = (launcher.resource_generators.storage)(node_id);
@@ -118,7 +125,9 @@ pub async fn build_system_handle_from_launcher<
118125
hotshot_config.known_da_nodes.clone(),
119126
)));
120127

121-
SystemContext::init(
128+
let node_key_map = launcher.metadata.build_node_key_map();
129+
130+
let (c, s, r) = SystemContext::init(
122131
public_key,
123132
private_key,
124133
node_id,
@@ -131,7 +140,9 @@ pub async fn build_system_handle_from_launcher<
131140
marketplace_config,
132141
)
133142
.await
134-
.expect("Could not init hotshot")
143+
.expect("Could not init hotshot");
144+
145+
(c, s, r, node_key_map)
135146
}
136147

137148
/// create certificate

‎hotshot-testing/src/test_builder.rs

+13-2
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@ use hotshot::{
1414
HotShotInitializer, MarketplaceConfig, SystemContext, TwinsHandlerState,
1515
};
1616
use hotshot_example_types::{
17-
auction_results_provider_types::TestAuctionResultsProvider, state_types::TestInstanceState,
18-
storage_types::TestStorage, testable_delay::DelayConfig,
17+
auction_results_provider_types::TestAuctionResultsProvider, node_types::TestTypes,
18+
state_types::TestInstanceState, storage_types::TestStorage, testable_delay::DelayConfig,
1919
};
2020
use hotshot_types::{
2121
consensus::ConsensusMetricsValue,
@@ -32,6 +32,7 @@ use super::{
3232
txn_task::TxnTaskDescription,
3333
};
3434
use crate::{
35+
helpers::{key_pair_for_id, TestNodeKeyMap},
3536
spinning_task::SpinningTaskDescription,
3637
test_launcher::{Network, ResourceGenerators, TestLauncher},
3738
test_task::TestTaskStateSeed,
@@ -441,6 +442,16 @@ impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> TestDescription
441442
..self
442443
}
443444
}
445+
446+
pub fn build_node_key_map(&self) -> Arc<TestNodeKeyMap> {
447+
let mut node_key_map = TestNodeKeyMap::new();
448+
for i in 0..self.test_config.num_nodes_with_stake.into() {
449+
let (private_key, public_key) = key_pair_for_id::<TestTypes>(i as u64);
450+
node_key_map.insert(public_key, private_key);
451+
}
452+
453+
Arc::new(node_key_map)
454+
}
444455
}
445456

446457
impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> Default

‎hotshot-testing/src/view_generator.rs

+52-6
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ use hotshot_types::{
3737
},
3838
traits::{
3939
consensus_api::ConsensusApi,
40+
election::Membership,
4041
node_implementation::{ConsensusTime, NodeType, Versions},
4142
BlockPayload,
4243
},
@@ -46,7 +47,7 @@ use rand::{thread_rng, Rng};
4647
use sha2::{Digest, Sha256};
4748

4849
use crate::helpers::{
49-
build_cert, build_da_certificate, build_vid_proposal, da_payload_commitment, key_pair_for_id,
50+
build_cert, build_da_certificate, build_vid_proposal, da_payload_commitment, TestNodeKeyMap,
5051
};
5152

5253
#[derive(Clone)]
@@ -57,6 +58,7 @@ pub struct TestView {
5758
pub view_number: ViewNumber,
5859
pub epoch_number: Option<EpochNumber>,
5960
pub membership: Arc<RwLock<<TestTypes as NodeType>::Membership>>,
61+
pub node_key_map: Arc<TestNodeKeyMap>,
6062
pub vid_disperse: Proposal<TestTypes, VidDisperse<TestTypes>>,
6163
pub vid_proposal: (
6264
Vec<Proposal<TestTypes, VidDisperseShare<TestTypes>>>,
@@ -73,8 +75,31 @@ pub struct TestView {
7375
}
7476

7577
impl TestView {
78+
async fn find_leader_key_pair(
79+
membership: &Arc<RwLock<<TestTypes as NodeType>::Membership>>,
80+
node_key_map: &Arc<TestNodeKeyMap>,
81+
view_number: <TestTypes as NodeType>::View,
82+
epoch: Option<<TestTypes as NodeType>::Epoch>,
83+
) -> (
84+
<<TestTypes as NodeType>::SignatureKey as SignatureKey>::PrivateKey,
85+
<TestTypes as NodeType>::SignatureKey,
86+
) {
87+
let membership_reader = membership.read().await;
88+
let leader = membership_reader
89+
.leader(view_number, epoch)
90+
.expect("expected Membership::leader to succeed");
91+
drop(membership_reader);
92+
93+
let sk = node_key_map
94+
.get(&leader)
95+
.expect("expected Membership::leader public key to be in node_key_map");
96+
97+
(sk.clone(), leader)
98+
}
99+
76100
pub async fn genesis<V: Versions>(
77101
membership: &Arc<RwLock<<TestTypes as NodeType>::Membership>>,
102+
node_key_map: Arc<TestNodeKeyMap>,
78103
) -> Self {
79104
let genesis_view = ViewNumber::new(1);
80105
let genesis_epoch = genesis_epoch_from_version::<V, TestTypes>();
@@ -96,7 +121,10 @@ impl TestView {
96121
&metadata,
97122
);
98123

99-
let (private_key, public_key) = key_pair_for_id::<TestTypes>(*genesis_view);
124+
//let (private_key, public_key) = key_pair_for_id::<TestTypes>(*genesis_view);
125+
let (private_key, public_key) =
126+
Self::find_leader_key_pair(membership, &node_key_map, genesis_view, genesis_epoch)
127+
.await;
100128

101129
let leader_public_key = public_key;
102130

@@ -198,6 +226,7 @@ impl TestView {
198226
view_number: genesis_view,
199227
epoch_number: genesis_epoch,
200228
membership: membership.clone(),
229+
node_key_map,
201230
vid_disperse,
202231
vid_proposal: (vid_proposal, public_key),
203232
da_certificate,
@@ -235,9 +264,19 @@ impl TestView {
235264
epoch: old_epoch,
236265
};
237266

238-
let (old_private_key, old_public_key) = key_pair_for_id::<TestTypes>(*old_view);
267+
//let (old_private_key, old_public_key) = key_pair_for_id::<TestTypes>(*old_view);
268+
let (old_private_key, old_public_key) =
269+
Self::find_leader_key_pair(&self.membership, &self.node_key_map, old_view, old_epoch)
270+
.await;
239271

240-
let (private_key, public_key) = key_pair_for_id::<TestTypes>(*next_view);
272+
//let (private_key, public_key) = key_pair_for_id::<TestTypes>(*next_view);
273+
let (private_key, public_key) = Self::find_leader_key_pair(
274+
&self.membership,
275+
&self.node_key_map,
276+
next_view,
277+
self.epoch_number,
278+
)
279+
.await;
241280

242281
let leader_public_key = public_key;
243282

@@ -441,6 +480,7 @@ impl TestView {
441480
view_number: next_view,
442481
epoch_number: self.epoch_number,
443482
membership: self.membership.clone(),
483+
node_key_map: self.node_key_map.clone(),
444484
vid_disperse,
445485
vid_proposal: (vid_proposal, public_key),
446486
da_certificate,
@@ -517,14 +557,19 @@ impl TestView {
517557
pub struct TestViewGenerator<V: Versions> {
518558
pub current_view: Option<TestView>,
519559
pub membership: Arc<RwLock<<TestTypes as NodeType>::Membership>>,
560+
pub node_key_map: Arc<TestNodeKeyMap>,
520561
pub _pd: PhantomData<fn(V)>,
521562
}
522563

523564
impl<V: Versions> TestViewGenerator<V> {
524-
pub fn generate(membership: Arc<RwLock<<TestTypes as NodeType>::Membership>>) -> Self {
565+
pub fn generate(
566+
membership: Arc<RwLock<<TestTypes as NodeType>::Membership>>,
567+
node_key_map: Arc<TestNodeKeyMap>,
568+
) -> Self {
525569
TestViewGenerator {
526570
current_view: None,
527571
membership,
572+
node_key_map,
528573
_pd: PhantomData,
529574
}
530575
}
@@ -603,12 +648,13 @@ impl<V: Versions> Stream for TestViewGenerator<V> {
603648

604649
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
605650
let mem = Arc::clone(&self.membership);
651+
let nkm = Arc::clone(&self.node_key_map);
606652
let curr_view = &self.current_view.clone();
607653

608654
let mut fut = if let Some(ref view) = curr_view {
609655
async move { TestView::next_view(view).await }.boxed()
610656
} else {
611-
async move { TestView::genesis::<V>(&mem).await }.boxed()
657+
async move { TestView::genesis::<V>(&mem, nkm).await }.boxed()
612658
};
613659

614660
match fut.as_mut().poll(cx) {

‎hotshot-testing/tests/tests_1/block_builder.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ use hotshot_testing::block_builder::{
2121
use hotshot_types::{
2222
network::RandomBuilderConfig,
2323
traits::{
24-
block_contents::advz_commitment, node_implementation::NodeType, signature_key::SignatureKey,
25-
BlockPayload,
24+
block_contents::advz_commitment, node_implementation::NodeType,
25+
signature_key::SignatureKey, BlockPayload,
2626
},
2727
};
2828
use tide_disco::Url;

‎hotshot-testing/tests/tests_1/da_task.rs

+8-8
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,8 @@ use vbs::version::{StaticVersionType, Version};
3535
async fn test_da_task() {
3636
hotshot::helpers::initialize_logging();
3737

38-
let handle = build_system_handle::<TestTypes, MemoryImpl, TestVersions>(2)
39-
.await
40-
.0;
38+
let (handle, _, _, node_key_map) =
39+
build_system_handle::<TestTypes, MemoryImpl, TestVersions>(2).await;
4140

4241
let membership = Arc::clone(&handle.hotshot.memberships);
4342
let default_version = Version { major: 0, minor: 0 };
@@ -52,7 +51,8 @@ async fn test_da_task() {
5251
default_version,
5352
);
5453

55-
let mut generator = TestViewGenerator::<TestVersions>::generate(membership.clone());
54+
let mut generator =
55+
TestViewGenerator::<TestVersions>::generate(membership.clone(), node_key_map);
5656

5757
let mut proposals = Vec::new();
5858
let mut leaders = Vec::new();
@@ -142,9 +142,8 @@ async fn test_da_task() {
142142
async fn test_da_task_storage_failure() {
143143
hotshot::helpers::initialize_logging();
144144

145-
let handle = build_system_handle::<TestTypes, MemoryImpl, TestVersions>(2)
146-
.await
147-
.0;
145+
let (handle, _, _, node_key_map) =
146+
build_system_handle::<TestTypes, MemoryImpl, TestVersions>(2).await;
148147

149148
// Set the error flag here for the system handle. This causes it to emit an error on append.
150149
handle.storage().write().await.should_return_err = true;
@@ -161,7 +160,8 @@ async fn test_da_task_storage_failure() {
161160
default_version,
162161
);
163162

164-
let mut generator = TestViewGenerator::<TestVersions>::generate(Arc::clone(&membership));
163+
let mut generator =
164+
TestViewGenerator::<TestVersions>::generate(Arc::clone(&membership), node_key_map);
165165

166166
let mut proposals = Vec::new();
167167
let mut leaders = Vec::new();

‎hotshot-testing/tests/tests_1/message.rs

+4-4
Original file line numberDiff line numberDiff line change
@@ -76,12 +76,12 @@ async fn test_certificate2_validity() {
7676
hotshot::helpers::initialize_logging();
7777

7878
let node_id = 1;
79-
let handle = build_system_handle::<TestTypes, MemoryImpl, TestVersions>(node_id)
80-
.await
81-
.0;
79+
let (handle, _, _, node_key_map) =
80+
build_system_handle::<TestTypes, MemoryImpl, TestVersions>(node_id).await;
8281
let membership = Arc::clone(&handle.hotshot.memberships);
8382

84-
let mut generator = TestViewGenerator::<TestVersions>::generate(Arc::clone(&membership));
83+
let mut generator =
84+
TestViewGenerator::<TestVersions>::generate(Arc::clone(&membership), node_key_map);
8585

8686
let mut proposals = Vec::new();
8787
let mut leaders = Vec::new();

0 commit comments

Comments
 (0)