Optimize Indexes (#1239)

Analogous to dbos-inc/dbos-transact-py#656

Author: kraftp

schemas/system_db_schema.ts

@@ -32,6 +32,7 @@ export interface workflow_status {
   parent_workflow_id?: string;
   serialization: SysDBSerializationFormat | null;
   delay_until_epoch_ms?: number | null;
+  rate_limited?: boolean;
 }
 
 export interface notifications {

src/sysdb_migrations/internal/migrations.ts

@@ -6,6 +6,7 @@ export function allMigrations(
 ): ReadonlyArray<DBMigration> {
   const useListenNotify = opts?.useListenNotify ?? true;
   const isCockroach = opts?.isCockroach ?? false;
+  const c = isCockroach ? '' : 'CONCURRENTLY';
   return [
     {
       name: '20240123182943_schema',
@@ -432,5 +433,85 @@ export function allMigrations(
         )`,
       ],
     },
+    // Migrations below replace broad indexes on workflow_status with partial
+    // indexes targeted at individual query patterns (recovery, troubleshooting,
+    // dequeue, rate-limit count). Each index DDL runs CONCURRENTLY on Postgres
+    // for safe online deployment; CockroachDB ignores the keyword.
+    {
+      online: true,
+      pg: [`DROP INDEX ${c} IF EXISTS "${schemaName}"."idx_workflow_status_forked_from"`],
+    },
+    {
+      online: true,
+      pg: [
+        `CREATE INDEX ${c} IF NOT EXISTS "idx_workflow_status_forked_from" ON "${schemaName}"."workflow_status" ("forked_from") WHERE "forked_from" IS NOT NULL`,
+      ],
+    },
+    {
+      online: true,
+      pg: [`DROP INDEX ${c} IF EXISTS "${schemaName}"."idx_workflow_status_parent_workflow_id"`],
+    },
+    {
+      online: true,
+      pg: [
+        `CREATE INDEX ${c} IF NOT EXISTS "idx_workflow_status_parent_workflow_id" ON "${schemaName}"."workflow_status" ("parent_workflow_id") WHERE "parent_workflow_id" IS NOT NULL`,
+      ],
+    },
+    {
+      online: true,
+      pg: [`DROP INDEX ${c} IF EXISTS "${schemaName}"."workflow_status_executor_id_index"`],
+    },
+    {
+      online: true,
+      pg: [
+        `CREATE UNIQUE INDEX ${c} IF NOT EXISTS "uq_workflow_status_dedup_id" ON "${schemaName}"."workflow_status" ("queue_name", "deduplication_id") WHERE "deduplication_id" IS NOT NULL`,
+      ],
+    },
+    {
+      // CockroachDB stores `UNIQUE (...)` constraints as unique indexes, so
+      // they must be dropped with DROP INDEX rather than ALTER TABLE DROP CONSTRAINT.
+      pg: isCockroach
+        ? [`DROP INDEX IF EXISTS "${schemaName}"."uq_workflow_status_queue_name_dedup_id" CASCADE`]
+        : [
+            `ALTER TABLE "${schemaName}"."workflow_status" DROP CONSTRAINT IF EXISTS "uq_workflow_status_queue_name_dedup_id"`,
+          ],
+    },
+    {
+      online: true,
+      pg: [
+        `CREATE INDEX ${c} IF NOT EXISTS "idx_workflow_status_pending" ON "${schemaName}"."workflow_status" ("created_at") WHERE "status" = 'PENDING'`,
+      ],
+    },
+    {
+      online: true,
+      pg: [
+        `CREATE INDEX ${c} IF NOT EXISTS "idx_workflow_status_failed" ON "${schemaName}"."workflow_status" ("status", "created_at") WHERE "status" IN ('ERROR', 'CANCELLED', 'MAX_RECOVERY_ATTEMPTS_EXCEEDED')`,
+      ],
+    },
+    {
+      online: true,
+      pg: [`DROP INDEX ${c} IF EXISTS "${schemaName}"."workflow_status_status_index"`],
+    },
+    {
+      online: true,
+      pg: [
+        `CREATE INDEX ${c} IF NOT EXISTS "idx_workflow_status_in_flight" ON "${schemaName}"."workflow_status" ("queue_name", "status", "priority", "created_at") WHERE "status" IN ('ENQUEUED', 'PENDING')`,
+      ],
+    },
+    {
+      pg: [
+        `ALTER TABLE "${schemaName}"."workflow_status" ADD COLUMN IF NOT EXISTS "rate_limited" BOOLEAN NOT NULL DEFAULT FALSE`,
+      ],
+    },
+    {
+      online: true,
+      pg: [
+        `CREATE INDEX ${c} IF NOT EXISTS "idx_workflow_status_rate_limited" ON "${schemaName}"."workflow_status" ("queue_name", "started_at_epoch_ms") WHERE "rate_limited" = TRUE`,
+      ],
+    },
+    {
+      online: true,
+      pg: [`DROP INDEX ${c} IF EXISTS "${schemaName}"."idx_workflow_status_queue_status_started"`],
+    },
   ];
 }

src/sysdb_migrations/migration_runner.ts

@@ -4,6 +4,14 @@ export type DBMigration = {
   name?: string;
   pg?: ReadonlyArray<string>;
   sqlite3?: ReadonlyArray<string>;
+  /**
+   * If true, the migration is executed without wrapping it in any helper that
+   * suppresses errors, and its statements must be safe to run outside a
+   * transaction (e.g. `CREATE INDEX CONCURRENTLY`). The runner also cleans up
+   * indexes left INVALID by a previously interrupted CONCURRENTLY build before
+   * running the migration.
+   */
+  online?: boolean;
 };
 
 /** Get the current DB version, or 0 if table is missing/empty. */
@@ -28,6 +36,7 @@ export async function getCurrentSysDBVersion(client: ClientBase, schemaName: str
 export type PgMigratorOptions = {
   ignoreErrorCodes?: ReadonlySet<string>;
   onWarn?: (msg: string, err?: unknown) => void;
+  isCockroach?: boolean;
 };
 
 const DEFAULT_IGNORABLE_CODES = new Set<string>([
@@ -73,11 +82,49 @@ async function runStatementsIgnoring(
   }
 }
 
+/**
+ * Drop indexes left INVALID by a prior interrupted `CREATE INDEX CONCURRENTLY`.
+ * Postgres marks such indexes invalid; they continue to consume write overhead
+ * without serving reads, so the next online migration cannot succeed by simply
+ * re-running `IF NOT EXISTS` (the name is taken).
+ */
+async function cleanupInvalidIndexes(client: ClientBase, schemaName: string, warn: (m: string) => void): Promise<void> {
+  const res = await client.query<{ relname: string }>(
+    `SELECT i.relname
+       FROM pg_index ix
+       JOIN pg_class i ON i.oid = ix.indexrelid
+       JOIN pg_class t ON t.oid = ix.indrelid
+       JOIN pg_namespace n ON n.oid = t.relnamespace
+      WHERE NOT ix.indisvalid AND n.nspname = $1`,
+    [schemaName],
+  );
+  for (const row of res.rows) {
+    warn(`Dropping invalid index "${schemaName}"."${row.relname}" left by a prior failed migration`);
+    await client.query(`DROP INDEX CONCURRENTLY IF EXISTS "${schemaName}"."${row.relname}"`);
+  }
+}
+
+async function bumpMigrationVersion(client: ClientBase, schemaName: string, version: number): Promise<void> {
+  // The earliest migrations create the schema and the dbos_migrations table itself,
+  // so the table is not available to update on the very first iterations of a fresh DB.
+  const reg = await client.query<{ table_name: string | null }>(
+    "SELECT table_name FROM information_schema.tables WHERE table_schema = $1 AND table_name = 'dbos_migrations'",
+    [schemaName],
+  );
+  if (!reg.rows[0]?.table_name) return;
+
+  const updateRes = await client.query(`UPDATE "${schemaName}"."dbos_migrations" SET "version" = $1`, [version]);
+  if (updateRes.rowCount === 0) {
+    await client.query(`INSERT INTO "${schemaName}"."dbos_migrations" ("version") VALUES ($1)`, [version]);
+  }
+}
+
 /**
  * Apply all migrations greater than the current DB version.
  * - Reads current version (0 if table missing/empty)
  * - Applies migrations in order
- * - Updates dbos_migrations.version at the end
+ * - After each migration, persists `dbos_migrations.version` so partial
+ *   progress is recorded
  * - Warns if current version > max known (likely newer software concurrently)
  */
 export async function runSysMigrationsPg(
@@ -92,7 +139,7 @@ export async function runSysMigrationsPg(
   skippedCount: number;
   notice?: string;
 }> {
-  const { ignoreErrorCodes = DEFAULT_IGNORABLE_CODES, onWarn = (m) => console.info(m) } = opts;
+  const { ignoreErrorCodes = DEFAULT_IGNORABLE_CODES, onWarn = (m) => console.info(m), isCockroach = false } = opts;
 
   const current = await getCurrentSysDBVersion(client, schemaName);
   const maxKnown = allMigrations.length;
@@ -131,26 +178,30 @@ export async function runSysMigrationsPg(
     const stmts = m.pg ?? [];
     if (stmts.length === 0) {
       onWarn(`Migration "${m.name}" has no Postgres statements; skipping.`);
+      await bumpMigrationVersion(client, schemaName, v);
       skipped++;
       lastAppliedVersion = v;
       continue;
     }
 
-    await runStatementsIgnoring(client, stmts, ignoreErrorCodes, (msg, e) =>
-      onWarn(`${msg}${e ? `\n  → ${String((e as { message?: string }).message ?? '')}` : ''}`),
-    );
+    // Run migrations in autocommit
+    if (m.online && !isCockroach) {
+      // Drop any indexes left INVALID by a prior interrupted run before
+      // attempting the next CONCURRENTLY statement.
+      await cleanupInvalidIndexes(client, schemaName, (msg) => onWarn(msg));
+      for (const s of stmts) {
+        await client.query(s, []);
+      }
+    } else {
+      await runStatementsIgnoring(client, stmts, ignoreErrorCodes, (msg, e) =>
+        onWarn(`${msg}${e ? `\n  → ${String((e as { message?: string }).message ?? '')}` : ''}`),
+      );
+    }
+    await bumpMigrationVersion(client, schemaName, v);
     applied++;
     lastAppliedVersion = v;
   }
 
-  // Update version at the end (insert or update)
-  const updateRes = await client.query(`UPDATE "${schemaName}"."dbos_migrations" SET "version" = $1`, [
-    lastAppliedVersion,
-  ]);
-  if (updateRes.rowCount === 0) {
-    await client.query(`INSERT into "${schemaName}"."dbos_migrations" ("version") values ($1)`, [lastAppliedVersion]);
-  }
-
   return {
     fromVersion: current,
     toVersion: lastAppliedVersion,

src/system_database.ts

@@ -294,6 +294,7 @@ export async function ensureSystemDatabase(
     const isCockroach = /cockroachdb/i.test(versionRes.rows[0]?.version ?? '');
     await runSysMigrationsPg(client, allMigrations(schemaName, { useListenNotify, isCockroach }), schemaName, {
       onWarn: (e: string) => logger.info(e),
+      isCockroach,
     });
   } finally {
     try {
@@ -2522,6 +2523,7 @@ export class SystemDatabase {
         const countResult = await client.query<{ count: string }>(
           `SELECT COUNT(*) FROM "${this.schemaName}".workflow_status
            WHERE queue_name = $1
+             AND rate_limited = TRUE
              AND status NOT IN ($2, $3)
              AND started_at_epoch_ms > $4
              ${partitionFilter.replace('$PARTITION', '$5')}`,
@@ -2574,7 +2576,6 @@ export class SystemDatabase {
       // Retrieve the first max_tasks workflows in the queue.
       // Only retrieve workflows of the local version (or without version set)
       const lockMode = queue.concurrency ? 'FOR UPDATE NOWAIT' : 'FOR UPDATE SKIP LOCKED';
-      const orderClause = queue.priorityEnabled ? 'ORDER BY priority ASC, created_at ASC' : 'ORDER BY created_at ASC';
       const limitClause = maxTasks !== Infinity ? `LIMIT ${maxTasks}` : '';
 
       const selectParams = [StatusString.ENQUEUED, queue.name, appVersion, ...partitionParams];
@@ -2585,7 +2586,7 @@ export class SystemDatabase {
           AND queue_name = $2
           AND (application_version IS NULL OR application_version = $3)
           ${partitionFilter.replace('$PARTITION', '$4')}
-        ${orderClause}
+        ORDER BY priority ASC, created_at ASC
         ${limitClause}
         ${lockMode}
       `;
@@ -2610,13 +2611,22 @@ export class SystemDatabase {
                executor_id = $2,
                application_version = $3,
                started_at_epoch_ms = $4,
+               rate_limited = $5,
                workflow_deadline_epoch_ms = CASE
                  WHEN workflow_timeout_ms IS NOT NULL AND workflow_deadline_epoch_ms IS NULL
                  THEN (EXTRACT(epoch FROM now()) * 1000)::bigint + workflow_timeout_ms
                  ELSE workflow_deadline_epoch_ms
                END
-           WHERE workflow_uuid = $5 AND status = $6`,
-          [StatusString.PENDING, executorID, appVersion, startTimeMs, id, StatusString.ENQUEUED],
+           WHERE workflow_uuid = $6 AND status = $7`,
+          [
+            StatusString.PENDING,
+            executorID,
+            appVersion,
+            startTimeMs,
+            queue.rateLimit !== undefined,
+            id,
+            StatusString.ENQUEUED,
+          ],
         );
         if ((updateRes.rowCount ?? 0) > 0) {
           claimedIDs.push(id);