2652 Membership::set_first_epoch (#2654)

* 2652 Membership::set_first_epoch

* fix for comments in PR review

* add logic to block upgrade

* fix

* rustfmt

* testing stuff

---------

Co-authored-by: ss-es <[email protected]>

builder/src/lib.rs

@@ -139,6 +139,7 @@ pub mod testing {
                 stop_proposing_time: 0,
                 stop_voting_time: 0,
                 epoch_height: 0,
+                epoch_start_block: 0,
             };
 
             Self {

hotshot-examples/infra/mod.rs

@@ -368,6 +368,8 @@ pub trait RunDa<
         let initializer = hotshot::HotShotInitializer::<TYPES>::from_genesis::<V>(
             TestInstanceState::default(),
             self.config().config.epoch_height,
+            self.config().config.epoch_start_block,
+            vec![],
         )
         .await
         .expect("Couldn't generate genesis block");

hotshot-orchestrator/run-config.toml

@@ -35,6 +35,7 @@ fixed_leader_for_gpuvid = 1
 next_view_timeout = 30000
 num_bootstrap = 5
 epoch_height = 0
+epoch_start_block = 0
 
 [random_builder]
 txn_in_block = 100

hotshot-query-service/examples/simple-server.rs

@@ -217,6 +217,7 @@ async fn init_consensus(
         start_voting_time: 0,
         stop_voting_time: 0,
         epoch_height: 0,
+        epoch_start_block: 0,
     };
 
     let nodes = join_all(priv_keys.into_iter().zip(data_sources).enumerate().map(
@@ -246,6 +247,8 @@ async fn init_consensus(
                     HotShotInitializer::from_genesis::<MockVersions>(
                         TestInstanceState::default(),
                         0,
+                        0,
+                        vec![],
                     )
                     .await
                     .unwrap(),

hotshot-query-service/src/testing/consensus.rs

@@ -149,6 +149,7 @@ impl<D: DataSourceLifeCycle + UpdateStatusData> MockNetwork<D> {
             start_voting_time: 0,
             stop_voting_time: 0,
             epoch_height: 0,
+            epoch_start_block: 0,
         };
         update_config(&mut config);
 
@@ -191,6 +192,8 @@ impl<D: DataSourceLifeCycle + UpdateStatusData> MockNetwork<D> {
                             HotShotInitializer::from_genesis::<MockVersions>(
                                 TestInstanceState::default(),
                                 0,
+                                0,
+                                vec![],
                             )
                             .await
                             .unwrap(),

hotshot-task-impls/src/helpers.rs

@@ -186,6 +186,7 @@ async fn decide_epoch_root<TYPES: NodeType>(
             TYPES::Epoch::new(epoch_from_block_number(decided_block_number, epoch_height) + 1);
 
         let write_callback = {
+            tracing::debug!("Calling add_epoch_root for epoch {:?}", next_epoch_number);
             let membership_reader = membership.read().await;
             membership_reader
                 .add_epoch_root(next_epoch_number, decided_leaf.block_header().clone())
@@ -364,7 +365,7 @@ pub async fn decide_from_proposal_2<TYPES: NodeType>(
 /// # Panics
 /// If the leaf chain contains no decided leaf while reaching a decided view, which should be
 /// impossible.
-pub async fn decide_from_proposal<TYPES: NodeType>(
+pub async fn decide_from_proposal<TYPES: NodeType, V: Versions>(
     proposal: &QuorumProposalWrapper<TYPES>,
     consensus: OuterConsensus<TYPES>,
     existing_upgrade_cert: Arc<RwLock<Option<UpgradeCertificate<TYPES>>>>,
@@ -477,10 +478,10 @@ pub async fn decide_from_proposal<TYPES: NodeType>(
         tracing::debug!("Leaf ascension failed; error={e}");
     }
 
-    if with_epochs && res.new_decided_view_number.is_some() {
-        let epoch_height = consensus_reader.epoch_height;
-        drop(consensus_reader);
+    let epoch_height = consensus_reader.epoch_height;
+    drop(consensus_reader);
 
+    if with_epochs && res.new_decided_view_number.is_some() {
         if let Some(decided_leaf_info) = res.leaf_views.last() {
             decide_epoch_root(&decided_leaf_info.leaf, epoch_height, membership).await;
         } else {

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

@@ -13,7 +13,7 @@ use committable::Committable;
 use hotshot_types::{
     consensus::OuterConsensus,
     data::{Leaf2, QuorumProposalWrapper, VidDisperseShare},
-    drb::{compute_drb_result, DrbResult},
+    drb::{compute_drb_result, DrbResult, INITIAL_DRB_RESULT},
     event::{Event, EventType},
     message::{Proposal, UpgradeLock},
     simple_vote::{HasEpoch, QuorumData2, QuorumVote2},
@@ -51,6 +51,7 @@ async fn notify_membership_of_drb_result<TYPES: NodeType>(
     epoch: <TYPES as NodeType>::Epoch,
     drb_result: DrbResult,
 ) {
+    tracing::debug!("Calling add_drb_result for epoch {:?}", epoch);
     membership.write().await.add_drb_result(epoch, drb_result);
 }
 
@@ -379,7 +380,7 @@ pub(crate) async fn handle_quorum_proposal_validated<
         )
         .await
     } else {
-        decide_from_proposal(
+        decide_from_proposal::<TYPES, V>(
             proposal,
             OuterConsensus::new(Arc::clone(&task_state.consensus.inner_consensus)),
             Arc::clone(&task_state.upgrade_lock.decided_upgrade_certificate),
@@ -390,21 +391,58 @@ pub(crate) async fn handle_quorum_proposal_validated<
         .await
     };
 
+    if let Some(cert) = &task_state.staged_epoch_upgrade_certificate {
+        if leaf_views.last().unwrap().leaf.height() >= task_state.epoch_upgrade_block_height {
+            let mut decided_certificate_lock = task_state
+                .upgrade_lock
+                .decided_upgrade_certificate
+                .write()
+                .await;
+            *decided_certificate_lock = Some(cert.clone());
+            drop(decided_certificate_lock);
+
+            let _ = task_state
+                .storage
+                .write()
+                .await
+                .update_decided_upgrade_certificate(Some(cert.clone()))
+                .await;
+
+            task_state.staged_epoch_upgrade_certificate = None;
+        }
+    };
+
     if let Some(cert) = decided_upgrade_cert.clone() {
-        let mut decided_certificate_lock = task_state
-            .upgrade_lock
-            .decided_upgrade_certificate
-            .write()
-            .await;
-        *decided_certificate_lock = Some(cert.clone());
-        drop(decided_certificate_lock);
+        if cert.data.new_version == V::Epochs::VERSION {
+            task_state.staged_epoch_upgrade_certificate = Some(cert);
 
-        let _ = task_state
-            .storage
-            .write()
-            .await
-            .update_decided_upgrade_certificate(Some(cert.clone()))
-            .await;
+            let epoch_height = task_state.consensus.read().await.epoch_height;
+            let first_epoch_number = TYPES::Epoch::new(epoch_from_block_number(
+                task_state.epoch_upgrade_block_height,
+                epoch_height,
+            ));
+            tracing::debug!("Calling set_first_epoch for epoch {:?}", first_epoch_number);
+            task_state
+                .membership
+                .write()
+                .await
+                .set_first_epoch(first_epoch_number, INITIAL_DRB_RESULT);
+        } else {
+            let mut decided_certificate_lock = task_state
+                .upgrade_lock
+                .decided_upgrade_certificate
+                .write()
+                .await;
+            *decided_certificate_lock = Some(cert.clone());
+            drop(decided_certificate_lock);
+
+            let _ = task_state
+                .storage
+                .write()
+                .await
+                .update_decided_upgrade_certificate(Some(cert.clone()))
+                .await;
+        }
     }
 
     let mut consensus_writer = task_state.consensus.write().await;

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

@@ -22,6 +22,7 @@ use hotshot_types::{
     drb::DrbComputation,
     event::Event,
     message::{Proposal, UpgradeLock},
+    simple_certificate::UpgradeCertificate,
     simple_vote::HasEpoch,
     traits::{
         block_contents::BlockHeader,
@@ -338,6 +339,12 @@ pub struct QuorumVoteTaskState<TYPES: NodeType, I: NodeImplementation<TYPES>, V:
 
     /// Number of blocks in an epoch, zero means there are no epochs
     pub epoch_height: u64,
+
+    /// Upgrade certificate to enable epochs, staged until we reach the specified block height
+    pub staged_epoch_upgrade_certificate: Option<UpgradeCertificate<TYPES>>,
+
+    /// Block height at which to enable the epoch upgrade
+    pub epoch_upgrade_block_height: u64,
 }
 
 impl<TYPES: NodeType, I: NodeImplementation<TYPES>, V: Versions> QuorumVoteTaskState<TYPES, I, V> {

hotshot-testing/src/helpers.rs

@@ -5,16 +5,14 @@
 // along with the HotShot repository. If not, see <https://mit-license.org/>.
 
 #![allow(clippy::panic)]
-use std::{collections::BTreeMap, fmt::Debug, hash::Hash, marker::PhantomData, sync::Arc};
-
 use async_broadcast::{Receiver, Sender};
 use async_lock::RwLock;
 use bitvec::bitvec;
 use committable::Committable;
 use hotshot::{
     traits::{BlockPayload, NodeImplementation, TestableNodeImplementation},
     types::{SignatureKey, SystemContextHandle},
-    HotShotInitializer, SystemContext,
+    HotShotInitializer, InitializerEpochInfo, SystemContext,
 };
 use hotshot_example_types::{
     auction_results_provider_types::TestAuctionResultsProvider,
@@ -24,9 +22,11 @@ use hotshot_example_types::{
     storage_types::TestStorage,
 };
 use hotshot_task_impls::events::HotShotEvent;
+use hotshot_types::traits::node_implementation::ConsensusTime;
 use hotshot_types::{
     consensus::ConsensusMetricsValue,
     data::{vid_commitment, Leaf2, VidCommitment, VidDisperse, VidDisperseShare},
+    drb::INITIAL_DRB_RESULT,
     message::{Proposal, UpgradeLock},
     simple_certificate::DaCertificate2,
     simple_vote::{DaData2, DaVote2, SimpleVote, VersionedVoteData},
@@ -41,6 +41,7 @@ use hotshot_types::{
 };
 use primitive_types::U256;
 use serde::Serialize;
+use std::{collections::BTreeMap, fmt::Debug, hash::Hash, marker::PhantomData, sync::Arc};
 use vbs::version::Version;
 
 use crate::{test_builder::TestDescription, test_launcher::TestLauncher};
@@ -105,6 +106,19 @@ pub async fn build_system_handle_from_launcher<
     let initializer = HotShotInitializer::<TYPES>::from_genesis::<V>(
         TestInstanceState::new(launcher.metadata.async_delay_config.clone()),
         launcher.metadata.test_config.epoch_height,
+        launcher.metadata.test_config.epoch_start_block,
+        vec![
+            InitializerEpochInfo::<TYPES> {
+                epoch: TYPES::Epoch::new(1),
+                drb_result: INITIAL_DRB_RESULT,
+                block_header: None,
+            },
+            InitializerEpochInfo::<TYPES> {
+                epoch: TYPES::Epoch::new(2),
+                drb_result: INITIAL_DRB_RESULT,
+                block_header: None,
+            },
+        ],
     )
     .await
     .unwrap();

hotshot-testing/src/spinning_task.rs

@@ -14,7 +14,8 @@ use async_lock::RwLock;
 use async_trait::async_trait;
 use futures::future::join_all;
 use hotshot::{
-    traits::TestableNodeImplementation, types::EventType, HotShotInitializer, SystemContext,
+    traits::TestableNodeImplementation, types::EventType, HotShotInitializer, InitializerEpochInfo,
+    SystemContext,
 };
 use hotshot_example_types::{
     auction_results_provider_types::TestAuctionResultsProvider,
@@ -57,6 +58,10 @@ pub struct SpinningTask<
 > {
     /// epoch height
     pub epoch_height: u64,
+    /// Epoch start block
+    pub epoch_start_block: u64,
+    /// Saved epoch information. This must be sorted ascending by epoch.
+    pub start_epoch_info: Vec<InitializerEpochInfo<TYPES>>,
     /// handle to the nodes
     pub(crate) handles: Arc<RwLock<Vec<Node<TYPES, I, V>>>>,
     /// late start nodes
@@ -161,6 +166,8 @@ where
                                         let initializer = HotShotInitializer::<TYPES>::load(
                                             TestInstanceState::new(self.async_delay_config.clone()),
                                             self.epoch_height,
+                                            self.epoch_start_block,
+                                            self.start_epoch_info.clone(),
                                             self.last_decided_leaf.clone(),
                                             (
                                                 TYPES::View::genesis(),
@@ -268,6 +275,8 @@ where
                                 let initializer = HotShotInitializer::<TYPES>::load(
                                     TestInstanceState::new(self.async_delay_config.clone()),
                                     self.epoch_height,
+                                    self.epoch_start_block,
+                                    self.start_epoch_info.clone(),
                                     self.last_decided_leaf.clone(),
                                     (start_view, start_epoch),
                                     (high_qc, next_epoch_high_qc),