Load membership from storage in MembershipCoordinator before triggering catchup

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

@@ -81,6 +81,7 @@ impl<
 {
     type Error = anyhow::Error;
     type StakeTableHash = NoStakeTableHash;
+    type StakeTableMetadata = ();
     type Storage = TestStorage<TYPES>;
 
     fn new<I: NodeImplementation<TYPES>>(
@@ -372,4 +373,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,

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

@@ -1,8 +1,9 @@
 use sha2::{Sha256, Digest};
-use hotshot_types::traits::storage::null_store_drb_progress_fn;
-use hotshot_types::traits::storage::null_load_drb_progress_fn;
 use hotshot_types::drb::compute_drb_result;
 use hotshot_types::drb::DrbInput;
+use hotshot_example_types::storage_types::TestStorage;
+use hotshot_example_types::node_types::TestTypes;
+use std::sync::Arc;
 
 #[cfg(test)]
 #[tokio::test(flavor = "multi_thread")]
@@ -25,7 +26,8 @@ async fn test_compute_drb_result() {
     expected_result.copy_from_slice(&hash);
     }
 
-    let actual_result = compute_drb_result(drb_input, null_store_drb_progress_fn(), null_load_drb_progress_fn()).await;
+    let storage = Arc::new(TestStorage::<TestTypes>::default());
+    let actual_result = compute_drb_result(drb_input, storage).await;
 
     assert_eq!(expected_result, actual_result);
 }
@@ -50,7 +52,8 @@ async fn test_compute_drb_result_2() {
     expected_result.copy_from_slice(&hash);
     }
 
-    let actual_result = compute_drb_result(drb_input, null_store_drb_progress_fn(), null_load_drb_progress_fn()).await;
+    let storage = Arc::new(TestStorage::<TestTypes>::default());
+    let actual_result = compute_drb_result(drb_input, storage).await;
 
     assert_eq!(expected_result, actual_result);
 }

crates/hotshot/types/src/drb.rs

@@ -14,7 +14,7 @@ use vbs::version::{StaticVersionType, Version};
 use crate::{
     traits::{
         node_implementation::{ConsensusTime, NodeType, Versions},
-        storage::{LoadDrbProgressFn, StoreDrbProgressFn},
+        storage::EpochStateStorage,
     },
     HotShotConfig,
 };
@@ -98,15 +98,14 @@ pub fn difficulty_level() -> u64 {
 /// # Arguments
 /// * `drb_seed_input` - Serialized QC signature.
 #[must_use]
-pub async fn compute_drb_result(
+pub async fn compute_drb_result<TYPES: NodeType>(
     drb_input: DrbInput,
-    store_drb_progress: StoreDrbProgressFn,
-    load_drb_progress: LoadDrbProgressFn,
+    storage: Arc<dyn EpochStateStorage<TYPES>>,
 ) -> DrbResult {
     tracing::warn!("Beginning DRB calculation with input {:?}", drb_input);
     let mut drb_input = drb_input;
 
-    if let Ok(loaded_drb_input) = load_drb_progress(drb_input.epoch).await {
+    if let Ok(loaded_drb_input) = storage.load_drb_input(drb_input.epoch).await {
         if loaded_drb_input.difficulty_level != drb_input.difficulty_level {
             tracing::error!(
                 "We are calculating the DRB result with input {drb_input:?}, but we had \
@@ -165,7 +164,7 @@ pub async fn compute_drb_result(
 
         let elapsed_time = last_time.elapsed().as_millis();
 
-        let store_drb_progress = store_drb_progress.clone();
+        let storage = Arc::clone(&storage);
         tokio::spawn(async move {
             tracing::warn!(
                 "Storing partial DRB progress: {:?}. Time elapsed since the previous iteration of \
@@ -174,7 +173,7 @@ pub async fn compute_drb_result(
                 last_iteration,
                 elapsed_time
             );
-            if let Err(e) = store_drb_progress(updated_drb_input).await {
+            if let Err(e) = storage.store_drb_input(updated_drb_input).await {
                 tracing::warn!("Failed to store DRB progress during calculation: {}", e);
             }
         });
@@ -212,10 +211,9 @@ pub async fn compute_drb_result(
 
     tracing::warn!("Completed DRB calculation. Result: {:?}", final_drb_input);
 
-    let store_drb_progress = store_drb_progress.clone();
     tokio::spawn(async move {
-        if let Err(e) = store_drb_progress(final_drb_input).await {
-            tracing::warn!("Failed to store DRB progress during calculation: {}", e);
+        if let Err(e) = storage.store_drb_input(final_drb_input).await {
+            tracing::warn!("Failed to store final DRB result: {}", e);
         }
     });