Merge pull request #2375 from input-output-hk/dlachaume/2364/avoid-database-vacuum-blocking-aggregator-api

dlachaume · web-flow · commit 6408d5c5de3f · 2025-03-18T09:47:30.000+01:00
Feat: move database vacuum to aggregator startup to avoid API interruptions during epoch transitions
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/mithril-aggregator/Cargo.toml b/mithril-aggregator/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "mithril-aggregator"
-version = "0.7.16"
+version = "0.7.17"
 description = "A Mithril Aggregator server"
 authors = { workspace = true }
 edition = { workspace = true }
diff --git a/mithril-aggregator/src/commands/database_command.rs b/mithril-aggregator/src/commands/database_command.rs
@@ -26,12 +26,16 @@ impl DatabaseCommand {
 pub enum DatabaseSubCommand {
     /// Migrate databases located in the given stores directory
     Migrate(MigrateCommand),
+
+    /// Vacuum the aggregator main database
+    Vacuum(VacuumCommand),
 }
 
 impl DatabaseSubCommand {
     pub async fn execute(&self, root_logger: Logger) -> StdResult<()> {
         match self {
             Self::Migrate(cmd) => cmd.execute(root_logger).await,
+            Self::Vacuum(cmd) => cmd.execute(root_logger).await,
         }
     }
 }
@@ -79,3 +83,38 @@ impl MigrateCommand {
         Ok(())
     }
 }
+
+#[derive(Parser, Debug, Clone)]
+pub struct VacuumCommand {
+    /// Stores directory
+    #[clap(long, env = "STORES_DIRECTORY")]
+    stores_directory: PathBuf,
+}
+
+impl VacuumCommand {
+    pub async fn execute(&self, root_logger: Logger) -> StdResult<()> {
+        let config = Configuration {
+            environment: ExecutionEnvironment::Production,
+            data_stores_directory: self.stores_directory.clone(),
+            // Temporary solution to avoid the need to provide a full configuration
+            ..Configuration::new_sample(std::env::temp_dir())
+        };
+        debug!(root_logger, "DATABASE VACUUM command"; "config" => format!("{config:?}"));
+        println!(
+            "Vacuuming database from stores directory: {}",
+            self.stores_directory.to_string_lossy()
+        );
+        let mut dependencies_builder =
+            DependenciesBuilder::new(root_logger.clone(), config.clone());
+
+        dependencies_builder
+            .get_upkeep_service()
+            .await
+            .with_context(|| "Dependencies Builder can not get upkeep service")?
+            .vacuum()
+            .await
+            .with_context(|| "Upkeep service can not vacuum")?;
+
+        Ok(())
+    }
+}
diff --git a/mithril-aggregator/src/commands/serve_command.rs b/mithril-aggregator/src/commands/serve_command.rs
@@ -1,14 +1,22 @@
+use std::{
+    net::IpAddr,
+    path::{Path, PathBuf},
+    time::Duration,
+};
+
 use anyhow::{anyhow, Context};
+use chrono::TimeDelta;
 use clap::Parser;
 use config::{builder::DefaultState, ConfigBuilder, Map, Source, Value, ValueKind};
-use mithril_common::StdResult;
-use mithril_metric::MetricsServer;
 use slog::{crit, debug, info, warn, Logger};
-use std::time::Duration;
-use std::{net::IpAddr, path::PathBuf};
 use tokio::{sync::oneshot, task::JoinSet};
 
-use crate::{dependency_injection::DependenciesBuilder, Configuration};
+use mithril_common::StdResult;
+use mithril_metric::MetricsServer;
+
+use crate::{dependency_injection::DependenciesBuilder, tools::VacuumTracker, Configuration};
+
+const VACUUM_MINIMUM_INTERVAL: TimeDelta = TimeDelta::weeks(1);
 
 /// Server runtime mode
 #[derive(Parser, Debug, Clone)]
