fix: Thread index database migrations are now atomic (#2280)

* fix: Thread index database migrations are now atomic

A single process will now migrate the database when needed. Migrations
will no longer be run every time the RPC service starts.

Author: dustinbyrne

packages/cli/package.json

@@ -49,6 +49,7 @@
     "@types/lunr": "^2.3.7",
     "@types/moo": "^0.5.5",
     "@types/node": "^16",
+    "@types/proper-lockfile": "^4.1.4",
     "@types/semver": "^7.3.10",
     "@types/sinon": "^10.0.2",
     "@types/tmp": "^0.2.3",
@@ -141,6 +142,7 @@
     "pkg": "^5.8.1",
     "port-pid": "^0.0.7",
     "pretty-bytes": "^5.6.0",
+    "proper-lockfile": "^4.1.2",
     "ps-node": "^0.1.6",
     "read-pkg-up": "^7.0.1",
     "reflect-metadata": "^0.2.2",

packages/cli/src/cmds/index/index.ts

@@ -136,7 +136,7 @@ export const handler = async (argv) => {
       const buildRemoteNavie = () => new NopNavie();
 
       const navieProvider = useLocalNavie() ? buildLocalNavie : buildRemoteNavie;
-      ThreadIndexService.useDefault();
+      await ThreadIndexService.useDefault();
       NavieService.bindNavieProvider(navieProvider);
 
       await configureRpcDirectories([process.cwd()]);

packages/cli/src/cmds/index/rpc.ts

@@ -116,7 +116,7 @@ export const handler = async (argv: HandlerArguments) => {
 
   const navie = buildNavieProvider(argv);
 
-  ThreadIndexService.useDefault();
+  await ThreadIndexService.useDefault();
   NavieService.bindNavieProvider(navie);
 
   let codeEditor: string | undefined = argv.codeEditor;

packages/cli/src/rpc/navie/services/threadIndexService.ts

@@ -3,12 +3,18 @@ import { homedir } from 'node:os';
 import { dirname, join } from 'node:path';
 import configuration from '../../configuration';
 import { container, inject, injectable, singleton } from 'tsyringe';
-import { mkdirSync } from 'node:fs';
+import { mkdir } from 'node:fs/promises';
+import { lock } from 'proper-lockfile';
+import { isNativeError } from 'node:util/types';
 
-const INITIALIZE_SQL = `
+const SCHEMA_VERSION = 1;
+
+const INITIALIZE_SESSION_SQL = `
+PRAGMA busy_timeout = 3000;
 PRAGMA foreign_keys = ON;
-PRAGMA journal_mode = WAL;
+`;
 
+const INITIALIZE_DB_SQL = `
 CREATE TABLE IF NOT EXISTS threads (
   id INTEGER PRIMARY KEY AUTOINCREMENT,
   uuid TEXT NOT NULL UNIQUE,
@@ -31,6 +37,8 @@ CREATE TABLE IF NOT EXISTS project_directories (
 );
 
 CREATE INDEX IF NOT EXISTS idx_thread_id ON project_directories (thread_id);
+
+PRAGMA user_version = ${SCHEMA_VERSION};
 `;
 
 const QUERY_INSERT_THREAD_SQL = `INSERT INTO threads (uuid, path, title) VALUES (?, ?, ?)
@@ -63,20 +71,70 @@ export interface ThreadIndexItem {
 @singleton()
 @injectable()
 export class ThreadIndexService {
+  static readonly MIGRATION_RETRIES = 5;
   static readonly DEFAULT_DATABASE_PATH = join(homedir(), '.appmap', 'navie', 'thread-index.db');
   static readonly DATABASE = 'ThreadIndexDatabase';
 
   constructor(@inject(ThreadIndexService.DATABASE) private readonly db: sqlite3.Database) {
-    this.db.exec(INITIALIZE_SQL);
+    this.db.exec(INITIALIZE_SESSION_SQL);
+  }
+
+  private shouldMigrate(): boolean {
+    const currentVersion = this.db.get('PRAGMA user_version')?.user_version;
+    return currentVersion !== SCHEMA_VERSION;
+  }
+
+  /**
+   * Migrates the database schema to the latest version. This should be called upon application
+   * initialization to ensure that the database is up to date.
+   */
+  public async migrate(databaseFilePath = ThreadIndexService.DEFAULT_DATABASE_PATH) {
+    if (!this.shouldMigrate()) return;
+
+    let lockRelease: (() => Promise<void>) | undefined;
+    try {
+      await mkdir(dirname(databaseFilePath), { recursive: true });
+
+      lockRelease = await lock(databaseFilePath, {
+        retries: ThreadIndexService.MIGRATION_RETRIES,
+        stale: 10000,
+        lockfilePath: `${databaseFilePath}.migration.lock`,
+      });
+
+      if (!this.shouldMigrate()) {
+        console.info('ThreadIndexService: migration completed by another process');
+        return;
+      }
+
+      console.info(`ThreadIndexService: migrating schema to v${SCHEMA_VERSION}`);
+      try {
+        this.db.exec('PRAGMA journal_mode = WAL');
+        this.db.exec('BEGIN TRANSACTION');
+        this.db.exec(INITIALIZE_DB_SQL);
+        this.db.exec('COMMIT');
+        console.info(`ThreadIndexService: successfully migrated schema to v${SCHEMA_VERSION}`);
+        return;
+      } catch (e) {
+        this.db.exec('ROLLBACK');
+        throw e;
+      }
+    } finally {
+      if (lockRelease) {
+        await lockRelease();
+      }
+    }
   }
 
   /**
    * Binds the database to a sqlite3 instance on disk at the default database path
    */
-  static useDefault() {
-    mkdirSync(dirname(this.DEFAULT_DATABASE_PATH), { recursive: true });
+  static async useDefault() {
+    await mkdir(dirname(this.DEFAULT_DATABASE_PATH), { recursive: true });
+
     const db = new sqlite3.Database(this.DEFAULT_DATABASE_PATH);
     container.registerInstance(this.DATABASE, db);
+
+    await container.resolve(ThreadIndexService).migrate();
   }
 
   /**

packages/cli/tests/unit/rpc/navie/services/threadIndexService.spec.ts

@@ -11,19 +11,31 @@ import {
 } from '../../../../../src/rpc/navie/services/threadIndexService';
 import sqlite3 from 'node-sqlite3-wasm';
 import configuration from '../../../../../src/rpc/configuration';
+import { mkdtemp, rm, writeFile } from 'fs/promises';
+import { join } from 'path';
 
 describe('ThreadIndexService', () => {
+  let tmpDir: string;
   let threadIndexService: ThreadIndexService;
   let db: sqlite3.Database;
   const threadId = '00000000-0000-0000-0000-000000000000';
 
-  beforeEach(() => {
+  beforeEach(async () => {
     container.reset();
     db = new sqlite3.Database(':memory:');
     container.registerInstance(ThreadIndexService.DATABASE, db);
     threadIndexService = container.resolve(ThreadIndexService);
+
+    // Create a fake database file on disk for file locking
+    tmpDir = await mkdtemp('thread-index-test-');
+    const mockDbFile = join(tmpDir, 'mock.db');
+    await writeFile(mockDbFile, '');
+
+    await threadIndexService.migrate(mockDbFile);
   });
 
+  afterEach(() => rm(tmpDir, { recursive: true, force: true }));
+
   describe('indexThread', () => {
     it('indexes a thread', () => {
       const path = 'example-thread-history.jsonl';

yarn.lock

@@ -171,6 +171,7 @@ __metadata:
     "@types/lunr": ^2.3.7
     "@types/moo": ^0.5.5
     "@types/node": ^16
+    "@types/proper-lockfile": ^4.1.4
     "@types/semver": ^7.3.10
     "@types/sinon": ^10.0.2
     "@types/tmp": ^0.2.3
@@ -232,6 +233,7 @@ __metadata:
     port-pid: ^0.0.7
     prettier: ^2.7.1
     pretty-bytes: ^5.6.0
+    proper-lockfile: ^4.1.2
     ps-node: ^0.1.6
     puppeteer: ^19.7.5
     read-pkg-up: ^7.0.1
@@ -11407,6 +11409,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@types/proper-lockfile@npm:^4.1.4":
+  version: 4.1.4
+  resolution: "@types/proper-lockfile@npm:4.1.4"
+  dependencies:
+    "@types/retry": "*"
+  checksum: b0d1b8e84a563b2c5f869f7ff7542b1d83dec03d1c9d980847cbb189865f44b4a854673cdde59767e41bcb8c31932e613ac43822d358a6f8eede6b79ccfceb1d
+  languageName: node
+  linkType: hard
+
 "@types/q@npm:^1.5.1":
   version: 1.5.5
   resolution: "@types/q@npm:1.5.5"
@@ -11448,6 +11459,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@types/retry@npm:*":
+  version: 0.12.5
+  resolution: "@types/retry@npm:0.12.5"
+  checksum: 3fb6bf91835ca0eb2987567d6977585235a7567f8aeb38b34a8bb7bbee57ac050ed6f04b9998cda29701b8c893f5dfe315869bc54ac17e536c9235637fe351a2
+  languageName: node
+  linkType: hard
+
 "@types/retry@npm:0.12.0":
   version: 0.12.0
   resolution: "@types/retry@npm:0.12.0"
@@ -36013,6 +36031,17 @@ __metadata:
   languageName: node
   linkType: hard
 
+"proper-lockfile@npm:^4.1.2":
+  version: 4.1.2
+  resolution: "proper-lockfile@npm:4.1.2"
+  dependencies:
+    graceful-fs: ^4.2.4
+    retry: ^0.12.0
+    signal-exit: ^3.0.2
+  checksum: 00078ee6a61c216a56a6140c7d2a98c6c733b3678503002dc073ab8beca5d50ca271de4c85fca13b9b8ee2ff546c36674d1850509b84a04a5d0363bcb8638939
+  languageName: node
+  linkType: hard
+
 "proto-list@npm:~1.2.1":
   version: 1.2.4
   resolution: "proto-list@npm:1.2.4"