load membership from storage before catchup

Author: imabdulbasit

crates/builder/src/non_permissioned.rs

@@ -72,7 +72,7 @@ pub fn build_instance_state<V: Versions>(
             genesis.epoch_height.unwrap_or_default(),
         ))),
         genesis.epoch_height.unwrap_or_default(),
-        &Arc::new(sequencer::persistence::no_storage::NoStorage),
+        Arc::new(sequencer::persistence::no_storage::NoStorage),
     );
 
     NodeState::new(

crates/hotshot/example-types/src/membership/strict_membership.rs

@@ -80,6 +80,7 @@ impl<
 {
     type Error = anyhow::Error;
     type StakeTableHash = NoStakeTableHash;
+    type StakeTableMetadata = ();
     type Storage = TestStorage<TYPES>;
 
     fn new<I: NodeImplementation<TYPES>>(
@@ -368,4 +369,14 @@ impl<
         self.inner
             .add_da_committee(first_epoch, committee.into_iter().map(Into::into).collect());
     }
+
+    fn insert_epoch_state(
+        &mut self,
+        epoch: TYPES::Epoch,
+        _stake_table: HSStakeTable<TYPES>,
+        _epoch_root: Option<TYPES::BlockHeader>,
+        _metadata: Self::StakeTableMetadata,
+    ) {
+        self.epochs.insert(epoch);
+    }
 }

crates/hotshot/example-types/src/storage_types.rs

@@ -12,7 +12,7 @@ use std::{
     },
 };
 
-use anyhow::{anyhow, bail, Result};
+use anyhow::{anyhow, bail, ensure, Result};
 use async_lock::RwLock;
 use async_trait::async_trait;
 use hotshot_types::{
@@ -28,9 +28,10 @@ use hotshot_types::{
         LightClientStateUpdateCertificateV2, NextEpochQuorumCertificate2, QuorumCertificate2,
         UpgradeCertificate,
     },
+    stake_table::HSStakeTable,
     traits::{
         node_implementation::{ConsensusTime, NodeType},
-        storage::Storage,
+        storage::{EpochStateStorage, Storage},
     },
     vote::HasViewNumber,
 };
