Bug 2055182 - Run database maintenance after migration

glean-core/src/database/sqlite.rs

@@ -171,11 +171,13 @@ impl Database {
 
         if sqlite_exists {
             log::debug!("SQLite database already exists. Not trying to migrate Rkv");
+            db.run_maintenance(false)?;
         } else {
             match migration::try_migrate(&path, &db) {
                 Ok(Some(state)) => {
                     log::debug!("Migration done. state={state:?}");
                     db.migration_state = Some(state);
+                    db.run_maintenance(true)?;
                 }
                 Ok(None) => {
                     log::debug!("No migration.");
@@ -204,6 +206,52 @@ impl Database {
         }
     }
 
+    /// Run periodic database maintenance.
+    ///
+    /// If `force=true` always run the full maintenance taks
+    pub fn run_maintenance(&self, force: bool) -> Result<()> {
+        let conn = self.conn.lock();
+        let conn = &*conn;
+
+        self.run_maintenance_vacuum(conn, force)?;
+        self.run_maintenance_optimize(conn)?;
+        self.run_maintenance_checkpoint(conn)?;
+
+        Ok(())
+    }
+
+    /// Run maintenance on the database (vacuum step)
+    ///
+    /// If `force_full: true` it _always_ runs a full `VACUUM`.
+    fn run_maintenance_vacuum(&self, conn: &rusqlite::Connection, force_full: bool) -> Result<()> {
+        let auto_vacuum_setting: u32 =
+            conn.query_row_and_then("PRAGMA auto_vacuum", [], |row| row.get(0))?;
+        if !force_full && auto_vacuum_setting == 2 {
+            // Ideally, we run an incremental vacuum to delete 2 pages
+            conn.execute("PRAGMA incremental_vacuum(2)", [])?;
+        } else {
+            // If auto_vacuum=incremental isn't set, configure it and run a full vacuum.
+            log::warn!(
+                "run_maintenance_vacuum: Need to run a full vacuum to set auto_vacuum=incremental"
+            );
+            conn.execute("PRAGMA auto_vacuum=incremental", [])?;
+            conn.execute("VACUUM", [])?;
+        }
+        Ok(())
+    }
+
+    /// Run maintenance on the database (optimize step)
+    fn run_maintenance_optimize(&self, conn: &rusqlite::Connection) -> Result<()> {
+        conn.execute("PRAGMA optimize", [])?;
+        Ok(())
+    }
+
+    /// Run maintenance on the database (checkpoint step)
+    fn run_maintenance_checkpoint(&self, conn: &rusqlite::Connection) -> Result<()> {
+        conn.query_row("PRAGMA wal_checkpoint(PASSIVE)", [], |_| Ok(()))?;
+        Ok(())
+    }
+
     /// Iterates with the provided transaction function
     /// over the requested data from the given storage.
     ///

glean-core/src/database/sqlite/connection.rs

@@ -7,7 +7,8 @@
 //! This module is inspired by, and borrows concepts from, the
 //! Application Services `sql-support` crate.
 
-use std::{fmt::Debug, num::NonZeroU32, path::Path, sync::Mutex};
+use std::sync::{Mutex, MutexGuard};
+use std::{fmt::Debug, num::NonZeroU32, path::Path};
 
 use rusqlite::{OpenFlags, Transaction, TransactionBehavior};
 
@@ -82,6 +83,11 @@ impl Connection {
         }
     }
 
+    /// Get ahold of the connection with no transaction opened.
+    pub fn lock<'a>(&'a self) -> MutexGuard<'a, rusqlite::Connection> {
+        self.conn.lock().unwrap()
+    }
+
     /// Accesses the database for reading.
     pub fn read<T, E>(&self, f: impl FnOnce(&Transaction<'_>) -> Result<T, E>) -> Result<T, E> {
         let mut conn = self.conn.lock().unwrap();

glean-core/tests/sqlite_migration.rs

@@ -1,5 +1,6 @@
 mod common;
 use std::fs;
+use std::os::unix::fs::MetadataExt;
 
 use crate::common::*;
 
@@ -10,6 +11,12 @@ use rkv::{Rkv, StoreOptions};
 use uuid::uuid;
 
 static RKV_DATABASE: &[u8] = include_bytes!("77ca0472-5124-4f6b-971d-4a2a928fb158.safe.bin");
+// This database is based on a submitted `metrics` ping from a real client.
+// Only metrics for that ping and a bare minimum of client info data is added,
+// a total of 61 metrics.
+// This is realistic enough to cause size differences in the non-vacuumed/vacuumed SQLite
+// databases after migration.
+static FILLED_RKV_DATABASE: &[u8] = include_bytes!("filled-rkv.data.safe.bin");
 
 fn clientid_metric() -> UuidMetric {
     UuidMetric::new(CommonMetricData {
@@ -218,3 +225,43 @@ fn migration_fails() {
     assert_eq!(None, metrics.failed_metrics.get_value(&glean, None));
     assert_eq!(None, metrics.migration_duration.get_value(&glean, None));
 }
+
+#[test]
+fn migration_checkpoints() {
+    let temp = tempfile::tempdir().unwrap();
+    let db_path = temp.path().join("db");
+    fs::create_dir_all(&db_path).unwrap();
+
+    let safe_bin = db_path.join("data.safe.bin");
+    // Reusing the same database file from above.
+    fs::write(safe_bin, FILLED_RKV_DATABASE).unwrap();
+    let exp_client_id = uuid!("3114d9df-9ae3-43a7-83b0-3540c3eba886");
+
+    let (glean, _temp) = new_glean(Some(temp));
+
+    let client_id = clientid_metric().get_value(&glean, None).unwrap();
+    assert_eq!(exp_client_id, client_id);
+
+    let metrics = MigrationMetrics::new();
+    assert_eq!(Some(61), metrics.migrated_metrics.get_value(&glean, None));
+    assert_eq!(Some(61), metrics.metrics_in_sqlite.get_value(&glean, None));
+
+    assert!(metrics.migration_duration.get_value(&glean, None).is_some());
+    assert_eq!(None, metrics.migration_error.get_value(&glean, None));
+
+    // Ensure we close the database connection.
+    drop(glean);
+
+    let db_file = db_path.join("glean.sqlite");
+    let db_file_size = fs::metadata(db_file).unwrap().size();
+
+    // This test is very vague, but it's hard to do better right now.
+    //
+    // Unvacuumed the database is _smaller_, around 20k bytes, because the migrated data is in the WAL file.
+    // Vacuumed & checkpointed the WAL transactions are merged into the database.
+    // As of writing that database is at least 8 pages big (8 * 4096 bytes = 32768 bytes).
+    // This might grow if we add more metrics.
+    // This might shrink if we remove metrics, in which case this test will break and needs adjustement.
+    let vacuumed_database_size = 32768;
+    assert!(db_file_size >= vacuumed_database_size);
+}