@@ -170,6 +178,15 @@ impl ServeCommand {
             .with_context(|| "Dependencies Builder can not create event store")?;
         let event_store_thread = tokio::spawn(async move { event_store.run().await.unwrap() });
 
+        // start the database vacuum operation, if needed
+        self.perform_database_vacuum_if_needed(
+            &config.data_stores_directory,
+            &mut dependencies_builder,
+            VACUUM_MINIMUM_INTERVAL,
+            &root_logger,
+        )
+        .await?;
+
         // start the aggregator runtime
         let mut runtime = dependencies_builder
             .create_aggregator_runner()
@@ -297,4 +314,51 @@ impl ServeCommand {
 
         Ok(())
     }
+
+    /// This function checks if a database vacuum is needed and performs it if necessary.
+    ///
+    /// Errors from [VacuumTracker] operations are logged but not propagated as errors.
+    async fn perform_database_vacuum_if_needed(
+        &self,
+        store_dir: &Path,
+        dependencies_builder: &mut DependenciesBuilder,
+        vacuum_min_interval: TimeDelta,
+        logger: &Logger,
+    ) -> StdResult<()> {
+        let vacuum_tracker = VacuumTracker::new(store_dir, vacuum_min_interval, logger.clone());
+        match vacuum_tracker.check_vacuum_needed() {
+            Ok((true, _)) => {
+                info!(logger, "Performing vacuum");
+
+                let upkeep = dependencies_builder
+                    .get_upkeep_service()
+                    .await
+                    .with_context(|| "Dependencies Builder can not create upkeep")?;
+
+                upkeep
+                    .vacuum()
+                    .await
+                    .with_context(|| "Upkeep service failed to vacuum database")?;
+
+                match vacuum_tracker.update_last_vacuum_time() {
+                    Ok(last_vacuum) => {
+                        info!(logger, "Vacuum performed"; "last_vacuum" => last_vacuum.to_rfc3339());
+                    }
+                    Err(e) => {
+                        warn!(logger, "Failed to update last vacuum time"; "error" => ?e);
+                    }
+                }
+            }
+            Ok((false, last_vacuum)) => {
+                let time_display =
+                    last_vacuum.map_or_else(|| "never".to_string(), |time| time.to_rfc3339());
+                info!(logger, "No vacuum needed"; "last_vacuum" => time_display);
+            }
+            Err(e) => {
+                warn!(logger, "Failed to check if vacuum is needed"; "error" => ?e);
+            }
+        }
+
+        Ok(())
+    }
 }
diff --git a/mithril-aggregator/src/services/upkeep.rs b/mithril-aggregator/src/services/upkeep.rs
@@ -24,6 +24,9 @@ use slog::{info, Logger};
 pub trait UpkeepService: Send + Sync {
     /// Run the upkeep service.
     async fn run(&self, epoch: Epoch) -> StdResult<()>;
+
+    /// Vacuum database.
+    async fn vacuum(&self) -> StdResult<()>;
 }
 
 /// Define the task responsible for pruning a datasource below a certain epoch threshold.
@@ -101,10 +104,7 @@ impl AggregatorUpkeepService {
             info!(db_upkeep_logger, "Cleaning main database");
             SqliteCleaner::new(&main_db_connection)
                 .with_logger(db_upkeep_logger.clone())
-                .with_tasks(&[
-                    SqliteCleaningTask::Vacuum,
-                    SqliteCleaningTask::WalCheckpointTruncate,
-                ])
+                .with_tasks(&[SqliteCleaningTask::WalCheckpointTruncate])
                 .run()?;
 
             info!(db_upkeep_logger, "Cleaning cardano transactions database");
@@ -127,6 +127,34 @@ impl AggregatorUpkeepService {
             .await
             .with_context(|| "Database Upkeep thread crashed")?
     }
+
+    async fn vacuum_main_database(&self) -> StdResult<()> {
+        if self.signed_entity_type_lock.has_locked_entities().await {
+            info!(
+                self.logger,
+                "Some entities are locked - Skipping main database vacuum"
+            );
+            return Ok(());
+        }
+
+        let main_db_connection = self.main_db_connection.clone();
+        let db_upkeep_logger = self.logger.clone();
+
+        // Run the database upkeep tasks in another thread to avoid blocking the tokio runtime
+        let db_upkeep_thread = tokio::task::spawn_blocking(move || -> StdResult<()> {
+            info!(db_upkeep_logger, "Vacuum main database");
+            SqliteCleaner::new(&main_db_connection)
+                .with_logger(db_upkeep_logger.clone())
+                .with_tasks(&[SqliteCleaningTask::Vacuum])
+                .run()?;
+
+            Ok(())
+        });
+
+        db_upkeep_thread
+            .await
+            .with_context(|| "Database Upkeep thread crashed")?
+    }
 }
 
 #[async_trait]
@@ -145,12 +173,23 @@ impl UpkeepService for AggregatorUpkeepService {
         info!(self.logger, "Upkeep finished");
         Ok(())
     }
