Bug 1813986 - Sqlite module connections should be able to register for vacuum on idle. r=asuth

Add a new vacuumOnIdle option to register a connection to the vacuum-participant
XPCOM category, so the VacuumManager can see it.
Add new openConnection options to set a database page size and incremental
vacuum mode.
Most of the tests for vacuum are covered in test_vacuum.js, so this just adds
basic tests to check vacuumOnIdle works in full and incremental mode.

Depends on D168298

Differential Revision: https://phabricator.services.mozilla.com/D170629

toolkit/modules/Sqlite.sys.mjs

@@ -255,6 +255,43 @@ XPCOMUtils.defineLazyGetter(lazy, "Barriers", () => {
   return Barriers;
 });
 
+const VACUUM_CATEGORY = "vacuum-participant";
+const VACUUM_CONTRACTID = "@sqlite.module.js/vacuum-participant;";
+var registeredVacuumParticipants = new Map();
+
+function registerVacuumParticipant(connectionData) {
+  let contractId = VACUUM_CONTRACTID + connectionData._identifier;
+  let factory = {
+    createInstance(iid) {
+      return connectionData.QueryInterface(iid);
+    },
+    QueryInterface: ChromeUtils.generateQI(["nsIFactory"]),
+  };
+  let cid = Services.uuid.generateUUID();
+  Components.manager
+    .QueryInterface(Ci.nsIComponentRegistrar)
+    .registerFactory(cid, contractId, contractId, factory);
+  Services.catMan.addCategoryEntry(
+    VACUUM_CATEGORY,
+    contractId,
+    contractId,
+    false,
+    false
+  );
+  registeredVacuumParticipants.set(contractId, { cid, factory });
+}
+
+function unregisterVacuumParticipant(connectionData) {
+  let contractId = VACUUM_CONTRACTID + connectionData._identifier;
+  let component = registeredVacuumParticipants.get(contractId);
+  if (component) {
+    Components.manager
+      .QueryInterface(Ci.nsIComponentRegistrar)
+      .unregisterFactory(component.cid, component.factory);
+    Services.catMan.deleteCategoryEntry(VACUUM_CATEGORY, contractId, false);
+  }
+}
+
 /**
  * Connection data with methods necessary for closing the connection.
  *
@@ -356,6 +393,33 @@ function ConnectionData(connection, identifier, options = {}) {
   this._timeoutPromise = null;
   // The last timestamp when we should consider using `this._timeoutPromise`.
   this._timeoutPromiseExpires = 0;
+
+  this._useIncrementalVacuum = !!options.incrementalVacuum;
+  if (this._useIncrementalVacuum) {
+    this._log.debug("Set auto_vacuum INCREMENTAL");
+    this.execute("PRAGMA auto_vacuum = 2").catch(ex => {
+      this._log.error("Setting auto_vacuum to INCREMENTAL failed.");
+      console.error(ex);
+    });
+  }
+
+  this._expectedPageSize = options.pageSize ?? 0;
+  if (this._expectedPageSize) {
+    this._log.debug("Set page_size to " + this._expectedPageSize);
+    this.execute("PRAGMA page_size = " + this._expectedPageSize).catch(ex => {
+      this._log.error(`Setting page_size to ${this._expectedPageSize} failed.`);
+      console.error(ex);
+    });
+  }
+
+  this._vacuumOnIdle = options.vacuumOnIdle;
+  if (this._vacuumOnIdle) {
+    this._log.debug("Register as vacuum participant");
+    this.QueryInterface = ChromeUtils.generateQI([
+      Ci.mozIStorageVacuumParticipant,
+    ]);
+    registerVacuumParticipant(this);
+  }
 }
 
 /**
@@ -371,6 +435,35 @@ function ConnectionData(connection, identifier, options = {}) {
 ConnectionData.byId = new Map();
 
 ConnectionData.prototype = Object.freeze({
+  get expectedDatabasePageSize() {
+    return this._expectedPageSize;
+  },
+
+  get useIncrementalVacuum() {
+    return this._useIncrementalVacuum;
+  },
+
+  /**
+   * This should only be used by the VacuumManager component.
+   * @see unsafeRawConnection for an official (but still unsafe) API.
+   */
+  get databaseConnection() {
+    if (this._vacuumOnIdle) {
+      return this._dbConn;
+    }
+    return null;
+  },
+
+  onBeginVacuum() {
+    let granted = !this.transactionInProgress;
+    this._log.debug("Begin Vacuum - " + granted ? "granted" : "denied");
+    return granted;
+  },
+
+  onEndVacuum(succeeded) {
+    this._log.debug("End Vacuum - " + succeeded ? "success" : "failure");
+  },
+
   /**
    * Run a task, ensuring that its execution will not be interrupted by shutdown.
    *
@@ -493,6 +586,11 @@ ConnectionData.prototype = Object.freeze({
     this._log.debug("Request to close connection.");
     this._clearIdleShrinkTimer();
 
+    if (this._vacuumOnIdle) {
+      this._log.debug("Unregister as vacuum participant");
+      unregisterVacuumParticipant(this);
+    }
+
     return this._barrier.wait().then(() => {
       if (!this._dbConn) {
         return undefined;
@@ -803,8 +901,9 @@ ConnectionData.prototype = Object.freeze({
 
   shrinkMemory() {
     this._log.debug("Shrinking memory usage.");
-    let onShrunk = this._clearIdleShrinkTimer.bind(this);
-    return this.execute("PRAGMA shrink_memory").then(onShrunk, onShrunk);
+    return this.execute("PRAGMA shrink_memory").finally(() => {
+      this._clearIdleShrinkTimer();
+    });
   },
 
   discardCachedStatements() {
@@ -1082,6 +1181,24 @@ ConnectionData.prototype = Object.freeze({
  *       USE WITH EXTREME CAUTION. This mode WILL produce incorrect results or
  *       return "false positive" corruption errors if other connections write
  *       to the DB at the same time.
+ *
+ *   vacuumOnIdle -- (bool) Whether to register this connection to be vacuumed
+ *       on idle by the VacuumManager component.
+ *       If you're vacuum-ing an incremental vacuum database, ensure to also
+ *       set incrementalVacuum to true, otherwise this will try to change it
+ *       to full vacuum mode.
+ *
+ *   incrementalVacuum -- (bool) if set to true auto_vacuum = INCREMENTAL will
+ *       be enabled for the database.
+ *       Changing auto vacuum of an already populated database requires a full
+ *       VACUUM. You can evaluate to enable vacuumOnIdle for that.
+ *
+ *   pageSize -- (integer) This allows to set a custom page size for the
+ *       database. It is usually not necessary to set it, since the default
+ *       value should be good for most consumers.
+ *       Changing the page size of an already populated database requires a full
+ *       VACUUM. You can evaluate to enable vacuumOnIdle for that.
+ *
  *   testDelayedOpenPromise -- (promise) Used by tests to delay the open
  *       callback handling and execute code between asyncOpen and its callback.
  *
@@ -1163,6 +1280,33 @@ function openConnection(options) {
     openedOptions.defaultTransactionType = defaultTransactionType;
   }
 
+  if ("vacuumOnIdle" in options) {
+    if (typeof options.vacuumOnIdle != "boolean") {
+      throw new Error("Invalid vacuumOnIdle: " + options.vacuumOnIdle);
+    }
+    openedOptions.vacuumOnIdle = options.vacuumOnIdle;
+  }
+
+  if ("incrementalVacuum" in options) {
+    if (typeof options.incrementalVacuum != "boolean") {
+      throw new Error(
+        "Invalid incrementalVacuum: " + options.incrementalVacuum
+      );
+    }
+    openedOptions.incrementalVacuum = options.incrementalVacuum;
+  }
+
+  if ("pageSize" in options) {
+    if (
+      ![512, 1024, 2048, 4096, 8192, 16384, 32768, 65536].includes(
+        options.pageSize
+      )
+    ) {
+      throw new Error("Invalid pageSize: " + options.pageSize);
+    }
+    openedOptions.pageSize = options.pageSize;
+  }
+
   let identifier = getIdentifierByFileName(PathUtils.filename(path));
 
   log.debug("Opening database: " + path + " (" + identifier + ")");

toolkit/modules/tests/xpcshell/test_sqlite.js

@@ -1381,3 +1381,19 @@ add_task(async function test_interrupt() {
     "Sqlite.interrupt() should throw on a closed connection"
   );
 });
+
+add_task(async function test_pageSize() {
+  // Testing the possibility to set the page size on database creation.
+  await Assert.rejects(
+    getDummyDatabase("pagesize", { pageSize: 1234 }),
+    /Invalid pageSize/,
+    "Check invalid pageSize value"
+  );
+  let c = await getDummyDatabase("pagesize", { pageSize: 8192 });
+  Assert.equal(
+    (await c.execute("PRAGMA page_size"))[0].getResultByIndex(0),
+    8192,
+    "Check page size was set"
+  );
+  await c.close();
+});

toolkit/modules/tests/xpcshell/test_sqlite_autoVacuum.js

@@ -0,0 +1,96 @@
+"use strict";
+
+const { Sqlite } = ChromeUtils.importESModule(
+  "resource://gre/modules/Sqlite.sys.mjs"
+);
+const { TestUtils } = ChromeUtils.importESModule(
+  "resource://testing-common/TestUtils.sys.mjs"
+);
+
+/**
+ * Sends a fake idle-daily notification to the VACUUM Manager.
+ */
+function synthesize_idle_daily() {
+  Cc["@mozilla.org/storage/vacuum;1"]
+    .getService(Ci.nsIObserver)
+    .observe(null, "idle-daily", null);
+}
+
+function unregister_vacuum_participants() {
+  // First unregister other participants.
+  for (let { data: entry } of Services.catMan.enumerateCategory(
+    "vacuum-participant"
+  )) {
+    Services.catMan.deleteCategoryEntry("vacuum-participant", entry, false);
+  }
+}
+
+function reset_vacuum_date(dbname) {
+  let date = parseInt(Date.now() / 1000 - 31 * 86400);
+  // Set last VACUUM to a date in the past.
+  Services.prefs.setIntPref(`storage.vacuum.last.${dbname}`, date);
+  return date;
+}
+
+function get_vacuum_date(dbname) {
+  return Services.prefs.getIntPref(`storage.vacuum.last.${dbname}`, 0);
+}
+
+async function get_freelist_count(conn) {
+  return (await conn.execute("PRAGMA freelist_count"))[0].getResultByIndex(0);
+}
+
+async function get_auto_vacuum(conn) {
+  return (await conn.execute("PRAGMA auto_vacuum"))[0].getResultByIndex(0);
+}
+
+async function test_vacuum(options = {}) {
+  unregister_vacuum_participants();
+  const dbName = "testVacuum.sqlite";
+  const dbFile = PathUtils.join(PathUtils.profileDir, dbName);
+  let lastVacuumDate = reset_vacuum_date(dbName);
+  let conn = await Sqlite.openConnection(
+    Object.assign(
+      {
+        path: dbFile,
+        vacuumOnIdle: true,
+      },
+      options
+    )
+  );
+  // Ensure the category manager is up-to-date.
+  await TestUtils.waitForTick();
+
+  Assert.equal(
+    await get_auto_vacuum(conn),
+    options.incrementalVacuum ? 2 : 0,
+    "Check auto_vacuum"
+  );
+
+  // Generate some freelist page.
+  await conn.execute("CREATE TABLE test (id INTEGER)");
+  await conn.execute("DROP TABLE test");
+  Assert.greater(await get_freelist_count(conn), 0, "Check freelist_count");
+
+  let promiseVacuumEnd = TestUtils.topicObserved(
+    "vacuum-end",
+    (_, d) => d == dbName
+  );
+  synthesize_idle_daily();
+  info("Await vacuum end");
+  await promiseVacuumEnd;
+
+  Assert.greater(get_vacuum_date(dbName), lastVacuumDate);
+
+  Assert.equal(await get_freelist_count(conn), 0, "Check freelist_count");
+
+  await conn.close();
+  await IOUtils.remove(dbFile);
+}
+
+add_task(async function test_vacuumOnIdle() {
+  info("Test full vacuum");
+  await test_vacuum();
+  info("Test incremental vacuum");
+  await test_vacuum({ incrementalVacuum: true });
+});

toolkit/modules/tests/xpcshell/xpcshell.ini

@@ -60,6 +60,8 @@ run-sequentially = very high failure rate in parallel
 [test_Services.js]
 [test_sqlite.js]
 skip-if = toolkit == 'android'
+[test_sqlite_autoVacuum.js]
+skip-if = toolkit == 'android'
 [test_sqlite_shutdown.js]
 [test_timer.js]
 [test_UpdateUtils_url.js]