refactor(aggregator): slave synchronizer updates the next signers in the epoch service

Instead of having to repeat the call on 'inform_new_epoch' in the state machine for the slave.

Author: jpraynaud

mithril-aggregator/src/dependency_injection/builder/protocol/certificates.rs

@@ -161,6 +161,7 @@ impl DependenciesBuilder {
         &mut self,
     ) -> Result<Arc<MithrilSignerRegistrationSlave>> {
         let registerer = MithrilSignerRegistrationSlave::new(
+            self.get_epoch_service().await?,
             self.get_verification_key_store().await?,
             self.get_signer_store().await?,
             self.get_signer_registration_verifier().await?,

mithril-aggregator/src/runtime/state_machine.rs

@@ -274,8 +274,6 @@ impl AggregatorRuntime {
                 self.runner
                     .synchronize_slave_aggregator_signer_registration()
                     .await?;
-                // Needed to recompute epoch data for the next signing round on the slave
-                self.runner.inform_new_epoch(new_time_point.epoch).await?;
             }
             self.runner.precompute_epoch_data().await?;
         }
@@ -940,7 +938,7 @@ mod tests {
             runner
                 .expect_inform_new_epoch()
                 .with(predicate::eq(new_time_point_clone.clone().epoch))
-                .times(2)
+                .once()
                 .returning(|_| Ok(()));
             runner
                 .expect_update_epoch_settings()

mithril-aggregator/src/services/epoch_service.rs

@@ -49,6 +49,9 @@ pub trait EpochService: Sync + Send {
     /// Note: must be called after `inform_epoch`.
     async fn update_epoch_settings(&mut self) -> StdResult<()>;
 
+    /// Update the next signers with stake for the next epoch.
+    async fn update_next_signers_with_stake(&mut self) -> StdResult<()>;
+
     /// Inform the service that it can precompute data for its current epoch.
     ///
     /// Note: must be called after `inform_epoch`.
@@ -383,6 +386,27 @@ impl EpochService for MithrilEpochService {
         self.insert_future_epoch_settings(data.epoch).await
     }
 
+    async fn update_next_signers_with_stake(&mut self) -> StdResult<()> {
+        debug!(self.logger, ">> update_next_signers_with_stake");
+
+        let data = self.unwrap_data().with_context(|| {
+            "can't update next signers with stake if inform_epoch has not been called first"
+        })?;
+
+        let next_signer_retrieval_epoch = data.epoch.offset_to_next_signer_retrieval_epoch();
+        let next_signers_with_stake = self
+            .get_signers_with_stake_at_epoch(next_signer_retrieval_epoch)
+            .await?;
+
+        self.epoch_data.as_mut().unwrap().next_signers_with_stake = next_signers_with_stake;
+
+        self.precompute_epoch_data()
+            .await
+            .with_context(|| "Epoch service failed to precompute epoch data")?;
+
+        Ok(())
+    }
+
     async fn precompute_epoch_data(&mut self) -> StdResult<()> {
         debug!(self.logger, ">> precompute_epoch_data");
 
@@ -523,6 +547,7 @@ pub(crate) struct FakeEpochService {
     inform_epoch_error: bool,
     update_epoch_settings_error: bool,
     precompute_epoch_data_error: bool,
+    update_next_signers_with_stake_error: bool,
 }
 
 #[cfg(test)]
@@ -615,6 +640,7 @@ impl FakeEpochServiceBuilder {
             inform_epoch_error: false,
             update_epoch_settings_error: false,
             precompute_epoch_data_error: false,
+            update_next_signers_with_stake_error: false,
         }
     }
 }
@@ -660,19 +686,21 @@ impl FakeEpochService {
             inform_epoch_error: false,
             update_epoch_settings_error: false,
             precompute_epoch_data_error: false,
+            update_next_signers_with_stake_error: false,
         }
     }
 
-    #[allow(dead_code)]
     pub fn toggle_errors(
         &mut self,
         inform_epoch: bool,
         update_protocol_parameters: bool,
         precompute_epoch: bool,
+        update_next_signers_with_stake: bool,
     ) {
         self.inform_epoch_error = inform_epoch;
         self.update_epoch_settings_error = update_protocol_parameters;
         self.precompute_epoch_data_error = precompute_epoch;
+        self.update_next_signers_with_stake_error = update_next_signers_with_stake;
     }
 
     fn unwrap_data(&self) -> Result<&EpochData, EpochServiceError> {
@@ -714,6 +742,13 @@ impl EpochService for FakeEpochService {
         Ok(())
     }
 
+    async fn update_next_signers_with_stake(&mut self) -> StdResult<()> {
+        if self.update_next_signers_with_stake_error {
+            anyhow::bail!("update_next_signers_with_stake fake error");
+        }
+        Ok(())
+    }
+
     fn cardano_era(&self) -> StdResult<CardanoEra> {
         Ok(self.unwrap_data()?.cardano_era.clone())
     }
@@ -1213,6 +1248,48 @@ mod tests {
         assert!(service.computed_epoch_data.is_none());
     }
 
+    #[tokio::test]
+    async fn update_next_signers_with_stake_succeeds() {
+        let fixture = MithrilFixtureBuilder::default().with_signers(3).build();
+        let next_fixture = MithrilFixtureBuilder::default().with_signers(5).build();
+        let next_avk = next_fixture.compute_avk();
+        let epoch = Epoch(4);
+        let mut service = EpochServiceBuilder {
+            next_signers_with_stake: next_fixture.signers_with_stake().clone(),
+            ..EpochServiceBuilder::new(epoch, fixture.clone())
+        }
+        .build()
+        .await;
+        service
+            .inform_epoch(epoch)
+            .await
+            .expect("inform_epoch should not fail");
+        service.epoch_data = Some(EpochData {
+            next_signers_with_stake: vec![],
+            ..service.epoch_data.unwrap()
+        });
+        service.computed_epoch_data = None;
+
+        service
+            .update_next_signers_with_stake()
+            .await
+            .expect("update_next_signers_with_stake should not fail");
+
+        let expected_next_signers_with_stake = next_fixture.signers_with_stake();
+        assert_eq!(
+            expected_next_signers_with_stake,
+            service.epoch_data.unwrap().next_signers_with_stake
+        );
+
+        assert_eq!(
+            next_avk,
+            service
+                .computed_epoch_data
+                .unwrap()
+                .next_aggregate_verification_key
+        );
+    }
+
     #[tokio::test]
     async fn update_epoch_settings_insert_future_epoch_settings_in_the_store() {
         let future_protocol_parameters = ProtocolParameters::new(6, 89, 0.124);

mithril-aggregator/src/services/signer_registration/error.rs

@@ -31,9 +31,13 @@ pub enum SignerRegistrationError {
     #[error("signer already registered")]
     ExistingSigner(Box<SignerWithStake>),
 
-    /// Store error.
+    /// Store.
     #[error("store error")]
-    StoreError(#[source] StdError),
+    Store(#[source] StdError),
+
+    /// Epoch service.
+    #[error("epoch service error")]
+    EpochService(#[source] StdError),
 
     /// Signer registration failed.
     #[error("signer registration failed")]

mithril-aggregator/src/services/signer_registration/master.rs

@@ -147,7 +147,7 @@ impl SignerRegisterer for MithrilSignerRegistrationMaster {
                     registration_round.epoch
                 )
             })
-            .map_err(|e| SignerRegistrationError::StoreError(anyhow!(e)))?
+            .map_err(|e| SignerRegistrationError::Store(anyhow!(e)))?
         {
             Some(_) => Err(SignerRegistrationError::ExistingSigner(Box::new(
                 signer_save,

mithril-aggregator/src/services/signer_registration/slave.rs

@@ -10,6 +10,7 @@ use mithril_common::{
 use mithril_persistence::store::StakeStorer;
 
 use crate::{
+    dependency_injection::EpochServiceWrapper,
     services::{AggregatorClient, EpochPruningTask},
     SignerRegistrationVerifier, VerificationKeyStorer,
 };
@@ -21,6 +22,9 @@ use super::{
 
 /// A [MithrilSignerRegistrationSlave] supports signer registrations in a slave aggregator
 pub struct MithrilSignerRegistrationSlave {
+    /// Epoch service
+    pub epoch_service: EpochServiceWrapper,
+
     /// Verification key store
     verification_key_store: Arc<dyn VerificationKeyStorer>,
 
@@ -44,6 +48,7 @@ pub struct MithrilSignerRegistrationSlave {
 impl MithrilSignerRegistrationSlave {
     /// MithrilSignerRegistererSlave factory
     pub fn new(
+        epoch_service: EpochServiceWrapper,
         verification_key_store: Arc<dyn VerificationKeyStorer>,
         signer_recorder: Arc<dyn SignerRecorder>,
         signer_registration_verifier: Arc<dyn SignerRegistrationVerifier>,
@@ -52,6 +57,7 @@ impl MithrilSignerRegistrationSlave {
         verification_key_epoch_retention_limit: Option<u64>,
     ) -> Self {
         Self {
+            epoch_service,
             verification_key_store,
             signer_recorder,
             signer_registration_verifier,
@@ -90,9 +96,16 @@ impl MithrilSignerRegistrationSlave {
                         epoch
                     )
                 })
-                .map_err(|e| SignerRegistrationError::StoreError(anyhow!(e)))?;
+                .map_err(|e| SignerRegistrationError::Store(anyhow!(e)))?;
         }
 
+        self.epoch_service
+            .write()
+            .await
+            .update_next_signers_with_stake()
+            .await
+            .map_err(|e| SignerRegistrationError::EpochService(anyhow!(e)))?;
+
         Ok(())
     }
 }
@@ -130,8 +143,8 @@ impl SignerSynchronizer for MithrilSignerRegistrationSlave {
             .get_stakes(registration_epoch)
             .await
             .with_context(|| "synchronize_all_signers failed")
-            .map_err(SignerRegistrationError::StoreError)?
-            .ok_or(SignerRegistrationError::StoreError(anyhow::anyhow!(
+            .map_err(SignerRegistrationError::Store)?
+            .ok_or(SignerRegistrationError::Store(anyhow::anyhow!(
                 "Slave aggregator did not return any stake distribution"
             )))?;
         self.synchronize_signers(registration_epoch, &next_signers, &stake_distribution)
@@ -214,8 +227,9 @@ mod tests {
         database::{repository::SignerRegistrationStore, test_helper::main_db_connection},
         message_adapters::FromEpochSettingsAdapter,
         services::{
-            AggregatorClient, AggregatorClientError, EpochPruningTask, MockAggregatorClient,
-            MockSignerRecorder, MockSignerRegistrationVerifier, SignerSynchronizer,
+            AggregatorClient, AggregatorClientError, EpochPruningTask, FakeEpochService,
+            MockAggregatorClient, MockSignerRecorder, MockSignerRegistrationVerifier,
+            SignerSynchronizer,
         },
         store::MockVerificationKeyStorer,
         tools::mocks::MockStakeStore,
@@ -226,10 +240,15 @@ mod tests {
     use test_utils::*;
 
     mod test_utils {
+        use tokio::sync::RwLock;
+
+        use crate::{dependency_injection::EpochServiceWrapper, services::FakeEpochService};
+
         use super::*;
 
         /// MithrilSignerRegistrationSlaveBuilder is a test builder for [MithrilSignerRegistrationSlave]
         pub struct MithrilSignerRegistrationSlaveBuilder {
+            epoch_service: EpochServiceWrapper,
             signer_recorder: Arc<dyn SignerRecorder>,
             signer_registration_verifier: Arc<dyn SignerRegistrationVerifier>,
             master_aggregator_client: Arc<dyn AggregatorClient>,
@@ -241,6 +260,7 @@ mod tests {
         impl Default for MithrilSignerRegistrationSlaveBuilder {
             fn default() -> Self {
                 Self {
+                    epoch_service: Arc::new(RwLock::new(FakeEpochService::without_data())),
                     signer_recorder: Arc::new(MockSignerRecorder::new()),
                     signer_registration_verifier: Arc::new(MockSignerRegistrationVerifier::new()),
                     master_aggregator_client: Arc::new(MockAggregatorClient::new()),
@@ -254,6 +274,13 @@ mod tests {
         }
 
         impl MithrilSignerRegistrationSlaveBuilder {
+            pub fn with_epoch_service(self, epoch_service: FakeEpochService) -> Self {
+                Self {
+                    epoch_service: Arc::new(RwLock::new(epoch_service)),
+                    ..self
+                }
+            }
+
             pub fn with_verification_key_store(
                 self,
                 verification_key_store: Arc<dyn VerificationKeyStorer>,
@@ -310,6 +337,7 @@ mod tests {
 
             pub fn build(self) -> MithrilSignerRegistrationSlave {
                 MithrilSignerRegistrationSlave {
+                    epoch_service: self.epoch_service,
                     verification_key_store: self.verification_key_store,
                     signer_recorder: self.signer_recorder,
                     signer_registration_verifier: self.signer_registration_verifier,
@@ -478,6 +506,73 @@ mod tests {
             .expect_err("synchronize_all_signers should fail");
     }
 
+    #[tokio::test]
+    async fn synchronize_all_signers_fails_if_epoch_service_update_next_signers_fails() {
+        let registration_epoch = Epoch(1);
+        let fixture = MithrilFixtureBuilder::default()
+            .with_signers(5)
+            .disable_signers_certification()
+            .build();
+        let signers = fixture.signers();
+        let stake_distribution = fixture.stake_distribution();
+        let epoch_settings_message = FromEpochSettingsAdapter::try_adapt(EpochSettingsMessage {
+            epoch: registration_epoch,
+            next_signers: SignerMessagePart::from_signers(signers),
+            ..EpochSettingsMessage::dummy()
+        })
+        .unwrap();
+
+        let signer_registration_slave = MithrilSignerRegistrationSlaveBuilder::default()
+            .with_epoch_service({
+                let mut epoch_service = FakeEpochService::without_data();
+                epoch_service.toggle_errors(false, false, false, true);
+
+                epoch_service
+            })
+            .with_signer_recorder({
+                let mut signer_recorder = MockSignerRecorder::new();
+                signer_recorder
+                    .expect_record_signer_registration()
+                    .returning(|_| Ok(()))
+                    .times(5);
+
+                Arc::new(signer_recorder)
+            })
+            .with_signer_registration_verifier({
+                let mut signer_registration_verifier = MockSignerRegistrationVerifier::new();
+                signer_registration_verifier
+                    .expect_verify()
+                    .returning(|signer, _| Ok(SignerWithStake::from_signer(signer.to_owned(), 123)))
+                    .times(5);
+
+                Arc::new(signer_registration_verifier)
+            })
+            .with_master_aggregator_client({
+                let mut aggregator_client = MockAggregatorClient::new();
+                aggregator_client
+                    .expect_retrieve_epoch_settings()
+                    .returning(move || Ok(Some(epoch_settings_message.clone())))
+                    .times(1);
+
+                Arc::new(aggregator_client)
+            })
+            .with_stake_store({
+                let mut stake_store = MockStakeStore::new();
+                stake_store
+                    .expect_get_stakes()
+                    .returning(move |_epoch| Ok(Some(stake_distribution.clone())))
+                    .times(1);
+
+                Arc::new(stake_store)
+            })
+            .build();
+
+        signer_registration_slave
+            .synchronize_all_signers()
+            .await
+            .expect_err("synchronize_all_signers should fail");
+    }
+
     #[tokio::test]
     async fn synchronize_all_signers_fails_if_fetching_epoch_settings_fails() {
         let signer_registration_slave = MithrilSignerRegistrationSlaveBuilder::default()