@@ -174,6 +175,8 @@ impl<TYPES: NodeType> TestStorage<TYPES> {
 
 #[async_trait]
 impl<TYPES: NodeType> Storage<TYPES> for TestStorage<TYPES> {
+    type StakeTableMetadata = ();
+
     async fn append_vid(&self, proposal: &Proposal<TYPES, ADVZDisperseShare<TYPES>>) -> Result<()> {
         if self.should_return_err.load(Ordering::Relaxed) {
             bail!("Failed to append VID proposal to storage");
@@ -426,41 +429,60 @@ impl<TYPES: NodeType> Storage<TYPES> for TestStorage<TYPES> {
 
         Ok(())
     }
+}
 
-    async fn store_drb_result(&self, epoch: TYPES::Epoch, drb_result: DrbResult) -> Result<()> {
-        let mut inner = self.inner.write().await;
-
-        inner.drb_results.insert(epoch, drb_result);
-
-        Ok(())
+#[async_trait]
+impl<TYPES: NodeType> EpochStateStorage<TYPES> for TestStorage<TYPES> {
+    async fn load_epoch_root(&self, epoch: TYPES::Epoch) -> Result<Option<TYPES::BlockHeader>> {
+        let inner = self.inner.read().await;
+        Ok(inner.epoch_roots.get(&epoch).cloned())
     }
 
-    async fn store_epoch_root(
+    async fn load_stake_table(
         &self,
-        epoch: TYPES::Epoch,
-        block_header: TYPES::BlockHeader,
-    ) -> Result<()> {
-        let mut inner = self.inner.write().await;
-
-        inner.epoch_roots.insert(epoch, block_header);
+        _epoch: TYPES::Epoch,
+    ) -> Result<Option<(HSStakeTable<TYPES>, Box<dyn std::any::Any + Send + Sync>)>> {
+        Ok(None)
+    }
 
-        Ok(())
+    async fn load_drb_result(&self, epoch: TYPES::Epoch) -> Result<DrbResult> {
+        match self.load_drb_input(*epoch).await {
+            Ok(drb_input) => {
+                ensure!(drb_input.iteration == drb_input.difficulty_level);
+                Ok(drb_input.value)
+            },
+            Err(e) => Err(e),
+        }
     }
 
     async fn store_drb_input(&self, drb_input: DrbInput) -> Result<()> {
         let mut inner = self.inner.write().await;
-
         inner.drb_inputs.insert(drb_input.epoch, drb_input);
-
         Ok(())
     }
 
     async fn load_drb_input(&self, epoch: u64) -> Result<DrbInput> {
         let inner = self.inner.read().await;
+        inner
+            .drb_inputs
+            .get(&epoch)
+            .cloned()
+            .ok_or_else(|| anyhow!("DRB input not found for epoch {}", epoch))
+    }
 
-        match inner.drb_inputs.get(&epoch) {
-            Some(drb_input) => Ok(drb_input.clone()),
-            None => Err(anyhow!("Missing DrbInput for epoch {}", epoch)),
-        }
+    async fn store_drb_result(&self, epoch: TYPES::Epoch, drb_result: DrbResult) -> Result<()> {
+        let mut inner = self.inner.write().await;
+        inner.drb_results.insert(epoch, drb_result);
+        Ok(())
+    }
+
+    async fn store_epoch_root(
+        &self,
+        epoch: TYPES::Epoch,
+        block_header: TYPES::BlockHeader,
+    ) -> Result<()> {
+        let mut inner = self.inner.write().await;
+        inner.epoch_roots.insert(epoch, block_header);
+        Ok(())
     }
 }

crates/hotshot/examples/infra/mod.rs

@@ -401,7 +401,7 @@ pub trait RunDa<
             state_sk,
             config.node_index,
             config.config,
-            EpochMembershipCoordinator::new(membership, epoch_height, &storage.clone()),
+            EpochMembershipCoordinator::new(membership, epoch_height, storage.clone()),
             network,
             initializer,
             ConsensusMetricsValue::default(),

crates/hotshot/hotshot/src/lib.rs

@@ -19,7 +19,8 @@ use hotshot_types::{
     simple_certificate::{CertificatePair, LightClientStateUpdateCertificateV2},
     traits::{
         block_contents::BlockHeader, election::Membership, network::BroadcastDelay,
-        node_implementation::Versions, signature_key::StateSignatureKey, storage::Storage,
+        node_implementation::Versions, signature_key::StateSignatureKey,
+        storage::EpochStateStorage,
     },
     utils::{epoch_from_block_number, is_ge_epoch_root},
 };

crates/hotshot/task-impls/src/helpers.rs

@@ -42,7 +42,7 @@ use hotshot_types::{
         signature_key::{
             LCV2StateSignatureKey, LCV3StateSignatureKey, SignatureKey, StakeTableEntryType,
         },
-        storage::Storage,
+        storage::{EpochStateStorage, Storage},
         BlockPayload, ValidatedState,
     },
     utils::{

crates/hotshot/testing/src/helpers.rs

@@ -142,7 +142,7 @@ where
     )));
 
     let coordinator =
-        EpochMembershipCoordinator::new(memberships, hotshot_config.epoch_height, &storage);
+        EpochMembershipCoordinator::new(memberships, hotshot_config.epoch_height, storage.clone());
     let node_key_map = launcher.metadata.build_node_key_map();
 
     let (c, s, r) = SystemContext::init(

crates/hotshot/testing/src/test_builder.rs

@@ -283,7 +283,7 @@ pub async fn create_test_handle<
     let public_key = validator_config.public_key.clone();
     let state_private_key = validator_config.state_private_key.clone();
     let membership_coordinator =
-        EpochMembershipCoordinator::new(memberships, config.epoch_height, &storage.clone());
+        EpochMembershipCoordinator::new(memberships, config.epoch_height, storage.clone());
 
     let behaviour = (metadata.behaviour)(node_id);
     match behaviour {

crates/hotshot/testing/src/test_runner.rs

@@ -561,7 +561,7 @@ where
             EpochMembershipCoordinator::new(
                 Arc::new(RwLock::new(memberships)),
                 epoch_height,
-                &storage.clone(),
+                storage.clone(),
             ),
             network,
             initializer,
@@ -602,7 +602,7 @@ where
             state_private_key,
             node_id,
             config,
-            EpochMembershipCoordinator::new(memberships, epoch_height, &storage.clone()),
+            EpochMembershipCoordinator::new(memberships, epoch_height, storage.clone()),
             network,
             initializer,
             ConsensusMetricsValue::default(),

crates/hotshot/testing/tests/tests_1/network_task.rs

@@ -70,7 +70,7 @@ async fn test_network_task() {
         config.epoch_height,
     )));
     let coordinator =
-        EpochMembershipCoordinator::new(membership, config.epoch_height, &storage.clone());
+        EpochMembershipCoordinator::new(membership, config.epoch_height, storage.clone());
     let network_state: NetworkEventTaskState<TestTypes, TestVersions, MemoryNetwork<_>, _> =
         NetworkEventTaskState {
             id: node_id,
@@ -250,7 +250,7 @@ async fn test_network_storage_fail() {
         config.epoch_height,
     )));
     let coordinator =
-        EpochMembershipCoordinator::new(membership, config.epoch_height, &storage.clone());
+        EpochMembershipCoordinator::new(membership, config.epoch_height, storage.clone());
     let network_state: NetworkEventTaskState<TestTypes, TestVersions, MemoryNetwork<_>, _> =
         NetworkEventTaskState {
             id: node_id,