feat: compute and store the total stake in the Epoch service

Author: dlachaume

mithril-aggregator/src/services/epoch_service.rs

@@ -113,6 +113,9 @@ pub trait EpochService: Sync + Send {
 
     /// Get the total number of SPOs for the current epoch.
     fn total_spo(&self) -> StdResult<&u32>;
+
+    /// Get the total stake for the current epoch.
+    fn total_stake(&self) -> StdResult<&u64>;
 }
 
 struct EpochData {
@@ -128,6 +131,7 @@ struct EpochData {
     next_signers: Vec<Signer>,
     signed_entity_config: SignedEntityConfig,
     total_spo: u32,
+    total_stake: u64,
 }
 
 struct ComputedEpochData {
@@ -215,14 +219,16 @@ impl MithrilEpochService {
         Ok(cardano_era)
     }
 
-    async fn get_total_spo(&self, epoch: Epoch) -> StdResult<u32> {
+    async fn get_total_spo_and_total_stake(&self, epoch: Epoch) -> StdResult<(u32, u64)> {
         let stake_distribution = self
             .stake_distribution_service
             .get_stake_distribution(epoch)
             .await
             .with_context(|| "Epoch service failed to obtain stake distribution")?;
 
-        Ok(stake_distribution.len() as u32)
+        let total_stake = stake_distribution.values().copied().sum::<u64>();
+
+        Ok((stake_distribution.len() as u32, total_stake))
     }
 
     async fn get_signers_with_stake_at_epoch(
@@ -333,7 +339,7 @@ impl EpochService for MithrilEpochService {
                 .clone(),
         };
 
-        let total_spo = self.get_total_spo(epoch).await?;
+        let (total_spo, total_stake) = self.get_total_spo_and_total_stake(epoch).await?;
 
         self.epoch_data = Some(EpochData {
             cardano_era,
@@ -348,6 +354,7 @@ impl EpochService for MithrilEpochService {
             next_signers,
             signed_entity_config,
             total_spo,
+            total_stake,
         });
         self.computed_epoch_data = None;
 
@@ -483,6 +490,10 @@ impl EpochService for MithrilEpochService {
     fn total_spo(&self) -> StdResult<&u32> {
         Ok(&self.unwrap_data()?.total_spo)
     }
+
+    fn total_stake(&self) -> StdResult<&u64> {
+        Ok(&self.unwrap_data()?.total_stake)
+    }
 }
 
 #[cfg(test)]
@@ -506,6 +517,7 @@ pub struct FakeEpochServiceBuilder {
     pub next_signers_with_stake: Vec<SignerWithStake>,
     pub signed_entity_config: SignedEntityConfig,
     pub total_spo: u32,
+    pub total_stake: u64,
 }
 
 #[cfg(test)]
@@ -525,6 +537,7 @@ impl FakeEpochServiceBuilder {
             next_signers_with_stake: signers,
             signed_entity_config: SignedEntityConfig::dummy(),
             total_spo: 0,
+            total_stake: 0,
         }
     }
 
@@ -561,6 +574,7 @@ impl FakeEpochServiceBuilder {
                 next_signers,
                 signed_entity_config: self.signed_entity_config,
                 total_spo: self.total_spo,
+                total_stake: self.total_stake,
             }),
             computed_epoch_data: Some(ComputedEpochData {
                 aggregate_verification_key: protocol_multi_signer
@@ -758,6 +772,10 @@ impl EpochService for FakeEpochService {
     fn total_spo(&self) -> StdResult<&u32> {
         Ok(&self.unwrap_data()?.total_spo)
     }
+
+    fn total_stake(&self) -> StdResult<&u64> {
+        Ok(&self.unwrap_data()?.total_stake)
+    }
 }
 
 #[cfg(test)]
@@ -767,7 +785,9 @@ mod tests {
         BlockNumber, CardanoTransactionsSigningConfig, PartyId, StakeDistribution,
     };
     use mithril_common::era::SupportedEra;
-    use mithril_common::test_utils::{fake_data, MithrilFixture, MithrilFixtureBuilder};
+    use mithril_common::test_utils::{
+        fake_data, MithrilFixture, MithrilFixtureBuilder, StakeDistributionGenerationMethod,
+    };
     use mithril_persistence::store::adapter::MemoryAdapter;
     use std::collections::{BTreeSet, HashMap};
 
@@ -778,13 +798,19 @@ mod tests {
 
     use super::*;
 
-    fn stake_distribution_with_given_number_of_spo(total_spo: u32) -> StakeDistribution {
-        let mut stake_distribution: StakeDistribution = StakeDistribution::new();
-        for i in 0..total_spo {
-            let party_id = format!("party-id-{i}");
-            stake_distribution.insert(party_id.clone(), 0);
+    fn build_stake_distribution(total_spo: u32, total_stake: u64) -> StakeDistribution {
+        if total_stake % total_spo as u64 != 0 {
+            panic!("'total_stake' must be a multiple of 'total_spo' to facilitate the creation of a uniform stake distribution");
         }
-        stake_distribution
+
+        let fixture = MithrilFixtureBuilder::default()
+            .with_signers(total_spo as usize)
+            .with_stake_distribution(StakeDistributionGenerationMethod::Uniform(
+                total_stake / total_spo as u64,
+            ))
+            .build();
+
+        fixture.stake_distribution()
     }
 
     #[derive(Debug, Clone, PartialEq)]
@@ -803,6 +829,7 @@ mod tests {
         next_signers: BTreeSet<Signer>,
         signed_entity_config: SignedEntityConfig,
         total_spo: u32,
+        total_stake: u64,
     }
 
     #[derive(Debug, Clone, PartialEq)]
@@ -842,6 +869,7 @@ mod tests {
                 next_signers: service.next_signers()?.clone().into_iter().collect(),
                 signed_entity_config: service.signed_entity_config()?.clone(),
                 total_spo: *service.total_spo()?,
+                total_stake: *service.total_stake()?,
             })
         }
     }
@@ -879,6 +907,7 @@ mod tests {
         stored_next_epoch_settings: AggregatorEpochSettings,
         stored_signer_registration_epoch_settings: AggregatorEpochSettings,
         total_spo: u32,
+        total_stake: u64,
     }
 
     impl EpochServiceBuilder {
@@ -905,7 +934,8 @@ mod tests {
                     protocol_parameters: epoch_fixture.protocol_parameters(),
                     cardano_transactions_signing_config: CardanoTransactionsSigningConfig::dummy(),
                 },
-                total_spo: 0,
+                total_spo: 1,
+                total_stake: 0,
             }
         }
 
@@ -946,7 +976,7 @@ mod tests {
 
             let era_checker = EraChecker::new(self.mithril_era, Epoch::default());
 
-            let stake_distribution = stake_distribution_with_given_number_of_spo(self.total_spo);
+            let stake_distribution = build_stake_distribution(self.total_spo, self.total_stake);
             let mut stake_distribution_service = MockStakeDistributionService::new();
             stake_distribution_service
                 .expect_get_stake_distribution()
@@ -995,7 +1025,8 @@ mod tests {
             allowed_discriminants: SignedEntityConfig::dummy().allowed_discriminants,
             cardano_era: "CardanoEra".to_string(),
             mithril_era: SupportedEra::eras()[1],
-            total_spo: 37,
+            total_spo: 40,
+            total_stake: 20_000_000,
             ..EpochServiceBuilder::new(epoch, current_epoch_fixture.clone())
         };
 
@@ -1032,7 +1063,8 @@ mod tests {
                 current_signers: current_epoch_fixture.signers().into_iter().collect(),
                 next_signers: next_epoch_fixture.signers().into_iter().collect(),
                 signed_entity_config: SignedEntityConfig::dummy(),
-                total_spo: 37,
+                total_spo: 40,
+                total_stake: 20_000_000,
             }
         );
     }
@@ -1250,6 +1282,7 @@ mod tests {
             ),
             ("signed_entity_config", service.signed_entity_config().err()),
             ("total_spo", service.total_spo().err()),
+            ("total_stake", service.total_stake().err()),
         ] {
             let error =
                 res.unwrap_or_else(|| panic!("getting {name} should have returned an error"));
@@ -1286,6 +1319,7 @@ mod tests {
         assert!(service.next_signers().is_ok());
         assert!(service.signed_entity_config().is_ok());
         assert!(service.total_spo().is_ok());
+        assert!(service.total_stake().is_ok());
 
         for (name, res) in [
             (