Skip to content

Commit

Permalink
feat: compute and store the total stake in the Epoch service
Browse files Browse the repository at this point in the history
  • Loading branch information
@dlachaume
dlachaume committed Nov 8, 2024
1 parent a723278 commit 5d48588
Showing 1 changed file with 59 additions and 19 deletions.
78 changes: 59 additions & 19 deletions mithril-aggregator/src/services/epoch_service.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -128,6 +131,7 @@ struct EpochData {
next_signers: Vec<Signer>,
signed_entity_config: SignedEntityConfig,
total_spo: u32,
total_stake: u64,
}

struct ComputedEpochData {
Expand Down Expand Up @@ -215,7 +219,7 @@ 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)
Expand All @@ -224,7 +228,9 @@ impl MithrilEpochService {
format!("Epoch service failed to obtain the stake distribution for epoch: {epoch}")
})?;

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(
Expand Down Expand Up @@ -335,7 +341,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,
Expand All @@ -350,6 +356,7 @@ impl EpochService for MithrilEpochService {
next_signers,
signed_entity_config,
total_spo,
total_stake,
});
self.computed_epoch_data = None;

Expand Down Expand Up @@ -485,6 +492,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)]
Expand All @@ -508,6 +519,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)]
Expand All @@ -527,6 +539,7 @@ impl FakeEpochServiceBuilder {
next_signers_with_stake: signers,
signed_entity_config: SignedEntityConfig::dummy(),
total_spo: 0,
total_stake: 0,
}
}

Expand Down Expand Up @@ -563,6 +576,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
Expand Down Expand Up @@ -760,6 +774,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)]
Expand All @@ -769,7 +787,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};

Expand All @@ -780,13 +800,13 @@ 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);
}
stake_distribution
fn build_uniform_stake_distribution(total_spo: u32, stake_by_spo: u64) -> StakeDistribution {
let fixture = MithrilFixtureBuilder::default()
.with_signers(total_spo as usize)
.with_stake_distribution(StakeDistributionGenerationMethod::Uniform(stake_by_spo))
.build();

fixture.stake_distribution()
}

#[derive(Debug, Clone, PartialEq)]
Expand All @@ -805,6 +825,7 @@ mod tests {
next_signers: BTreeSet<Signer>,
signed_entity_config: SignedEntityConfig,
total_spo: u32,
total_stake: u64,
}

#[derive(Debug, Clone, PartialEq)]
Expand Down Expand Up @@ -844,6 +865,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()?,
})
}
}
Expand Down Expand Up @@ -881,6 +903,7 @@ mod tests {
stored_next_epoch_settings: AggregatorEpochSettings,
stored_signer_registration_epoch_settings: AggregatorEpochSettings,
total_spo: u32,
total_stake: u64,
}

impl EpochServiceBuilder {
Expand All @@ -907,7 +930,8 @@ mod tests {
protocol_parameters: epoch_fixture.protocol_parameters(),
cardano_transactions_signing_config: CardanoTransactionsSigningConfig::dummy(),
},
total_spo: 0,
total_spo: 1,
total_stake: 0,
}
}

Expand Down Expand Up @@ -948,11 +972,23 @@ 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 mut stake_distribution_service = MockStakeDistributionService::new();
stake_distribution_service
.expect_get_stake_distribution()
.returning(move |_| Ok(stake_distribution.clone()));
let stake_distribution_service = {
assert!(
self.total_stake % self.total_spo as u64 == 0,
"'total_stake' must be a multiple of 'total_spo' to create a uniform stake distribution"
);
let stake_per_spo = self.total_stake / self.total_spo as u64;

let stake_distribution =
build_uniform_stake_distribution(self.total_spo, stake_per_spo);

let mut stake_distribution_service = MockStakeDistributionService::new();
stake_distribution_service
.expect_get_stake_distribution()
.returning(move |_| Ok(stake_distribution.clone()));

stake_distribution_service
};

MithrilEpochService::new(
AggregatorEpochSettings {
Expand Down Expand Up @@ -997,7 +1033,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())
};

Expand Down Expand Up @@ -1034,7 +1071,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,
}
);
}
Expand Down Expand Up @@ -1252,6 +1290,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"));
Expand Down Expand Up @@ -1288,6 +1327,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 [
(
Expand Down

0 comments on commit 5d48588

Please sign in to comment.