Add method to vkey store to get registered stake distribution for an epoch

Author: Alenar

mithril-aggregator/src/database/provider/signer_registration.rs

@@ -8,7 +8,7 @@ use mithril_common::{
     crypto_helper::KESPeriod,
     entities::{
         Epoch, HexEncodedOpCert, HexEncodedVerificationKey, HexEncodedVerificationKeySignature,
-        PartyId, Signer, SignerWithStake, Stake,
+        PartyId, Signer, SignerWithStake, Stake, StakeDistribution,
     },
     sqlite::{
         EntityCursor, HydrationError, Projection, Provider, SourceAlias, SqLiteEntity,
@@ -458,6 +458,26 @@ impl VerificationKeyStorer for SignerRegistrationStore {
 
         Ok(())
     }
+
+    async fn get_stake_distribution_for_epoch(
+        &self,
+        epoch: Epoch,
+    ) -> Result<Option<StakeDistribution>, StoreError> {
+        let connection = &*self.connection.lock().await;
+        let provider = SignerRegistrationRecordProvider::new(connection);
+        let cursor = provider
+            .get_by_epoch(&epoch)
+            .map_err(|e| AdapterError::GeneralError(format!("{e}")))?;
+
+        let stake_distribution = StakeDistribution::from_iter(
+            cursor.map(|r| (r.signer_id, r.stake.unwrap_or_default())),
+        );
+
+        match stake_distribution.is_empty() {
+            true => Ok(None),
+            false => Ok(Some(stake_distribution)),
+        }
+    }
 }
 
 #[cfg(test)]

mithril-aggregator/src/store/verification_key_store.rs

@@ -2,12 +2,15 @@ use async_trait::async_trait;
 use std::collections::HashMap;
 use tokio::sync::RwLock;
 
-use mithril_common::entities::{Epoch, PartyId, Signer, SignerWithStake};
+use mithril_common::entities::{Epoch, PartyId, Signer, SignerWithStake, StakeDistribution};
 use mithril_common::store::{adapter::StoreAdapter, StoreError};
 
 type Adapter = Box<dyn StoreAdapter<Key = Epoch, Record = HashMap<PartyId, SignerWithStake>>>;
 
 /// Store and get signers verification keys for given epoch.
+///
+/// Important note: This store works on the **recording** epoch, the epoch at which the signers
+/// are signed into a certificate so they can sign single signatures at the next epoch.
 #[async_trait]
 pub trait VerificationKeyStorer: Sync + Send {
     /// Save the verification key, for the given [Signer] for the given [Epoch], returns the
@@ -26,6 +29,12 @@ pub trait VerificationKeyStorer: Sync + Send {
 
     /// Prune all verification keys that are at or below the given epoch.
     async fn prune_verification_keys(&self, max_epoch_to_prune: Epoch) -> Result<(), StoreError>;
+
+    /// Return the parties that are stored at the given epoch.
+    async fn get_stake_distribution_for_epoch(
+        &self,
+        epoch: Epoch,
+    ) -> Result<Option<StakeDistribution>, StoreError>;
 }
 
 /// Store for the `VerificationKey`.
@@ -85,6 +94,19 @@ impl VerificationKeyStorer for VerificationKeyStore {
 
         Ok(())
     }
+
+    async fn get_stake_distribution_for_epoch(
+        &self,
+        epoch: Epoch,
+    ) -> Result<Option<StakeDistribution>, StoreError> {
+        let record = self.adapter.read().await.get_record(&epoch).await?;
+        Ok(record.map(|r| {
+            StakeDistribution::from_iter(
+                r.into_iter()
+                    .map(|(party_id, signer)| (party_id, signer.stake)),
+            )
+        }))
+    }
 }
 
 /// Macro that generate tests that a [VerificationKeyStorer] must pass
@@ -110,11 +132,21 @@ macro_rules! test_verification_key_storer {
                 test_suite::get_verification_keys_for_empty_epoch(&$store_builder).await;
             }
 
+            #[tokio::test]
+            async fn get_stake_distribution_for_empty_epoch() {
+                test_suite::get_stake_distribution_for_empty_epoch(&$store_builder).await;
+            }
+
             #[tokio::test]
             async fn get_verification_keys_for_existing_epoch() {
                 test_suite::get_verification_keys_for_existing_epoch(&$store_builder).await;
             }
 
+            #[tokio::test]
+            async fn get_stake_distribution_for_existing_epoch() {
+                test_suite::get_stake_distribution_for_existing_epoch(&$store_builder).await;
+            }
+
             #[tokio::test]
             async fn can_prune_keys_from_given_epoch_retention_limit() {
                 test_suite::can_prune_keys_from_given_epoch_retention_limit(&$store_builder).await;
@@ -129,7 +161,7 @@ pub(crate) use test_verification_key_storer;
 #[macro_use]
 #[cfg(test)]
 pub mod test_suite {
-    use mithril_common::entities::{Epoch, PartyId, Signer, SignerWithStake};
+    use mithril_common::entities::{Epoch, PartyId, Signer, SignerWithStake, StakeDistribution};
     use mithril_common::test_utils::fake_keys;
     use std::collections::{BTreeMap, HashMap};
     use std::sync::Arc;
@@ -233,6 +265,17 @@ pub mod test_suite {
         assert!(res.is_none());
     }
 
+    pub async fn get_stake_distribution_for_empty_epoch(store_builder: &StoreBuilder) {
+        let signers = build_signers(2, 1);
+        let store = store_builder(signers);
+        let res = store
+            .get_stake_distribution_for_epoch(Epoch(0))
+            .await
+            .unwrap();
+
+        assert!(res.is_none());
+    }
+
     pub async fn get_verification_keys_for_existing_epoch(store_builder: &StoreBuilder) {
         let signers = build_signers(2, 2);
         let store = store_builder(signers.clone());
@@ -253,6 +296,26 @@ pub mod test_suite {
         assert_eq!(expected_signers, res);
     }
 
+    pub async fn get_stake_distribution_for_existing_epoch(store_builder: &StoreBuilder) {
+        let signers = build_signers(2, 2);
+        let store = store_builder(signers.clone());
+
+        let expected: Option<StakeDistribution> = signers
+            .into_iter()
+            .filter(|(e, _)| e == 1)
+            .map(|(_, signers)| {
+                StakeDistribution::from_iter(signers.into_iter().map(|(p, s)| (p, s.stake)))
+            })
+            .next();
+        let res = store
+            .get_stake_distribution_for_epoch(Epoch(1))
+            .await
+            .unwrap()
+            .map(|x| BTreeMap::from_iter(x.into_iter()));
+
+        assert_eq!(expected, res);
+    }
+
     pub async fn can_prune_keys_from_given_epoch_retention_limit(store_builder: &StoreBuilder) {
         let signers = build_signers(6, 2);
         let store = store_builder(signers);