Remove associate_signers_with_stake from runner and improve epoch_service implementation

Author: sfauvel

mithril-aggregator/src/message_adapters/to_certificate_pending_message.rs

@@ -1,6 +1,6 @@
 use mithril_common::entities::CardanoDbBeacon;
 use mithril_common::{
-    entities::{CertificatePending, SignedEntityType, Signer},
+    entities::{CertificatePending, SignedEntityType},
     messages::{CertificatePendingMessage, SignerMessagePart},
 };
 

mithril-aggregator/src/message_adapters/to_epoch_settings_message.rs

@@ -1,4 +1,4 @@
-use mithril_common::entities::{EpochSettings, Signer};
+use mithril_common::entities::EpochSettings;
 use mithril_common::messages::{EpochSettingsMessage, SignerMessagePart, ToMessageAdapter};
 
 /// Adapter to spawn [EpochSettingsMessage] from [EpochSettings] instances.

mithril-signer/src/runtime/runner.rs

@@ -1,6 +1,6 @@
 use anyhow::Context;
 use async_trait::async_trait;
-use slog_scope::{debug, info, trace, warn};
+use slog_scope::{debug, info, warn};
 use thiserror::Error;
 
 #[cfg(test)]
@@ -55,14 +55,6 @@ pub trait Runner: Send + Sync {
     /// Register epoch information
     async fn inform_epoch_settings(&self, epoch_settings: EpochSettings) -> StdResult<()>;
 
-    /// From a list of signers, associate them with the stake read on the
-    /// Cardano node.
-    async fn associate_signers_with_stake(
-        &self,
-        epoch: Epoch,
-        signers: &[Signer],
-    ) -> StdResult<Vec<SignerWithStake>>;
-
     /// Create the message to be signed with the single signature.
     async fn compute_message(
         &self,
@@ -164,8 +156,6 @@ impl Runner for SignerRunner {
             .cloned()
     }
 
-    // TODO do we need to be async ?
-    // TODO do we return a vec or a &vec ?
     async fn get_next_signers(&self) -> StdResult<Vec<Signer>> {
         debug!("RUNNER: get_next_signers");
 
@@ -359,45 +349,6 @@ impl Runner for SignerRunner {
             .await
     }
 
-    // TODO do it in EpochService and store SignerWithStack
-    // Inject stake_store in EpochService (epoch + current protocol parameters)
-    async fn associate_signers_with_stake(
-        &self,
-        epoch: Epoch,
-        signers: &[Signer],
-    ) -> StdResult<Vec<SignerWithStake>> {
-        debug!("RUNNER: associate_signers_with_stake");
-
-        let stakes = self
-            .services
-            .stake_store
-            .get_stakes(epoch)
-            .await?
-            .ok_or_else(|| RunnerError::NoValueError(format!("stakes at epoch {epoch}")))?;
-        let mut signers_with_stake = vec![];
-
-        for signer in signers {
-            let stake = stakes
-                .get(&*signer.party_id)
-                .ok_or_else(|| RunnerError::NoStakeForSigner(signer.party_id.to_string()))?;
-
-            signers_with_stake.push(SignerWithStake::new(
-                signer.party_id.to_owned(),
-                signer.verification_key.to_owned(),
-                signer.verification_key_signature.to_owned(),
-                signer.operational_certificate.to_owned(),
-                signer.kes_period.to_owned(),
-                *stake,
-            ));
-            trace!(
-                " > associating signer_id {} with stake {}",
-                signer.party_id,
-                *stake
-            );
-        }
-        Ok(signers_with_stake)
-    }
-
     async fn compute_message(
         &self,
         signed_entity_type: &SignedEntityType,
@@ -551,7 +502,7 @@ mod tests {
             MKMap, MKMapNode, MKTreeNode, MKTreeStoreInMemory, MKTreeStorer, ProtocolInitializer,
         },
         digesters::{DumbImmutableDigester, DumbImmutableFileObserver},
-        entities::{BlockNumber, BlockRange, CardanoDbBeacon, Epoch, StakeDistribution},
+        entities::{BlockNumber, BlockRange, CardanoDbBeacon, Epoch},
         era::{adapters::EraReaderBootstrapAdapter, EraChecker, EraReader},
         signable_builder::{
             BlockRangeRootRetriever, CardanoImmutableFilesFullSignableBuilder,
@@ -857,36 +808,6 @@ mod tests {
         );
     }
 
-    #[tokio::test]
-    async fn test_associate_signers_with_stake() {
-        let services = init_services().await;
-        let stake_store = services.stake_store.clone();
-        let runner = init_runner(Some(services), None).await;
-        let epoch = Epoch(12);
-        let expected = fake_data::signers_with_stakes(5);
-        let signers = expected
-            .clone()
-            .into_iter()
-            .map(|s| s.into())
-            .collect::<Vec<_>>();
-        let stake_distribution = expected
-            .clone()
-            .iter()
-            .map(|s| s.into())
-            .collect::<StakeDistribution>();
-
-        stake_store
-            .save_stakes(epoch, stake_distribution)
-            .await
-            .expect("save_stakes should not fail");
-
-        let result = runner
-            .associate_signers_with_stake(epoch, &signers)
-            .await
-            .expect("associate_signers_with_stake should not fail");
-        assert_eq!(expected, result);
-    }
-
     #[tokio::test]
     async fn test_compute_message() {
         let mut services = init_services().await;

mithril-signer/src/runtime/state_machine.rs

@@ -295,7 +295,6 @@ impl StateMachine {
         })
     }
 
-    // TODO do we remove epoch_settings information when unregistered ???
     async fn transition_from_signed_to_unregistered(
         &self,
         epoch: Epoch,
@@ -518,7 +517,6 @@ mod tests {
             epoch: Epoch(3),
             protocol_parameters: fake_data::protocol_parameters(),
             next_protocol_parameters: fake_data::protocol_parameters(),
-            // TODO : put some data ?
             current_signers: vec![],
             next_signers: vec![],
         };
@@ -555,7 +553,6 @@ mod tests {
             .once()
             .returning(|| Ok(Some(fake_data::epoch_settings())));
 
-        // TODO do we check the epoch_setting is the one returning by get_epoch_settings
         runner
             .expect_inform_epoch_settings()
             .with(predicate::eq(fake_data::epoch_settings()))

mithril-signer/src/services/epoch_service.rs

@@ -80,8 +80,6 @@ impl MithrilEpochService {
         }
     }
 
-    // TODO do it in EpochService and store SignerWithStack
-    // Inject stake_store in EpochService (epoch + current protocol parameters)
     async fn associate_signers_with_stake(
         &self,
         epoch: Epoch,
@@ -162,24 +160,20 @@ impl EpochService for MithrilEpochService {
 
     async fn current_signers_with_stake(&self) -> StdResult<Vec<SignerWithStake>> {
         let current_epoch = self.epoch_of_current_data()?;
-        let (retrieval_epoch, _next_retrieval_epoch) = (
+        self.associate_signers_with_stake(
             current_epoch.offset_to_signer_retrieval_epoch()?,
-            current_epoch.offset_to_next_signer_retrieval_epoch(),
-        );
-        let current_signers = self.current_signers()?;
-        self.associate_signers_with_stake(retrieval_epoch, current_signers)
-            .await
+            self.current_signers()?,
+        )
+        .await
     }
 
     async fn next_signers_with_stake(&self) -> StdResult<Vec<SignerWithStake>> {
         let current_epoch = self.epoch_of_current_data()?;
-        let (_retrieval_epoch, _next_retrieval_epoch) = (
-            current_epoch.offset_to_signer_retrieval_epoch()?,
+        self.associate_signers_with_stake(
             current_epoch.offset_to_next_signer_retrieval_epoch(),
-        );
-        let next_signers = self.next_signers()?;
-        self.associate_signers_with_stake(_next_retrieval_epoch, next_signers)
-            .await
+            self.next_signers()?,
+        )
+        .await
     }
 }
 
@@ -191,12 +185,15 @@ mod tests {
 
     use mithril_common::{
         entities::{Epoch, EpochSettings, StakeDistribution},
-        test_utils::fake_data,
+        test_utils::fake_data::{self},
+    };
+    use mithril_persistence::store::{
+        adapter::{DumbStoreAdapter, MemoryAdapter},
+        StakeStore, StakeStorer,
     };
-    use mithril_persistence::store::{adapter::DumbStoreAdapter, StakeStore, StakeStorer};
 
     #[tokio::test]
-    async fn test_signers_are_not_available_before_register_epoch_settings_was_call() {
+    async fn test_retrieve_data_return_error_before_register_epoch_settings_was_call() {
         let epoch = Epoch(12);
         // Signers and stake distribution
         let signers = fake_data::signers(10);
@@ -215,12 +212,16 @@ mod tests {
 
         // Build service and register epoch settings
         let service = MithrilEpochService::new(stake_store);
+        assert!(service.epoch_of_current_data().is_err());
+        assert!(service.next_protocol_parameters().is_err());
         assert!(service.current_signers().is_err());
         assert!(service.next_signers().is_err());
+        assert!(service.current_signers_with_stake().await.is_err());
+        assert!(service.next_signers_with_stake().await.is_err());
     }
 
     #[tokio::test]
-    async fn test_signers_are_available_after_register_epoch_settings_call() {
+    async fn test_data_are_available_after_register_epoch_settings_call() {
         let epoch = Epoch(12);
         // Signers and stake distribution
         let signers = fake_data::signers(10);
@@ -245,15 +246,19 @@ mod tests {
             .unwrap();
 
         // Check current_signers
-        let current_signers = service.current_signers().unwrap();
-        let expected_current_signers = epoch_settings.current_signers.clone();
-        assert_eq!(expected_current_signers, *current_signers);
-
+        {
+            let current_signers = service.current_signers().unwrap();
+            let expected_current_signers = epoch_settings.current_signers.clone();
+            assert_eq!(expected_current_signers, *current_signers);
+        }
         // Check next_signers
-        let next_signers = service.next_signers().unwrap();
-        let expected_next_signers = epoch_settings.next_signers.clone();
-        assert_eq!(expected_next_signers, *next_signers);
+        {
+            let next_signers = service.next_signers().unwrap();
+            let expected_next_signers = epoch_settings.next_signers.clone();
+            assert_eq!(expected_next_signers, *next_signers);
+        }
 
+        // Check other data
         assert_eq!(
             epoch_settings.epoch,
             service.epoch_of_current_data().unwrap()
@@ -264,78 +269,75 @@ mod tests {
         );
     }
 
-    // TODO try to simplify this test
     #[tokio::test]
     async fn test_signers_with_stake_are_available_after_register_epoch_settings_call() {
+        fn build_stake_distribution(signers: &Vec<Signer>, first_stake: u64) -> StakeDistribution {
+            signers
+                .iter()
+                .enumerate()
+                .map(|(i, signer)| (signer.party_id.clone(), first_stake + i as u64))
+                .collect()
+        }
+
         let epoch = Epoch(12);
+
         // Signers and stake distribution
         let signers = fake_data::signers(10);
-
-        // Init stake_store
-        let stake_store = Arc::new(StakeStore::new(Box::new(DumbStoreAdapter::new()), None));
+        let stake_distribution: StakeDistribution = build_stake_distribution(&signers, 100);
+        let next_stake_distribution: StakeDistribution = build_stake_distribution(&signers, 500);
+
+        let stake_store = Arc::new(StakeStore::new(
+            Box::new(
+                MemoryAdapter::<Epoch, StakeDistribution>::new(Some(vec![
+                    (
+                        epoch.offset_to_signer_retrieval_epoch().unwrap(),
+                        stake_distribution.clone(),
+                    ),
+                    (
+                        epoch.offset_to_next_signer_retrieval_epoch(),
+                        next_stake_distribution.clone(),
+                    ),
+                ]))
+                .unwrap(),
+            ),
+            None,
+        ));
 
         // Epoch settings
-        let epoch_settings = fake_data::epoch_settings();
         let epoch_settings = EpochSettings {
             epoch,
             current_signers: signers[2..5].to_vec(),
             next_signers: signers[3..7].to_vec(),
-            ..epoch_settings.clone()
+            ..fake_data::epoch_settings().clone()
         };
+
         // Build service and register epoch settings
-        let mut service = MithrilEpochService::new(stake_store.clone());
+        let mut service = MithrilEpochService::new(stake_store);
         service
             .inform_epoch_settings(epoch_settings.clone())
             .await
             .unwrap();
 
-        // Check current_signers
-        let stake_distribution: StakeDistribution = signers
-            .iter()
-            .enumerate()
-            .map(|(i, signer)| (signer.party_id.clone(), (i + 1) as u64 * 100))
-            .collect();
-        stake_store
-            .save_stakes(
-                epoch.offset_to_signer_retrieval_epoch().unwrap(),
-                stake_distribution.clone(),
-            )
-            .await
-            .expect("save_stakes should not fail");
-        let current_signers = service.current_signers_with_stake().await.unwrap();
-        let expected_current_signers = epoch_settings.current_signers.clone();
-        assert_eq!(expected_current_signers.len(), current_signers.len());
-        for signer in current_signers {
-            assert_eq!(
-                stake_distribution.get(&signer.party_id).unwrap(),
-                &signer.stake
-            );
+        // Check current signers with stake
+        {
+            let current_signers = service.current_signers_with_stake().await.unwrap();
+
+            assert_eq!(epoch_settings.current_signers.len(), current_signers.len());
+            for signer in current_signers {
+                let expected_stake = stake_distribution.get(&signer.party_id).unwrap();
+                assert_eq!(expected_stake, &signer.stake);
+            }
         }
 
-        // Check next_signers
-        let next_stake_distribution: StakeDistribution = signers
-            .iter()
-            .enumerate()
-            .map(|(i, signer)| (signer.party_id.clone(), (i + 1) as u64 * 1000))
-            .collect();
-        stake_store
-            .save_stakes(
-                epoch.offset_to_next_signer_retrieval_epoch(),
-                next_stake_distribution.clone(),
-            )
-            .await
-            .expect("save_stakes should not fail");
-        let next_signers = service.next_signers_with_stake().await.unwrap();
-        let expected_next_signers = epoch_settings.next_signers.clone();
-        assert_eq!(expected_next_signers.len(), next_signers.len());
-        for signer in next_signers {
-            assert_eq!(
-                next_stake_distribution.get(&signer.party_id).unwrap(),
-                &signer.stake
-            );
+        // Check next signers with stake
+        {
+            let next_signers = service.next_signers_with_stake().await.unwrap();
+
+            assert_eq!(epoch_settings.next_signers.len(), next_signers.len());
+            for signer in next_signers {
+                let expected_stake = next_stake_distribution.get(&signer.party_id).unwrap();
+                assert_eq!(expected_stake, &signer.stake);
+            }
         }
     }
-
-    // TODO check update of signer after register_epoch_settings call
-    // TODO should we provide a "reset" function ?
 }