Create a task to prune StakePool and add it to the UpkeepService.

Author: sfauvel

mithril-aggregator/src/database/repository/stake_pool_store.rs

@@ -15,6 +15,7 @@ use crate::database::query::{
     DeleteStakePoolQuery, GetStakePoolQuery, InsertOrReplaceStakePoolQuery,
 };
 use crate::database::record::StakePool;
+use crate::services::EpochPruningTask;
 
 /// Service to deal with stake pools (read & write).
 pub struct StakePoolStore {
@@ -53,17 +54,6 @@ impl StakeStorer for StakePoolStore {
             .with_context(|| format!("persist stakes failure, epoch: {epoch}"))
             .map_err(AdapterError::GeneralError)?;
 
-        // Prune useless old stake distributions.
-        if let Some(threshold) = self.retention_limit {
-            let _ = self
-                .connection
-                .fetch(DeleteStakePoolQuery::below_epoch_threshold(
-                    epoch - threshold,
-                ))
-                .map_err(AdapterError::QueryError)?
-                .count();
-        }
-
         Ok(Some(StakeDistribution::from_iter(
             pools.into_iter().map(|p| (p.stake_pool_id, p.stake)),
         )))
@@ -95,27 +85,43 @@ impl StakeDistributionRetriever for StakePoolStore {
     }
 }
 
+#[async_trait]
+impl EpochPruningTask for StakePoolStore {
+    fn pruned_data(&self) -> &'static str {
+        "Stake pool"
+    }
+
+    async fn prune(&self, epoch: Epoch) -> StdResult<()> {
+        if let Some(threshold) = self.retention_limit {
+            self.connection
+                .apply(DeleteStakePoolQuery::below_epoch_threshold(
+                    epoch - threshold,
+                ))
+                .map_err(AdapterError::QueryError)?;
+        }
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::database::test_helper::{insert_stake_pool, main_db_connection};
 
     use super::*;
 
     #[tokio::test]
-    async fn save_protocol_parameters_prune_older_epoch_settings() {
+    async fn prune_epoch_settings_older_than_threshold() {
         let connection = main_db_connection().unwrap();
         const STAKE_POOL_PRUNE_EPOCH_THRESHOLD: u64 = 10;
         insert_stake_pool(&connection, &[1, 2]).unwrap();
         let store =
             StakePoolStore::new(Arc::new(connection), Some(STAKE_POOL_PRUNE_EPOCH_THRESHOLD));
 
         store
-            .save_stakes(
-                Epoch(2) + STAKE_POOL_PRUNE_EPOCH_THRESHOLD,
-                StakeDistribution::from_iter([("pool1".to_string(), 100)]),
-            )
+            .prune(Epoch(2) + STAKE_POOL_PRUNE_EPOCH_THRESHOLD)
             .await
-            .expect("saving stakes should not fails");
+            .unwrap();
+
         let epoch1_stakes = store.get_stakes(Epoch(1)).await.unwrap();
         let epoch2_stakes = store.get_stakes(Epoch(2)).await.unwrap();
 
@@ -129,6 +135,27 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn without_threshold_nothing_is_pruned() {
+        let connection = main_db_connection().unwrap();
+        insert_stake_pool(&connection, &[1, 2]).unwrap();
+        let store = StakePoolStore::new(Arc::new(connection), None);
+
+        store.prune(Epoch(100)).await.unwrap();
+
+        let epoch1_stakes = store.get_stakes(Epoch(1)).await.unwrap();
+        let epoch2_stakes = store.get_stakes(Epoch(2)).await.unwrap();
+
+        assert!(
+            epoch1_stakes.is_some(),
+            "Stakes at epoch 1 should have been pruned",
+        );
+        assert!(
+            epoch2_stakes.is_some(),
+            "Stakes at epoch 2 should still exist",
+        );
+    }
+
     #[tokio::test]
     async fn retrieve_with_no_stakes_returns_none() {
         let connection = main_db_connection().unwrap();

mithril-aggregator/src/dependency_injection/builder.rs

@@ -1337,13 +1337,15 @@ impl DependenciesBuilder {
     }
 
     async fn build_upkeep_service(&mut self) -> Result<Arc<dyn UpkeepService>> {
+        let stake_pool_pruning_task = self.get_stake_store().await?;
+
         let upkeep_service = Arc::new(AggregatorUpkeepService::new(
             self.get_sqlite_connection().await?,
             self.get_sqlite_connection_cardano_transaction_pool()
                 .await?,
             self.get_event_store_sqlite_connection().await?,
             self.get_signed_entity_lock().await?,
-            vec![], // TODO XXX Add pruning tasks here
+            vec![stake_pool_pruning_task],
             self.root_logger(),
         ));