Pruning of old vkeys is now done manually instead of when saving

This in order to prune only when needed, reducing the risk of doing
it on data that are not yet to be pruned. The manual call will be
likely done when a signer registration round open.

Author: Alenar

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

@@ -394,17 +394,12 @@ impl<'conn> DeleteSignerRegistrationRecordProvider<'conn> {
 /// Service to deal with signer_registration (read & write).
 pub struct SignerRegistrationStore {
     connection: Arc<Mutex<Connection>>,
-    /// Number of epoch before previous records will be deleted at the next save
-    epoch_retention_limit: Option<u64>,
 }
 
 impl SignerRegistrationStore {
     /// Create a new [SignerRegistrationStore] service
-    pub fn new(connection: Arc<Mutex<Connection>>, epoch_retention_limit: Option<u64>) -> Self {
-        Self {
-            connection,
-            epoch_retention_limit,
-        }
+    pub fn new(connection: Arc<Mutex<Connection>>) -> Self {
+        Self { connection }
     }
 }
 
@@ -428,17 +423,6 @@ impl VerificationKeyStorer for SignerRegistrationStore {
             ))
             .map_err(|e| AdapterError::GeneralError(format!("{e}")))?;
 
-        if let Some(threshold) = self.epoch_retention_limit {
-            // Note: this means that if called with an epoch in the future this could remove all
-            // current records, the caller should check that the given epoch is the current one
-            // or we must get it to do the work ourself here.
-            let _deleted_records = DeleteSignerRegistrationRecordProvider::new(connection)
-                // we want to prune including the given epoch (+1)
-                .prune(epoch - threshold + 1)
-                .map_err(|e| AdapterError::QueryError(e))?
-                .collect::<Vec<_>>();
-        }
-
         match existing_record {
             None => Ok(None),
             Some(previous_record) => Ok(Some(previous_record.into())),
@@ -463,6 +447,17 @@ impl VerificationKeyStorer for SignerRegistrationStore {
             false => Ok(Some(signer_with_stakes)),
         }
     }
+
+    async fn prune_verification_keys(&self, max_epoch_to_prune: Epoch) -> Result<(), StoreError> {
+        let connection = &*self.connection.lock().await;
+        let _deleted_records = DeleteSignerRegistrationRecordProvider::new(connection)
+            // we want to prune including the given epoch (+1)
+            .prune(max_epoch_to_prune + 1)
+            .map_err(|e| AdapterError::QueryError(e))?
+            .collect::<Vec<_>>();
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -805,7 +800,6 @@ mod tests {
 
     pub fn init_signer_registration_store(
         initial_data: Vec<(Epoch, HashMap<PartyId, SignerWithStake>)>,
-        retention_limit: Option<usize>,
     ) -> Arc<dyn VerificationKeyStorer> {
         let connection = Connection::open(":memory:").unwrap();
         let initial_data: Vec<(Epoch, Vec<SignerWithStake>)> = initial_data
@@ -815,10 +809,9 @@ mod tests {
 
         setup_signer_registration_db(&connection, initial_data).unwrap();
 
-        Arc::new(SignerRegistrationStore::new(
-            Arc::new(Mutex::new(connection)),
-            retention_limit.map(|threshold| threshold as u64),
-        ))
+        Arc::new(SignerRegistrationStore::new(Arc::new(Mutex::new(
+            connection,
+        ))))
     }
 
     test_verification_key_storer!(

mithril-aggregator/src/dependency_injection/builder.rs

@@ -404,7 +404,6 @@ impl DependenciesBuilder {
     async fn build_verification_key_store(&mut self) -> Result<Arc<dyn VerificationKeyStorer>> {
         Ok(Arc::new(SignerRegistrationStore::new(
             self.get_sqlite_connection().await?,
-            self.configuration.store_retention_limit.map(|l| l as u64),
         )))
     }
 

mithril-aggregator/src/multi_signer.rs

@@ -595,10 +595,9 @@ mod tests {
 
     async fn setup_multi_signer() -> MultiSignerImpl {
         let beacon = fake_data::beacon();
-        let verification_key_store = VerificationKeyStore::new(
-            Box::new(MemoryAdapter::<Epoch, HashMap<PartyId, SignerWithStake>>::new(None).unwrap()),
-            None,
-        );
+        let verification_key_store = VerificationKeyStore::new(Box::new(
+            MemoryAdapter::<Epoch, HashMap<PartyId, SignerWithStake>>::new(None).unwrap(),
+        ));
         let stake_store = StakeStore::new(
             Box::new(MemoryAdapter::<Epoch, StakeDistribution>::new(None).unwrap()),
             None,

mithril-aggregator/src/signer_registerer.rs

@@ -295,10 +295,9 @@ mod tests {
     #[tokio::test]
     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, SignerWithStake>>::new(None).unwrap()),
-            None,
-        ));
+        let verification_key_store = Arc::new(VerificationKeyStore::new(Box::new(
+            MemoryAdapter::<Epoch, HashMap<PartyId, SignerWithStake>>::new(None).unwrap(),
+        )));
         let mut signer_recorder = MockSignerRecorder::new();
         signer_recorder
             .expect_record_signer_id()
@@ -341,10 +340,9 @@ mod tests {
     #[tokio::test]
     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, SignerWithStake>>::new(None).unwrap()),
-            None,
-        ));
+        let verification_key_store = Arc::new(VerificationKeyStore::new(Box::new(
+            MemoryAdapter::<Epoch, HashMap<PartyId, SignerWithStake>>::new(None).unwrap(),
+        )));
         let mut signer_recorder = MockSignerRecorder::new();
         signer_recorder
             .expect_record_signer_id()
@@ -390,10 +388,9 @@ mod tests {
     #[tokio::test]
     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, SignerWithStake>>::new(None).unwrap()),
-            None,
-        ));
+        let verification_key_store = Arc::new(VerificationKeyStore::new(Box::new(
+            MemoryAdapter::<Epoch, HashMap<PartyId, SignerWithStake>>::new(None).unwrap(),
+        )));
         let signer_recorder = MockSignerRecorder::new();
         let signer_registerer = MithrilSignerRegisterer::new(
             Arc::new(chain_observer),

mithril-aggregator/src/store/verification_key_store.rs

@@ -1,5 +1,4 @@
 use async_trait::async_trait;
-use mithril_common::store::StorePruner;
 use std::collections::HashMap;
 use tokio::sync::RwLock;
 
@@ -24,40 +23,25 @@ pub trait VerificationKeyStorer: Sync + Send {
         &self,
         epoch: Epoch,
     ) -> Result<Option<HashMap<PartyId, Signer>>, StoreError>;
+
+    /// 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>;
 }
 
 /// Store for the `VerificationKey`.
 pub struct VerificationKeyStore {
     adapter: RwLock<Adapter>,
-    retention_limit: Option<usize>,
 }
 
 impl VerificationKeyStore {
     /// Create a new instance.
-    pub fn new(adapter: Adapter, retention_limit: Option<usize>) -> Self {
+    pub fn new(adapter: Adapter) -> Self {
         Self {
             adapter: RwLock::new(adapter),
-            retention_limit,
         }
     }
 }
 
-#[async_trait]
-impl StorePruner for VerificationKeyStore {
-    type Key = Epoch;
-    type Record = HashMap<PartyId, SignerWithStake>;
-
-    fn get_adapter(
-        &self,
-    ) -> &RwLock<Box<dyn StoreAdapter<Key = Self::Key, Record = Self::Record>>> {
-        &self.adapter
-    }
-
-    fn get_max_records(&self) -> Option<usize> {
-        self.retention_limit
-    }
-}
-
 #[async_trait]
 impl VerificationKeyStorer for VerificationKeyStore {
     async fn save_verification_key(
@@ -75,7 +59,6 @@ impl VerificationKeyStorer for VerificationKeyStore {
             .await
             .store_record(&epoch, &signers)
             .await?;
-        self.prune().await?;
 
         Ok(prev_signer)
     }
@@ -87,6 +70,21 @@ impl VerificationKeyStorer for VerificationKeyStore {
         let record = self.adapter.read().await.get_record(&epoch).await?;
         Ok(record.map(|h| h.into_iter().map(|(k, v)| (k, v.into())).collect()))
     }
+
+    async fn prune_verification_keys(&self, max_epoch_to_prune: Epoch) -> Result<(), StoreError> {
+        let mut adapter = self.adapter.write().await;
+
+        for (epoch, _record) in adapter
+            .get_last_n_records(usize::MAX)
+            .await?
+            .into_iter()
+            .filter(|(e, _)| e <= &max_epoch_to_prune)
+        {
+            adapter.remove(&epoch).await?;
+        }
+
+        Ok(())
+    }
 }
 
 /// Macro that generate tests that a [VerificationKeyStorer] must pass
@@ -118,8 +116,8 @@ macro_rules! test_verification_key_storer {
             }
 
             #[tokio::test]
-            async fn check_retention_limit() {
-                test_suite::check_retention_limit(&$store_builder).await;
+            async fn can_prune_keys_from_given_epoch_retention_limit() {
+                test_suite::can_prune_keys_from_given_epoch_retention_limit(&$store_builder).await;
             }
         }
     };
@@ -139,11 +137,8 @@ pub mod test_suite {
 
     /// A builder of [VerificationKeyStorer], the arguments are:
     /// * initial_data
-    /// * optional retention limit
-    type StoreBuilder = dyn Fn(
-        Vec<(Epoch, HashMap<PartyId, SignerWithStake>)>,
-        Option<usize>,
-    ) -> Arc<dyn VerificationKeyStorer>;
+    type StoreBuilder =
+        dyn Fn(Vec<(Epoch, HashMap<PartyId, SignerWithStake>)>) -> Arc<dyn VerificationKeyStorer>;
 
     fn build_signers(
         nb_epoch: u64,
@@ -177,7 +172,7 @@ pub mod test_suite {
 
     pub async fn save_key_in_empty_store(store_builder: &StoreBuilder) {
         let signers = build_signers(0, 0);
-        let store = store_builder(signers, None);
+        let store = store_builder(signers);
         let res = store
             .save_verification_key(
                 Epoch(0),
@@ -198,7 +193,7 @@ pub mod test_suite {
 
     pub async fn update_signer_in_store(store_builder: &StoreBuilder) {
         let signers = build_signers(1, 1);
-        let store = store_builder(signers, None);
+        let store = store_builder(signers);
         let res = store
             .save_verification_key(
                 Epoch(1),
@@ -229,23 +224,23 @@ pub mod test_suite {
 
     pub async fn get_verification_keys_for_empty_epoch(store_builder: &StoreBuilder) {
         let signers = build_signers(2, 1);
-        let store = store_builder(signers, None);
+        let store = store_builder(signers);
         let res = store.get_verification_keys(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());
+
         let expected_signers: Option<BTreeMap<PartyId, Signer>> = signers
-            .iter()
+            .into_iter()
             .filter(|(e, _)| e == 1)
-            .cloned()
             .map(|(_, signers)| {
                 BTreeMap::from_iter(signers.into_iter().map(|(p, s)| (p, s.into())))
             })
             .next();
-        let store = store_builder(signers, None);
         let res = store
             .get_verification_keys(Epoch(1))
             .await
@@ -255,30 +250,27 @@ pub mod test_suite {
         assert_eq!(expected_signers, res);
     }
 
-    pub async fn check_retention_limit(store_builder: &StoreBuilder) {
-        let signers = build_signers(2, 2);
-        let store = store_builder(signers, Some(2));
-        assert!(store
-            .get_verification_keys(Epoch(1))
-            .await
-            .unwrap()
-            .is_some());
-        let _ = store
-            .save_verification_key(
-                Epoch(3),
-                SignerWithStake {
-                    party_id: "party_id".to_string(),
-                    verification_key: "whatever".to_string(),
-                    verification_key_signature: None,
-                    operational_certificate: None,
-                    kes_period: None,
-                    stake: 10,
-                },
-            )
-            .await
-            .unwrap();
-        let first_epoch_keys = store.get_verification_keys(Epoch(1)).await.unwrap();
-        assert_eq!(None, first_epoch_keys);
+    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);
+
+        for epoch in 1..6 {
+            assert!(
+                store
+                    .get_verification_keys(Epoch(epoch))
+                    .await
+                    .unwrap()
+                    .is_some(),
+                "Keys should exist before pruning"
+            );
+            store
+                .prune_verification_keys(Epoch(epoch))
+                .await
+                .expect("Pruning should not fail");
+
+            let pruned_epoch_keys = store.get_verification_keys(Epoch(epoch)).await.unwrap();
+            assert_eq!(None, pruned_epoch_keys);
+        }
     }
 }
 
@@ -294,7 +286,6 @@ mod tests {
 
     pub fn init_store(
         initial_data: Vec<(Epoch, HashMap<PartyId, SignerWithStake>)>,
-        retention_limit: Option<usize>,
     ) -> Arc<dyn VerificationKeyStorer> {
         let values = if initial_data.is_empty() {
             None
@@ -305,10 +296,7 @@ mod tests {
         let adapter: MemoryAdapter<Epoch, HashMap<PartyId, SignerWithStake>> =
             MemoryAdapter::new(values).unwrap();
 
-        Arc::new(VerificationKeyStore::new(
-            Box::new(adapter),
-            retention_limit,
-        ))
+        Arc::new(VerificationKeyStore::new(Box::new(adapter)))
     }
 
     test_verification_key_storer!(