Store progress at DRB checkpoints (#3229)

Author: ss-es

Cargo.lock

@@ -24,6 +24,7 @@ hotshot = { workspace = true }
 hotshot-builder-api = { workspace = true }
 hotshot-builder-core = { workspace = true }
 hotshot-events-service = { workspace = true }
+hotshot-example-types = { workspace = true }
 hotshot-state-prover = { workspace = true }
 hotshot-types = { workspace = true }
 marketplace-builder-shared = { workspace = true }

builder/Cargo.toml

@@ -60,14 +60,15 @@ pub fn build_instance_state<V: Versions>(
         l1_client.clone(),
         chain_config,
     );
+
     let coordinator = EpochMembershipCoordinator::new(
         Arc::new(RwLock::new(EpochCommittees::new_stake(
             vec![],
             vec![],
             fetcher,
         ))),
-        None,
         100,
+        &Arc::new(sequencer::persistence::no_storage::NoStorage),
     );
 
     NodeState::new(

builder/src/non_permissioned.rs

@@ -63,7 +63,6 @@ use hotshot_types::{
         network::ConnectedNetwork,
         node_implementation::{ConsensusTime, NodeType, Versions},
         states::TestableState,
-        storage::storage_add_drb_result,
     },
     utils::genesis_epoch_from_version,
     HotShotConfig, PeerConfig, ValidatorConfig,
@@ -388,7 +387,6 @@ pub trait RunDa<
             fallback_builder_url: config.config.builder_urls.first().clone(),
         };
         let epoch_height = config.config.epoch_height;
-
         let storage = TestStorage::<TYPES>::default();
 
         SystemContext::init(
@@ -397,11 +395,7 @@ pub trait RunDa<
             state_sk,
             config.node_index,
             config.config,
-            EpochMembershipCoordinator::new(
-                membership,
-                Some(storage_add_drb_result(storage.clone())),
-                epoch_height,
-            ),
+            EpochMembershipCoordinator::new(membership, epoch_height, &storage.clone()),
             Arc::from(network),
             initializer,
             ConsensusMetricsValue::default(),

crates/hotshot/examples/infra/mod.rs

@@ -18,7 +18,7 @@ use hotshot_task::dependency::{Dependency, EventDependency};
 use hotshot_types::{
     consensus::OuterConsensus,
     data::{Leaf2, QuorumProposalWrapper, VidDisperseShare, ViewChangeEvidence2},
-    drb::{DrbResult, DrbSeedInput},
+    drb::{DrbInput, DrbResult, DrbSeedInput},
     epoch_membership::EpochMembershipCoordinator,
     event::{Event, EventType, LeafInfo},
     message::{Proposal, UpgradeLock},
@@ -34,7 +34,7 @@ use hotshot_types::{
         election::Membership,
         node_implementation::{ConsensusTime, NodeImplementation, NodeType, Versions},
         signature_key::{SignatureKey, StakeTableEntryType, StateSignatureKey},
-        storage::Storage,
+        storage::{store_drb_progress_fn, Storage},
         BlockPayload, ValidatedState,
     },
     utils::{
@@ -177,10 +177,16 @@ fn start_drb_task<TYPES: NodeType, I: NodeImplementation<TYPES>>(
 ) {
     let membership = membership.clone();
     let storage = storage.clone();
+    let store_drb_progress_fn = store_drb_progress_fn(storage.clone());
     let consensus = consensus.clone();
+    let drb_input = DrbInput {
+        epoch: *epoch,
+        iteration: 0,
+        value: seed,
+    };
     tokio::spawn(async move {
         let drb_result = tokio::task::spawn_blocking(move || {
-            hotshot_types::drb::compute_drb_result::<TYPES>(seed)
+            hotshot_types::drb::compute_drb_result(drb_input, store_drb_progress_fn)
         })
         .await
         .unwrap();

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

@@ -36,7 +36,6 @@ use hotshot_types::{
     traits::{
         election::Membership,
         node_implementation::{NodeType, Versions},
-        storage::storage_add_drb_result,
         EncodeBytes,
     },
     utils::{option_epoch_from_block_number, View, ViewInner},
@@ -129,11 +128,8 @@ pub async fn build_system_handle_from_launcher<
         hotshot_config.known_da_nodes.clone(),
     )));
 
-    let coordinator = EpochMembershipCoordinator::new(
-        memberships,
-        Some(storage_add_drb_result(storage.clone())),
-        hotshot_config.epoch_height,
-    );
+    let coordinator =
+        EpochMembershipCoordinator::new(memberships, hotshot_config.epoch_height, &storage);
     let node_key_map = launcher.metadata.build_node_key_map();
 
     let (c, s, r) = SystemContext::init(

crates/hotshot/testing/src/helpers.rs

@@ -21,10 +21,7 @@ use hotshot_example_types::{
 use hotshot_types::{
     consensus::ConsensusMetricsValue,
     epoch_membership::EpochMembershipCoordinator,
-    traits::{
-        node_implementation::{NodeType, Versions},
-        storage::storage_add_drb_result,
-    },
+    traits::node_implementation::{NodeType, Versions},
     HotShotConfig, PeerConfig, ValidatorConfig,
 };
 use hotshot_utils::anytrace::*;
@@ -262,11 +259,8 @@ pub async fn create_test_handle<
     let private_key = validator_config.private_key.clone();
     let public_key = validator_config.public_key.clone();
     let state_private_key = validator_config.state_private_key.clone();
-    let membership_coordinator = EpochMembershipCoordinator::new(
-        memberships,
-        Some(storage_add_drb_result(storage.clone())),
-        config.epoch_height,
-    );
+    let membership_coordinator =
+        EpochMembershipCoordinator::new(memberships, config.epoch_height, &storage.clone());
 
     let behaviour = (metadata.behaviour)(node_id);
     match behaviour {

crates/hotshot/testing/src/test_builder.rs

@@ -39,7 +39,6 @@ use hotshot_types::{
         election::Membership,
         network::ConnectedNetwork,
         node_implementation::{ConsensusTime, NodeImplementation, NodeType, Versions},
-        storage::storage_add_drb_result,
     },
     HotShotConfig, ValidatorConfig,
 };
@@ -615,8 +614,8 @@ where
             config,
             EpochMembershipCoordinator::new(
                 Arc::new(RwLock::new(memberships)),
-                Some(storage_add_drb_result(storage.clone())),
                 epoch_height,
+                &storage.clone(),
             ),
             network,
             initializer,
@@ -658,11 +657,7 @@ where
             state_private_key,
             node_id,
             config,
-            EpochMembershipCoordinator::new(
-                memberships,
-                Some(storage_add_drb_result(storage.clone())),
-                epoch_height,
-            ),
+            EpochMembershipCoordinator::new(memberships, epoch_height, &storage.clone()),
             network,
             initializer,
             ConsensusMetricsValue::default(),

crates/hotshot/testing/src/test_runner.rs

@@ -22,7 +22,7 @@ use hotshot_types::{
     message::UpgradeLock,
     traits::{
         election::Membership,
-        node_implementation::{ConsensusTime, NodeType},storage::storage_add_drb_result
+        node_implementation::{ConsensusTime, NodeType}
     },
 };
 use tokio::time::timeout;
@@ -62,7 +62,7 @@ async fn test_network_task() {
         all_nodes.clone(),
         all_nodes,
     )));
-    let coordinator = EpochMembershipCoordinator::new(membership, Some(storage_add_drb_result(storage.clone())),config.epoch_height);
+    let coordinator = EpochMembershipCoordinator::new(membership, config.epoch_height, &storage.clone());
     let network_state: NetworkEventTaskState<TestTypes, TestVersions, MemoryNetwork<_>, _> =
         NetworkEventTaskState {
             id: node_id,
@@ -238,7 +238,7 @@ async fn test_network_storage_fail() {
         all_nodes.clone(),
         all_nodes,
     )));
-    let coordinator = EpochMembershipCoordinator::new(membership, Some(storage_add_drb_result(storage.clone())),config.epoch_height);
+    let coordinator = 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/network_task.rs

@@ -0,0 +1,51 @@
+use sha2::{Sha256, Digest};
+use hotshot_types::traits::storage::null_store_drb_progress_fn;
+use hotshot_types::drb::compute_drb_result;
+use hotshot_types::drb::DIFFICULTY_LEVEL;
+use hotshot_types::drb::DrbInput;
+
+#[cfg(test)]
+#[tokio::test(flavor = "multi_thread")]
+async fn test_compute_drb_result() {
+    let drb_input = DrbInput {
+        epoch: 0,
+        iteration: 0,
+        value: [0u8; 32],
+    };
+
+    let mut expected_result = [0u8; 32];
+    {
+    let mut hash = drb_input.value.to_vec().clone();
+        for _ in 0..DIFFICULTY_LEVEL {
+            hash = Sha256::digest(hash).to_vec();
+        }
+    expected_result.copy_from_slice(&hash);
+    }
+
+    let actual_result = compute_drb_result(drb_input, null_store_drb_progress_fn());
+
+    assert_eq!(expected_result, actual_result);
+}
+
+#[cfg(test)]
+#[tokio::test(flavor = "multi_thread")]
+async fn test_compute_drb_result_2() {
+    let drb_input = DrbInput {
+        epoch: 0,
+        iteration: 2,
+        value: [0u8; 32],
+    };
+
+    let mut expected_result = [0u8; 32];
+    {
+    let mut hash = drb_input.value.to_vec().clone();
+        for _ in 2..DIFFICULTY_LEVEL {
+            hash = Sha256::digest(hash).to_vec();
+        }
+    expected_result.copy_from_slice(&hash);
+    }
+
+    let actual_result = compute_drb_result(drb_input, null_store_drb_progress_fn());
+
+    assert_eq!(expected_result, actual_result);
+}