Update 'VerificationKeyStore' record type

And use 'SignerWithStake' instead of 'Signer' for new adapter to save stake as well.

Author: jpraynaud

mithril-aggregator/src/dependency_injection/builder.rs

@@ -12,8 +12,8 @@ use mithril_common::{
         ImmutableFileObserver, ImmutableFileSystemObserver,
     },
     entities::{
-        Beacon, Certificate, CertificatePending, Epoch, PartyId, ProtocolParameters, Signer,
-        SingleSignatures,
+        Beacon, Certificate, CertificatePending, Epoch, PartyId, ProtocolParameters,
+        SignerWithStake, SingleSignatures,
     },
     era::{
         adapters::{EraReaderAdapterBuilder, EraReaderDummyAdapter},
@@ -419,30 +419,31 @@ impl DependenciesBuilder {
     }
 
     async fn build_verification_key_store(&mut self) -> Result<Arc<VerificationKeyStore>> {
-        let adapter: Box<dyn StoreAdapter<Key = Epoch, Record = HashMap<PartyId, Signer>>> =
-            match self.configuration.environment {
-                ExecutionEnvironment::Production => {
-                    let adapter =
-                        SQLiteAdapter::new("verification_key", self.get_sqlite_connection().await?)
-                            .map_err(|e| DependenciesBuilderError::Initialization {
-                                message: "Cannot create SQLite adapter for VerificationKeyStore."
-                                    .to_string(),
-                                error: Some(e.into()),
-                            })?;
-
-                    Box::new(adapter)
-                }
-                _ => {
-                    let adapter = MemoryAdapter::new(None).map_err(|e| {
-                        DependenciesBuilderError::Initialization {
-                            message: "Cannot create Memory adapter for VerificationKeyStore."
+        let adapter: Box<
+            dyn StoreAdapter<Key = Epoch, Record = HashMap<PartyId, SignerWithStake>>,
+        > = match self.configuration.environment {
+            ExecutionEnvironment::Production => {
+                let adapter =
+                    SQLiteAdapter::new("verification_key", self.get_sqlite_connection().await?)
+                        .map_err(|e| DependenciesBuilderError::Initialization {
+                            message: "Cannot create SQLite adapter for VerificationKeyStore."
                                 .to_string(),
                             error: Some(e.into()),
-                        }
-                    })?;
-                    Box::new(adapter)
-                }
-            };
+                        })?;
+
+                Box::new(adapter)
+            }
+            _ => {
+                let adapter = MemoryAdapter::new(None).map_err(|e| {
+                    DependenciesBuilderError::Initialization {
+                        message: "Cannot create Memory adapter for VerificationKeyStore."
+                            .to_string(),
+                        error: Some(e.into()),
+                    }
+                })?;
+                Box::new(adapter)
+            }
+        };
 
         Ok(Arc::new(VerificationKeyStore::new(
             adapter,

mithril-aggregator/src/multi_signer.rs

@@ -652,8 +652,10 @@ mod tests {
         let beacon = fake_data::beacon();
         let verification_key_store = VerificationKeyStore::new(
             Box::new(
-                MemoryAdapter::<Epoch, HashMap<entities::PartyId, entities::Signer>>::new(None)
-                    .unwrap(),
+                MemoryAdapter::<Epoch, HashMap<entities::PartyId, entities::SignerWithStake>>::new(
+                    None,
+                )
+                .unwrap(),
             ),
             None,
         );

mithril-aggregator/src/signer_registerer.rs

@@ -223,7 +223,7 @@ mod tests {
 
     use mithril_common::{
         chain_observer::FakeObserver,
-        entities::{Epoch, PartyId, Signer},
+        entities::{Epoch, PartyId, Signer, SignerWithStake},
         store::adapter::MemoryAdapter,
         test_utils::MithrilFixtureBuilder,
     };
@@ -237,7 +237,7 @@ mod tests {
     async fn can_register_signer_if_registration_round_is_opened_with_operational_certificate() {
         let chain_observer = FakeObserver::default();
         let verification_key_store = Arc::new(VerificationKeyStore::new(
-            Box::new(MemoryAdapter::<Epoch, HashMap<PartyId, Signer>>::new(None).unwrap()),
+            Box::new(MemoryAdapter::<Epoch, HashMap<PartyId, SignerWithStake>>::new(None).unwrap()),
             None,
         ));
         let signer_registerer =
@@ -275,7 +275,7 @@ mod tests {
     async fn can_register_signer_if_registration_round_is_opened_without_operational_certificate() {
         let chain_observer = FakeObserver::default();
         let verification_key_store = Arc::new(VerificationKeyStore::new(
-            Box::new(MemoryAdapter::<Epoch, HashMap<PartyId, Signer>>::new(None).unwrap()),
+            Box::new(MemoryAdapter::<Epoch, HashMap<PartyId, SignerWithStake>>::new(None).unwrap()),
             None,
         ));
         let signer_registerer =
@@ -316,7 +316,7 @@ mod tests {
     async fn cant_register_signer_if_registration_round_is_not_opened() {
         let chain_observer = FakeObserver::default();
         let verification_key_store = Arc::new(VerificationKeyStore::new(
-            Box::new(MemoryAdapter::<Epoch, HashMap<PartyId, Signer>>::new(None).unwrap()),
+            Box::new(MemoryAdapter::<Epoch, HashMap<PartyId, SignerWithStake>>::new(None).unwrap()),
             None,
         ));
         let signer_registerer =

mithril-aggregator/src/store/verification_key_store.rs

@@ -6,7 +6,7 @@ use tokio::sync::RwLock;
 use mithril_common::entities::{Epoch, PartyId, Signer, SignerWithStake};
 use mithril_common::store::{adapter::StoreAdapter, StoreError};
 
-type Adapter = Box<dyn StoreAdapter<Key = Epoch, Record = HashMap<PartyId, Signer>>>;
+type Adapter = Box<dyn StoreAdapter<Key = Epoch, Record = HashMap<PartyId, SignerWithStake>>>;
 
 /// Mocking trait for `VerificationKeyStore`.
 #[async_trait]
@@ -43,7 +43,7 @@ impl VerificationKeyStore {
 #[async_trait]
 impl StorePruner for VerificationKeyStore {
     type Key = Epoch;
-    type Record = HashMap<PartyId, Signer>;
+    type Record = HashMap<PartyId, SignerWithStake>;
 
     fn get_adapter(
         &self,
@@ -67,23 +67,23 @@ impl VerificationKeyStorer for VerificationKeyStore {
             Some(s) => s,
             None => HashMap::new(),
         };
-        let prev_signer = signers.insert(signer.party_id.to_owned(), signer.clone().into());
+        let prev_signer = signers.insert(signer.party_id.to_owned(), signer.clone());
         self.adapter
             .write()
             .await
             .store_record(&epoch, &signers)
             .await?;
         self.prune().await?;
 
-        Ok(prev_signer.map(|prev_signer| SignerWithStake::from_signer(prev_signer, signer.stake)))
+        Ok(prev_signer)
     }
 
     async fn get_verification_keys(
         &self,
         epoch: Epoch,
     ) -> Result<Option<HashMap<PartyId, Signer>>, StoreError> {
         let record = self.adapter.read().await.get_record(&epoch).await?;
-        Ok(record)
+        Ok(record.map(|h| h.into_iter().map(|(k, v)| (k, v.into())).collect()))
     }
 }
 
@@ -98,21 +98,22 @@ mod tests {
         signers_per_epoch: u64,
         retention_limit: Option<usize>,
     ) -> VerificationKeyStore {
-        let mut values: Vec<(Epoch, HashMap<PartyId, Signer>)> = Vec::new();
+        let mut values: Vec<(Epoch, HashMap<PartyId, SignerWithStake>)> = Vec::new();
 
         for epoch in 1..=nb_epoch {
-            let mut signers: HashMap<PartyId, Signer> = HashMap::new();
+            let mut signers: HashMap<PartyId, SignerWithStake> = HashMap::new();
 
             for party_idx in 1..=signers_per_epoch {
                 let party_id = format!("{party_idx}");
                 signers.insert(
                     party_id.clone(),
-                    Signer {
+                    SignerWithStake {
                         party_id: party_id.clone(),
                         verification_key: format!("vkey {party_id}"),
                         verification_key_signature: None,
                         operational_certificate: None,
                         kes_period: None,
+                        stake: 10,
                     },
                 );
             }
@@ -124,7 +125,7 @@ mod tests {
         } else {
             None
         };
-        let adapter: MemoryAdapter<Epoch, HashMap<PartyId, Signer>> =
+        let adapter: MemoryAdapter<Epoch, HashMap<PartyId, SignerWithStake>> =
             MemoryAdapter::new(values).unwrap();
         VerificationKeyStore::new(Box::new(adapter), retention_limit)
     }