Implement pruning behavior in the aggregator

Author: sfauvel

mithril-aggregator/src/dependency_injection/builder.rs

@@ -1343,6 +1343,7 @@ impl DependenciesBuilder {
                 .await?,
             self.get_event_store_sqlite_connection().await?,
             self.get_signed_entity_lock().await?,
+            vec![], // TODO XXX Add pruning tasks here
             self.root_logger(),
         ));
 

mithril-aggregator/src/runtime/runner.rs

@@ -110,7 +110,7 @@ pub trait AggregatorRunnerTrait: Sync + Send {
     async fn inform_new_epoch(&self, epoch: Epoch) -> StdResult<()>;
 
     /// Perform the upkeep tasks.
-    async fn upkeep(&self) -> StdResult<()>;
+    async fn upkeep(&self, epoch: Epoch) -> StdResult<()>;
 
     /// Precompute what doesn't change for the actual epoch
     async fn precompute_epoch_data(&self) -> StdResult<()>;
@@ -522,9 +522,9 @@ impl AggregatorRunnerTrait for AggregatorRunner {
         Ok(())
     }
 
-    async fn upkeep(&self) -> StdResult<()> {
+    async fn upkeep(&self, epoch: Epoch) -> StdResult<()> {
         debug!(self.logger, ">> upkeep");
-        self.dependencies.upkeep_service.run().await
+        self.dependencies.upkeep_service.run(epoch).await
     }
 
     async fn create_open_message(
@@ -1010,16 +1010,20 @@ pub mod tests {
     }
 
     #[tokio::test]
-    async fn test_upkeep() {
+    async fn test_upkeep_calls_run_on_upkeep_service() {
         let mut upkeep_service = MockUpkeepService::new();
-        upkeep_service.expect_run().returning(|| Ok(())).times(1);
+        upkeep_service
+            .expect_run()
+            .with(eq(Epoch(5)))
+            .returning(|_| Ok(()))
+            .times(1);
 
         let mut deps = initialize_dependencies().await;
         deps.upkeep_service = Arc::new(upkeep_service);
 
         let runner = AggregatorRunner::new(Arc::new(deps));
 
-        runner.upkeep().await.unwrap();
+        runner.upkeep(Epoch(5)).await.unwrap();
     }
 
     #[tokio::test]

mithril-aggregator/src/runtime/state_machine.rs

@@ -245,7 +245,7 @@ impl AggregatorRuntime {
         trace!(self.logger, "Trying transition from IDLE to READY state");
 
         if maybe_current_time_point.is_none()
-            || maybe_current_time_point.unwrap().epoch < new_time_point.epoch
+            || maybe_current_time_point.as_ref().unwrap().epoch < new_time_point.epoch
         {
             self.runner.close_signer_registration_round().await?;
             self.runner
@@ -256,7 +256,7 @@ impl AggregatorRuntime {
                 .update_stake_distribution(&new_time_point)
                 .await?;
             self.runner.inform_new_epoch(new_time_point.epoch).await?;
-            self.runner.upkeep().await?;
+            self.runner.upkeep(new_time_point.epoch).await?;
             self.runner
                 .open_signer_registration_round(&new_time_point)
                 .await?;
@@ -450,7 +450,11 @@ mod tests {
             .expect_precompute_epoch_data()
             .once()
             .returning(|| Ok(()));
-        runner.expect_upkeep().once().returning(|| Ok(()));
+        runner
+            .expect_upkeep()
+            .with(predicate::eq(TimePoint::dummy().epoch))
+            .once()
+            .returning(|_| Ok(()));
         runner
             .expect_increment_runtime_cycle_total_since_startup_counter()
             .once()
@@ -514,7 +518,11 @@ mod tests {
             .expect_precompute_epoch_data()
             .once()
             .returning(|| Ok(()));
-        runner.expect_upkeep().once().returning(|| Ok(()));
+        runner
+            .expect_upkeep()
+            .with(predicate::eq(TimePoint::dummy().epoch))
+            .once()
+            .returning(|_| Ok(()));
         runner
             .expect_increment_runtime_cycle_total_since_startup_counter()
             .once()

mithril-aggregator/src/services/upkeep.rs

@@ -9,21 +9,32 @@ use std::sync::Arc;
 
 use anyhow::Context;
 use async_trait::async_trait;
-use slog::{info, Logger};
-
+use mithril_common::entities::Epoch;
 use mithril_common::logging::LoggerExtensions;
 use mithril_common::signed_entity_type_lock::SignedEntityTypeLock;
 use mithril_common::StdResult;
 use mithril_persistence::sqlite::{
     SqliteCleaner, SqliteCleaningTask, SqliteConnection, SqliteConnectionPool,
 };
+use slog::{info, Logger};
 
 /// Define the service responsible for the upkeep of the application.
 #[cfg_attr(test, mockall::automock)]
 #[async_trait]
 pub trait UpkeepService: Send + Sync {
     /// Run the upkeep service.
-    async fn run(&self) -> StdResult<()>;
+    async fn run(&self, epoch: Epoch) -> StdResult<()>;
+}
+
+/// Define the task responsible for pruning a datasource below a certain epoch threshold.
+#[cfg_attr(test, mockall::automock)]
+#[async_trait]
+pub trait EpochPruningTask: Send + Sync {
+    /// Get the name of the data that will be pruned.
+    fn pruned_data(&self) -> &'static str;
+
+    /// Prune the datasource based on the given current epoch.
+    async fn prune(&self, current_epoch: Epoch) -> StdResult<()>;
 }
 
 /// Implementation of the upkeep service for the aggregator.
@@ -35,6 +46,7 @@ pub struct AggregatorUpkeepService {
     cardano_tx_connection_pool: Arc<SqliteConnectionPool>,
     event_store_connection: Arc<SqliteConnection>,
     signed_entity_type_lock: Arc<SignedEntityTypeLock>,
+    pruning_tasks: Vec<Arc<dyn EpochPruningTask>>,
     logger: Logger,
 }
 
@@ -45,17 +57,31 @@ impl AggregatorUpkeepService {
         cardano_tx_connection_pool: Arc<SqliteConnectionPool>,
         event_store_connection: Arc<SqliteConnection>,
         signed_entity_type_lock: Arc<SignedEntityTypeLock>,
+        pruning_tasks: Vec<Arc<dyn EpochPruningTask>>,
         logger: Logger,
     ) -> Self {
         Self {
             main_db_connection,
             cardano_tx_connection_pool,
             event_store_connection,
             signed_entity_type_lock,
+            pruning_tasks,
             logger: logger.new_with_component_name::<Self>(),
         }
     }
 
+    async fn execute_pruning_tasks(&self, current_epoch: Epoch) -> StdResult<()> {
+        for task in &self.pruning_tasks {
+            info!(
+                self.logger, "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<()> {
         if self.signed_entity_type_lock.has_locked_entities().await {
             info!(
@@ -105,9 +131,13 @@ impl AggregatorUpkeepService {
 
 #[async_trait]
 impl UpkeepService for AggregatorUpkeepService {
-    async fn run(&self) -> StdResult<()> {
+    async fn run(&self, current_epoch: Epoch) -> StdResult<()> {
         info!(self.logger, "Start upkeep of the application");
 
+        self.execute_pruning_tasks(current_epoch)
+            .await
+            .with_context(|| "Pruning tasks failed")?;
+
         self.upkeep_all_databases()
             .await
             .with_context(|| "Database upkeep failed")?;
@@ -121,6 +151,7 @@ impl UpkeepService for AggregatorUpkeepService {
 mod tests {
     use mithril_common::entities::SignedEntityTypeDiscriminants;
     use mithril_common::test_utils::TempDir;
+    use mockall::predicate::eq;
 
     use crate::database::test_helper::{
         cardano_tx_db_connection, cardano_tx_db_file_connection, main_db_connection,
@@ -133,6 +164,26 @@ mod tests {
 
     use super::*;
 
+    fn mock_epoch_pruning_task(
+        mock_config: impl FnOnce(&mut MockEpochPruningTask),
+    ) -> Arc<dyn EpochPruningTask> {
+        let mut task_mock = MockEpochPruningTask::new();
+        task_mock.expect_pruned_data().return_const("mock_data");
+        mock_config(&mut task_mock);
+        Arc::new(task_mock)
+    }
+
+    fn default_upkeep_service() -> AggregatorUpkeepService {
+        AggregatorUpkeepService::new(
+            Arc::new(main_db_connection().unwrap()),
+            Arc::new(SqliteConnectionPool::build(1, cardano_tx_db_connection).unwrap()),
+            Arc::new(event_store_db_connection().unwrap()),
+            Arc::new(SignedEntityTypeLock::default()),
+            vec![],
+            TestLogger::stdout(),
+        )
+    }
+
     #[tokio::test]
     async fn test_cleanup_database() {
         let (main_db_path, ctx_db_path, event_store_db_path, log_path) = {
@@ -158,10 +209,11 @@ mod tests {
                 )),
                 Arc::new(event_store_connection),
                 Arc::new(SignedEntityTypeLock::default()),
+                vec![],
                 TestLogger::file(&log_path),
             );
 
-            service.run().await.expect("Upkeep service failed");
+            service.run(Epoch(5)).await.expect("Upkeep service failed");
         }
 
         let logs = std::fs::read_to_string(&log_path).unwrap();
@@ -195,14 +247,12 @@ mod tests {
 
         // Separate block to force log flushing by dropping the service that owns the logger
         {
-            let service = AggregatorUpkeepService::new(
-                Arc::new(main_db_connection().unwrap()),
-                Arc::new(SqliteConnectionPool::build(1, cardano_tx_db_connection).unwrap()),
-                Arc::new(event_store_db_connection().unwrap()),
-                signed_entity_type_lock.clone(),
-                TestLogger::file(&log_path),
-            );
-            service.run().await.expect("Upkeep service failed");
+            let service = AggregatorUpkeepService {
+                signed_entity_type_lock: signed_entity_type_lock.clone(),
+                logger: TestLogger::file(&log_path),
+                ..default_upkeep_service()
+            };
+            service.run(Epoch(5)).await.expect("Upkeep service failed");
         }
 
         let logs = std::fs::read_to_string(&log_path).unwrap();
@@ -218,4 +268,26 @@ mod tests {
             0,
         );
     }
+    #[tokio::test]
+    async fn test_execute_all_pruning_tasks() {
+        let task1 = mock_epoch_pruning_task(|mock| {
+            mock.expect_prune()
+                .once()
+                .with(eq(Epoch(14)))
+                .returning(|_| Ok(()));
+        });
+        let task2 = mock_epoch_pruning_task(|mock| {
+            mock.expect_prune()
+                .once()
+                .with(eq(Epoch(14)))
+                .returning(|_| Ok(()));
+        });
+
+        let service = AggregatorUpkeepService {
+            pruning_tasks: vec![task1, task2],
+            ..default_upkeep_service()
+        };
+
+        service.run(Epoch(14)).await.expect("Upkeep service failed");
+    }
 }