+
+    async fn vacuum(&self) -> StdResult<()> {
+        info!(self.logger, "Start database vacuum");
+
+        self.vacuum_main_database()
+            .await
+            .with_context(|| "Vacuuming main database failed")?;
+
+        info!(self.logger, "Vacuum finished");
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use mithril_common::entities::SignedEntityTypeDiscriminants;
-    use mithril_common::test_utils::TempDir;
+    use mithril_common::{entities::SignedEntityTypeDiscriminants, temp_dir_create};
     use mockall::predicate::eq;
 
     use crate::database::test_helper::{
@@ -187,7 +226,7 @@ mod tests {
     #[tokio::test]
     async fn test_cleanup_database() {
         let (main_db_path, ctx_db_path, event_store_db_path, log_path) = {
-            let db_dir = TempDir::create("aggregator_upkeep", "test_cleanup_database");
+            let db_dir = temp_dir_create!();
             (
                 db_dir.join("main.db"),
                 db_dir.join("cardano_tx.db"),
@@ -218,27 +257,23 @@ mod tests {
 
         let logs = std::fs::read_to_string(&log_path).unwrap();
 
-        assert_eq!(
-            logs.matches(SqliteCleaningTask::Vacuum.log_message())
-                .count(),
-            1,
-            "Should have run only once since only the main database has a `Vacuum` cleanup"
-        );
         assert_eq!(
             logs.matches(SqliteCleaningTask::WalCheckpointTruncate.log_message())
                 .count(),
             3,
             "Should have run three times since the three databases have a `WalCheckpointTruncate` cleanup"
         );
+        assert_eq!(
+            logs.matches(SqliteCleaningTask::Vacuum.log_message())
+                .count(),
+            0,
+            "Upkeep operation should not include Vacuum tasks"
+        );
     }
 
     #[tokio::test]
     async fn test_doesnt_cleanup_db_if_any_entity_is_locked() {
-        let log_path = TempDir::create(
-            "aggregator_upkeep",
-            "test_doesnt_cleanup_db_if_any_entity_is_locked",
-        )
-        .join("upkeep.log");
+        let log_path = temp_dir_create!().join("upkeep.log");
 
         let signed_entity_type_lock = Arc::new(SignedEntityTypeLock::default());
         signed_entity_type_lock
@@ -257,11 +292,6 @@ mod tests {
 
         let logs = std::fs::read_to_string(&log_path).unwrap();
 
-        assert_eq!(
-            logs.matches(SqliteCleaningTask::Vacuum.log_message())
-                .count(),
-            0,
-        );
         assert_eq!(
             logs.matches(SqliteCleaningTask::WalCheckpointTruncate.log_message())
                 .count(),
@@ -290,4 +320,59 @@ mod tests {
 
         service.run(Epoch(14)).await.expect("Upkeep service failed");
     }
+
+    #[tokio::test]
+    async fn test_doesnt_vacuum_db_if_any_entity_is_locked() {
+        let log_path = temp_dir_create!().join("vacuum.log");
+
+        let signed_entity_type_lock = Arc::new(SignedEntityTypeLock::default());
+        signed_entity_type_lock
+            .lock(SignedEntityTypeDiscriminants::CardanoTransactions)
+            .await;
+
+        // Separate block to force log flushing by dropping the service that owns the logger
+        {
+            let service = AggregatorUpkeepService {
+                signed_entity_type_lock: signed_entity_type_lock.clone(),
+                logger: TestLogger::file(&log_path),
+                ..default_upkeep_service()
+            };
+            service.vacuum().await.expect("Vacuum failed");
+        }
+
+        let logs = std::fs::read_to_string(&log_path).unwrap();
+
+        assert_eq!(
+            logs.matches(SqliteCleaningTask::Vacuum.log_message())
+                .count(),
+            0,
+        );
+    }
+
+    #[tokio::test]
+    async fn test_vacuum_database() {
+        let log_path = temp_dir_create!().join("vacuum.log");
+
+        let signed_entity_type_lock = Arc::new(SignedEntityTypeLock::default());
+        signed_entity_type_lock
+            .lock(SignedEntityTypeDiscriminants::CardanoTransactions)
+            .await;
+
+        // Separate block to force log flushing by dropping the service that owns the logger
+        {
+            let service = AggregatorUpkeepService {
+                logger: TestLogger::file(&log_path),
+                ..default_upkeep_service()
+            };
+            service.vacuum().await.expect("Vacuum failed");
+        }
+
+        let logs = std::fs::read_to_string(&log_path).unwrap();
+
+        assert_eq!(
+            logs.matches(SqliteCleaningTask::Vacuum.log_message())
+                .count(),
+            1,
+        );
+    }
 }
diff --git a/mithril-aggregator/src/tools/mod.rs b/mithril-aggregator/src/tools/mod.rs
@@ -9,6 +9,7 @@ pub mod mocks;
 mod signer_importer;
 mod single_signature_authenticator;
 pub mod url_sanitizer;
+mod vacuum_tracker;
 
 pub use certificates_hash_migrator::CertificatesHashMigrator;
 pub use digest_helpers::extract_digest_from_path;
@@ -18,6 +19,7 @@ pub use signer_importer::{
     CExplorerSignerRetriever, SignersImporter, SignersImporterPersister, SignersImporterRetriever,
 };
 pub use single_signature_authenticator::*;
+pub use vacuum_tracker::VacuumTracker;
 
 /// Downcast the error to the specified error type and check if the error satisfies the condition.
 pub(crate) fn downcast_check<E>(
diff --git a/mithril-aggregator/src/tools/vacuum_tracker.rs b/mithril-aggregator/src/tools/vacuum_tracker.rs
