refactor: make upkeep pruning tasks apply themself the retention limit

Alenar · Alenar · commit 8fdf36beab77 · 2024-10-02T17:20:25.000+02:00
Instead of making the upkeep service itself compute the threshold based
on the limit.

This design allow each task to use different limit if needed.
diff --git a/mithril-signer/src/database/repository/signed_beacon_repository.rs b/mithril-signer/src/database/repository/signed_beacon_repository.rs
@@ -15,12 +15,16 @@ use crate::services::{EpochPruningTask, SignedBeaconStore};
 /// A [SignedBeaconStore] implementation using SQLite.
 pub struct SignedBeaconRepository {
     connection: Arc<SqliteConnection>,
+    store_retention_limit: Option<u64>,
 }
 
 impl SignedBeaconRepository {
     /// Create a new instance of the `SignedBeaconRepository`.
-    pub fn new(connection: Arc<SqliteConnection>) -> Self {
-        Self { connection }
+    pub fn new(connection: Arc<SqliteConnection>, store_retention_limit: Option<u64>) -> Self {
+        Self {
+            connection,
+            store_retention_limit,
+        }
     }
 
     /// Get the last signed beacon.
@@ -71,8 +75,14 @@ impl EpochPruningTask for SignedBeaconRepository {
         "Signed Beacon"
     }
 
-    async fn prune_below_epoch_threshold(&self, epoch_threshold: Epoch) -> StdResult<()> {
-        self.prune_below_epoch(epoch_threshold)
+    async fn prune(&self, current_epoch: Epoch) -> StdResult<()> {
+        match self
+            .store_retention_limit
+            .map(|limit| current_epoch - limit)
+        {
+            Some(threshold) if *threshold > 0 => self.prune_below_epoch(threshold),
+            _ => Ok(()),
+        }
     }
 }
 
@@ -102,7 +112,7 @@ mod tests {
     #[test]
     fn get_last_stored_signed_beacon() {
         let connection = Arc::new(main_db_connection().unwrap());
-        let repository = SignedBeaconRepository::new(connection.clone());
+        let repository = SignedBeaconRepository::new(connection.clone(), None);
 
         let last_signed_beacon = repository.get_last().unwrap();
         assert_eq!(None, last_signed_beacon);
@@ -151,7 +161,7 @@ mod tests {
     #[tokio::test]
     async fn filter_out_nothing_if_nothing_was_previously_signed() {
         let connection = Arc::new(main_db_connection().unwrap());
-        let repository = SignedBeaconRepository::new(connection.clone());
+        let repository = SignedBeaconRepository::new(connection.clone(), None);
 
         let to_filter = all_signed_entity_type_for(&TimePoint::dummy());
         let available_entities = repository
@@ -165,7 +175,7 @@ mod tests {
     #[tokio::test]
     async fn filter_out_nothing_if_previously_signed_entities_doesnt_match_passed_entities() {
         let connection = Arc::new(main_db_connection().unwrap());
-        let repository = SignedBeaconRepository::new(connection.clone());
+        let repository = SignedBeaconRepository::new(connection.clone(), None);
 
         let time_point = TimePoint::dummy();
         insert_signed_beacons(
@@ -189,7 +199,7 @@ mod tests {
     #[tokio::test]
     async fn filter_out_everything_if_previously_signed_entities_match_all_passed_entities() {
         let connection = Arc::new(main_db_connection().unwrap());
-        let repository = SignedBeaconRepository::new(connection.clone());
+        let repository = SignedBeaconRepository::new(connection.clone(), None);
 
         let to_filter = all_signed_entity_type_for(&TimePoint::dummy());
         insert_signed_beacons(
@@ -210,7 +220,7 @@ mod tests {
     #[tokio::test]
     async fn filter_out_partially_if_some_previously_signed_entities_match_passed_entities() {
         let connection = Arc::new(main_db_connection().unwrap());
-        let repository = SignedBeaconRepository::new(connection.clone());
+        let repository = SignedBeaconRepository::new(connection.clone(), None);
 
         let time_point = TimePoint::dummy();
         let signed_beacons = [
@@ -253,7 +263,7 @@ mod tests {
     #[tokio::test]
     async fn mark_beacon_as_signed() {
         let connection = Arc::new(main_db_connection().unwrap());
-        let repository = SignedBeaconRepository::new(connection.clone());
+        let repository = SignedBeaconRepository::new(connection.clone(), None);
 
         let beacon_to_sign = BeaconToSign {
             epoch: Epoch(13),
@@ -278,10 +288,50 @@ mod tests {
         assert_eq!(beacon_to_sign, signed_beacon);
     }
 
-    #[test]
-    fn test_prune() {
+    #[tokio::test]
+    async fn test_dont_execute_pruning_tasks_if_no_retention_limit_set() {
+        let connection = Arc::new(main_db_connection().unwrap());
+        let repository = SignedBeaconRepository::new(connection.clone(), None);
+        insert_signed_beacons(
+            &connection,
+            SignedBeaconRecord::fakes(&[(
+                Epoch(8),
+                vec![SignedEntityType::MithrilStakeDistribution(Epoch(8))],
+            )]),
+        );
+
+        EpochPruningTask::prune(&repository, Epoch(1000))
+            .await
+            .unwrap();
+
+        let cursor = connection.fetch(GetSignedBeaconQuery::all()).unwrap();
+        assert_eq!(1, cursor.count(),);
+    }
+
+    #[tokio::test]
+    async fn test_dont_execute_pruning_tasks_if_current_epoch_minus_retention_limit_is_0() {
         let connection = Arc::new(main_db_connection().unwrap());
-        let repository = SignedBeaconRepository::new(connection.clone());
+        let repository = SignedBeaconRepository::new(connection.clone(), Some(10));
+        insert_signed_beacons(
+            &connection,
+            SignedBeaconRecord::fakes(&[(
+                Epoch(8),
+                vec![SignedEntityType::MithrilStakeDistribution(Epoch(8))],
+            )]),
+        );
+
+        EpochPruningTask::prune(&repository, Epoch(9))
+            .await
+            .unwrap();
+
+        let cursor = connection.fetch(GetSignedBeaconQuery::all()).unwrap();
+        assert_eq!(1, cursor.count(),);
+    }
+
+    #[tokio::test]
+    async fn test_prune_task_substract_set_retention_limit_to_given_epoch() {
+        let connection = Arc::new(main_db_connection().unwrap());
+        let repository = SignedBeaconRepository::new(connection.clone(), Some(10));
         insert_signed_beacons(
             &connection,
             SignedBeaconRecord::fakes(&[
@@ -299,7 +349,9 @@ mod tests {
             ]),
         );
 
-        repository.prune_below_epoch(Epoch(8)).unwrap();
+        EpochPruningTask::prune(&repository, Epoch(18))
+            .await
+            .unwrap();
 
         let signed_beacons: Vec<SignedBeaconRecord> = connection
             .fetch_collect(GetSignedBeaconQuery::all())
diff --git a/mithril-signer/src/dependency_injection/builder.rs b/mithril-signer/src/dependency_injection/builder.rs
@@ -347,14 +347,15 @@ impl<'a> DependenciesBuilder<'a> {
             slog_scope::logger(),
             Arc::new(preloader_activation),
         ));
-        let signed_beacon_repository =
-            Arc::new(SignedBeaconRepository::new(sqlite_connection.clone()));
+        let signed_beacon_repository = Arc::new(SignedBeaconRepository::new(
+            sqlite_connection.clone(),
+            self.config.store_retention_limit.map(|limit| limit as u64),
+        ));
         let upkeep_service = Arc::new(SignerUpkeepService::new(
             sqlite_connection.clone(),
             sqlite_connection_cardano_transaction_pool,
             signed_entity_type_lock.clone(),
             vec![signed_beacon_repository.clone()],
-            self.config.store_retention_limit.map(|limit| limit as u64),
             slog_scope::logger(),
         ));
         let certifier = Arc::new(SignerCertifierService::new(
diff --git a/mithril-signer/src/runtime/runner.rs b/mithril-signer/src/runtime/runner.rs
@@ -642,7 +642,7 @@ mod tests {
         let upkeep_service = Arc::new(MockUpkeepService::new());
         let certifier = Arc::new(SignerCertifierService::new(
             ticker_service.clone(),
-            Arc::new(SignedBeaconRepository::new(sqlite_connection.clone())),
+            Arc::new(SignedBeaconRepository::new(sqlite_connection.clone(), None)),
             Arc::new(SignerSignedEntityConfigProvider::new(
                 network,
                 epoch_service.clone(),
diff --git a/mithril-signer/src/services/upkeep_service.rs b/mithril-signer/src/services/upkeep_service.rs
@@ -33,8 +33,8 @@ pub trait EpochPruningTask: Send + Sync {
     /// Get the name of the data that will be pruned.
     fn pruned_data(&self) -> &'static str;
 
-    /// Prune the datasource below the given epoch threshold.
-    async fn prune_below_epoch_threshold(&self, epoch_threshold: Epoch) -> StdResult<()>;
+    /// Prune the datasource based on the given current epoch.
+    async fn prune(&self, current_epoch: Epoch) -> StdResult<()>;
 }
 
 /// Implementation of the upkeep service for the signer.
@@ -46,7 +46,6 @@ pub struct SignerUpkeepService {
     cardano_tx_connection_pool: Arc<SqliteConnectionPool>,
     signed_entity_type_lock: Arc<SignedEntityTypeLock>,
     pruning_tasks: Vec<Arc<dyn EpochPruningTask>>,
-    store_retention_limit: Option<u64>,
     logger: Logger,
 }
 
@@ -57,37 +56,27 @@ impl SignerUpkeepService {
         cardano_tx_connection_pool: Arc<SqliteConnectionPool>,
         signed_entity_type_lock: Arc<SignedEntityTypeLock>,
         pruning_tasks: Vec<Arc<dyn EpochPruningTask>>,
-        store_retention_limit: Option<u64>,
         logger: Logger,
     ) -> Self {
         Self {
             main_db_connection,
             cardano_tx_connection_pool,
             signed_entity_type_lock,
             pruning_tasks,
-            store_retention_limit,
             logger,
         }
     }
 
     async fn execute_pruning_tasks(&self, current_epoch: Epoch) -> StdResult<()> {
-        match self
-            .store_retention_limit
-            .map(|limit| current_epoch - limit)
-        {
-            Some(threshold) if *threshold > 0 => {
-                for task in &self.pruning_tasks {
-                    info!(
-                        self.logger, "UpkeepService::Pruning stale data";
-                        "pruned_data" => task.pruned_data(), "below_epoch_threshold" => ?threshold
-                    );
-
-                    task.prune_below_epoch_threshold(threshold).await?;
-                }
-                Ok(())
-            }
-            _ => Ok(()),
+        for task in &self.pruning_tasks {
+            info!(
+                self.logger, "UpkeepService::Pruning stale data";
+                "pruned_data" => task.pruned_data(), "current_epoch" => ?current_epoch
+            );
+            task.prune(current_epoch).await?;
         }
+
+        Ok(())
     }
 
     async fn upkeep_all_databases(&self) -> StdResult<()> {
@@ -198,7 +187,6 @@ mod tests {
                 )),
                 Arc::new(SignedEntityTypeLock::default()),
                 vec![],
-                None,
                 TestLogger::file(&log_path),
             );
 
@@ -241,7 +229,6 @@ mod tests {
                 Arc::new(SqliteConnectionPool::build(1, cardano_tx_db_connection).unwrap()),
                 signed_entity_type_lock.clone(),
                 vec![],
-                None,
                 TestLogger::file(&log_path),
             );
 
@@ -265,15 +252,15 @@ mod tests {
     #[tokio::test]
     async fn test_execute_all_pruning_tasks() {
         let task1 = mock_epoch_pruning_task(|mock| {
-            mock.expect_prune_below_epoch_threshold()
+            mock.expect_prune()
                 .once()
-                .with(eq(Epoch(4)))
+                .with(eq(Epoch(14)))
                 .returning(|_| Ok(()));
         });
         let task2 = mock_epoch_pruning_task(|mock| {
-            mock.expect_prune_below_epoch_threshold()
+            mock.expect_prune()
                 .once()
-                .with(eq(Epoch(4)))
+                .with(eq(Epoch(14)))
                 .returning(|_| Ok(()));
         });
 
@@ -282,52 +269,9 @@ mod tests {
             Arc::new(SqliteConnectionPool::build(1, cardano_tx_db_connection).unwrap()),
             Arc::new(SignedEntityTypeLock::default()),
             vec![task1, task2],
-            Some(10),
             TestLogger::stdout(),
         );
 
         service.run(Epoch(14)).await.expect("Upkeep service failed");
     }
-
-    #[tokio::test]
-    async fn test_dont_execute_pruning_tasks_if_current_epoch_minus_retention_limit_is_0() {
-        let task = mock_epoch_pruning_task(|mock| {
-            mock.expect_prune_below_epoch_threshold()
-                .never()
-                .returning(|_| Ok(()));
-        });
-
-        let service = SignerUpkeepService::new(
-            Arc::new(main_db_connection().unwrap()),
-            Arc::new(SqliteConnectionPool::build(1, cardano_tx_db_connection).unwrap()),
-            Arc::new(SignedEntityTypeLock::default()),
-            vec![task],
-            Some(17),
-            TestLogger::stdout(),
-        );
-
-        service.run(Epoch(1)).await.expect("Upkeep service failed");
-        service.run(Epoch(16)).await.expect("Upkeep service failed");
-        service.run(Epoch(17)).await.expect("Upkeep service failed");
-    }
-
-    #[tokio::test]
-    async fn test_dont_execute_pruning_tasks_if_no_retention_limit_set() {
-        let task = mock_epoch_pruning_task(|mock| {
-            mock.expect_prune_below_epoch_threshold()
-                .never()
-                .returning(|_| Ok(()));
-        });
-
-        let service = SignerUpkeepService::new(
-            Arc::new(main_db_connection().unwrap()),
-            Arc::new(SqliteConnectionPool::build(1, cardano_tx_db_connection).unwrap()),
-            Arc::new(SignedEntityTypeLock::default()),
-            vec![task],
-            None,
-            TestLogger::stdout(),
-        );
-
-        service.run(Epoch(29)).await.expect("Upkeep service failed");
-    }
 }
diff --git a/mithril-signer/tests/test_extensions/state_machine_tester.rs b/mithril-signer/tests/test_extensions/state_machine_tester.rs
@@ -247,11 +247,10 @@ impl StateMachineTester {
             sqlite_connection_cardano_transaction_pool,
             signed_entity_type_lock.clone(),
             vec![],
-            None,
             slog_scope::logger(),
         ));
         let signed_beacon_repository =
-            Arc::new(SignedBeaconRepository::new(sqlite_connection.clone()));
+            Arc::new(SignedBeaconRepository::new(sqlite_connection.clone(), None));
         let certifier = Arc::new(SignerCertifierService::new(
             ticker_service.clone(),
             signed_beacon_repository.clone(),
