fix(server): hold migration advisory lock across the whole migration (PERF-027)

packages/server/src/__tests__/bootstrap-migration-lock.test.ts

@@ -1,19 +1,24 @@
 import postgres from "postgres";
-import { afterEach, describe, expect, it } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import { sslOptions } from "../db/connection.js";
 import { runMigrations } from "../db/migrate.js";
 
+const LOCK_KEY_SQL = "hashtext('drizzle_migrations')";
+
 /**
- * Pins the `hashtext('drizzle_migrations')` key used by
- * `preflightMigrationLock`. If a future drizzle bump changes the lock key,
- * this test stays green only as long as the preflight constant matches the
- * one we hold here — meaning the preflight and the holder both speak to the
- * same advisory-lock slot.
+ * Pins the `hashtext('drizzle_migrations')` advisory lock that
+ * `runMigrations` acquires and HOLDS for the entire migration (see
+ * `packages/server/src/db/migrate.ts` `MIGRATION_LOCK_KEY_SQL`):
  *
- * See `packages/server/src/db/migrate.ts` `MIGRATION_LOCK_KEY_SQL` comment
- * for the verification path through drizzle-orm.
+ *   1. contention: an external holder makes `runMigrations` fail with a
+ *      clear error after the timeout instead of running DDL in parallel;
+ *   2. waiting: once the holder releases, a waiting replica acquires the
+ *      lock and completes normally (journal no-op on an up-to-date DB);
+ *   3. full-span locking: while drizzle `migrate()` is executing, no other
+ *      session can grab the key — the startup-migration serialization that
+ *      PERF-027 requires across replicas.
  */
-describe("preflightMigrationLock (T10)", () => {
+describe("runMigrations advisory lock (T10 / PERF-027)", () => {
   const databaseUrl = process.env.DATABASE_URL;
 
   afterEach(() => {
@@ -30,23 +35,109 @@ describe("preflightMigrationLock (T10)", () => {
     // contender so it returns immediately.
     const holder = postgres(url, { max: 1, ...sslOptions(url) });
     try {
-      await holder`SELECT pg_advisory_lock(hashtext('drizzle_migrations'))`;
+      await holder.unsafe(`SELECT pg_advisory_lock(${LOCK_KEY_SQL})`);
 
-      // Use a tight 2s preflight timeout so the contention path completes
-      // fast. The 30s production default is the right answer for boot, not
-      // for a unit test that just needs to assert the error path.
+      // Use a tight 2s lock timeout so the contention path completes fast.
+      // The 15s production default is the right answer for boot, not for a
+      // unit test that just needs to assert the error path.
       await expect(runMigrations(url, { lockTimeoutMs: 2_000 })).rejects.toThrow(/migration lock contention/);
     } finally {
       // Release the lock so subsequent test runs (and other test files in
       // the same worker process) aren't blocked.
-      await holder`SELECT pg_advisory_unlock(hashtext('drizzle_migrations'))`;
+      await holder.unsafe(`SELECT pg_advisory_unlock(${LOCK_KEY_SQL})`);
       await holder.end();
     }
   });
 
+  it("waits for a holder to release, then completes the migration", async () => {
+    expect(databaseUrl, "DATABASE_URL must be set by global setup").toBeTruthy();
+    const url = databaseUrl ?? "";
+
+    // Simulate the multi-replica rollout: another "replica" owns the lock
+    // when we start. runMigrations must poll (1s interval) rather than
+    // fail immediately, and proceed once the owner releases.
+    const holder = postgres(url, { max: 1, ...sslOptions(url) });
+    let released = false;
+    try {
+      await holder.unsafe(`SELECT pg_advisory_lock(${LOCK_KEY_SQL})`);
+
+      const waiter = runMigrations(url, { lockTimeoutMs: 10_000 });
+
+      // Keep the lock across at least one poll cycle, then release. 1.5s
+      // straddles the first retry so the waiter observes real contention.
+      await new Promise((r) => setTimeout(r, 1_500));
+      await holder.unsafe(`SELECT pg_advisory_unlock(${LOCK_KEY_SQL})`);
+      released = true;
+
+      // The DB is already migrated (template clone), so the waiter's
+      // migrate() is a journal no-op — success here proves the waiting
+      // path, not re-application.
+      await expect(waiter).resolves.toBeGreaterThan(0);
+    } finally {
+      if (!released) {
+        await holder.unsafe(`SELECT pg_advisory_unlock(${LOCK_KEY_SQL})`);
+      }
+      await holder.end();
+    }
+  });
+
+  it("holds the advisory lock while drizzle migrate() is executing", async () => {
+    expect(databaseUrl, "DATABASE_URL must be set by global setup").toBeTruthy();
+    const url = databaseUrl ?? "";
+
+    // Swap drizzle's migrate() for a probe that checks — from a SECOND
+    // session — whether the lock is grabbable mid-migration. If the lock
+    // were released before migrate() (the old preflight-then-unlock bug),
+    // the probe would acquire it and this test would fail.
+    let lockGrabbableDuringMigrate: boolean | undefined;
+    vi.resetModules();
+    vi.doMock("drizzle-orm/postgres-js/migrator", () => ({
+      migrate: vi.fn(async () => {
+        const probe = postgres(url, { max: 1, ...sslOptions(url) });
+        try {
+          const rows = (await probe.unsafe(`SELECT pg_try_advisory_lock(${LOCK_KEY_SQL}) AS acquired`)) as Array<{
+            acquired: boolean;
+          }>;
+          lockGrabbableDuringMigrate = rows[0]?.acquired;
+          if (lockGrabbableDuringMigrate) {
+            // Never expected — but don't strand the lock if the invariant
+            // breaks, or later cases in this worker would time out.
+            await probe.unsafe(`SELECT pg_advisory_unlock(${LOCK_KEY_SQL})`);
+          }
+        } finally {
+          await probe.end();
+        }
+      }),
+    }));
+
+    try {
+      const { runMigrations: runMigrationsWithMockedMigrator } = await import("../db/migrate.js");
+      const tableCount = await runMigrationsWithMockedMigrator(url);
+
+      expect(lockGrabbableDuringMigrate).toBe(false);
+      expect(tableCount).toBeGreaterThan(0);
+
+      // And the lock must be released again once runMigrations returns
+      // (session closed), so the next replica can proceed.
+      const after = postgres(url, { max: 1, ...sslOptions(url) });
+      try {
+        const rows = (await after.unsafe(`SELECT pg_try_advisory_lock(${LOCK_KEY_SQL}) AS acquired`)) as Array<{
+          acquired: boolean;
+        }>;
+        expect(rows[0]?.acquired).toBe(true);
+        await after.unsafe(`SELECT pg_advisory_unlock(${LOCK_KEY_SQL})`);
+      } finally {
+        await after.end();
+      }
+    } finally {
+      vi.doUnmock("drizzle-orm/postgres-js/migrator");
+      vi.resetModules();
+    }
+  });
+
   it("succeeds when the advisory lock is free", async () => {
     expect(databaseUrl, "DATABASE_URL must be set by global setup").toBeTruthy();
-    // Sanity case: with no holder, preflight returns fast and migrate runs.
+    // Sanity case: with no holder, the lock is acquired fast and migrate runs.
     const tableCount = await runMigrations(databaseUrl ?? "");
     expect(tableCount).toBeGreaterThan(0);
   });

packages/server/src/bootstrap-server.ts

@@ -87,8 +87,10 @@ export async function startServer(deps: ServerBootstrapDeps = {}): Promise<void>
   // failed}` logs emitted by `runStage`, which are sufficient for boot
   // analysis without dragging a context onto every downstream span.
 
-  // Run Drizzle migrations before the app comes up. Idempotent under
-  // multi-replica startup (Drizzle journal table); cold-start cost is a few
+  // Run Drizzle migrations before the app comes up. Serialized across
+  // replicas by an advisory lock held for the whole migration (see
+  // db/migrate.ts): one replica applies, the others wait, then no-op off
+  // the advanced journal; cold-start cost is a few
   // hundred ms when there's nothing new to apply. The 20s budget matches
   // the Dockerfile HEALTHCHECK start-period so a migration that truly
   // exceeds it fails the boot fast rather than letting docker judge

packages/server/src/db/migrate.ts

@@ -53,102 +53,101 @@ function validateJournalOrder(migrationsFolder: string): void {
 }
 
 /**
- * Advisory-lock key used by the preflight check.
+ * Advisory-lock key serializing startup migrations across replicas.
  *
  * Verified against `drizzle-orm@0.44.7`:
  *   - `node_modules/drizzle-orm/postgres-js/migrator.js` → delegates to
  *     `node_modules/drizzle-orm/pg-core/dialect.js::migrate()`, which
  *     **does NOT acquire any advisory lock** in this version; it only wraps
  *     `INSERT INTO drizzle.__drizzle_migrations` in a transaction.
- *   - So this preflight is **purely defensive**: it surfaces *external*
- *     holders (an operator's `SELECT pg_advisory_lock(...)`, a stale
- *     prior-replica session that exited mid-migration with the lock held,
- *     or a future drizzle release that re-introduces locking) within the
- *     timeout window instead of letting drizzle hang silently.
+ *   - So `runMigrations` acquires this session-level lock itself and holds
+ *     it on the same single connection for the *entire* migration (journal
+ *     read + DDL/backfills + journal insert). Concurrent replicas wait on
+ *     the key; when the owner finishes and the lock is released, the next
+ *     replica acquires it, sees the advanced journal, and no-ops.
  *
  * If you bump `drizzle-orm`, re-read `pg-core/dialect.js::migrate()` and
- * update this key (and `MIGRATION_LOCK_TIMEOUT_MS`) accordingly. The
- * integration test `bootstrap-migration-lock.test.ts` pins the contention
- * behavior using the same key.
+ * confirm it still does not take a conflicting lock of its own. The
+ * integration test `bootstrap-migration-lock.test.ts` pins both the
+ * contention error path and the held-for-the-whole-migration behavior
+ * using the same key.
  */
 const MIGRATION_LOCK_KEY_SQL = "hashtext('drizzle_migrations')";
-// Sits inside the 20s `runMigrations` stage budget set in `index.ts`; 15s
-// preflight + ≥5s for the actual drizzle migrate call. If you raise either,
-// raise the other in lockstep (and re-evaluate the Dockerfile HEALTHCHECK
-// `--start-period`).
+// Sits inside the 20s `runMigrations` stage budget set in
+// `bootstrap-server.ts`; up to 15s waiting for the lock + ≥5s for the actual
+// drizzle migrate call. If you raise either, raise the other in lockstep
+// (and re-evaluate the Dockerfile HEALTHCHECK `--start-period`).
 const DEFAULT_MIGRATION_LOCK_TIMEOUT_MS = 15_000;
 const MIGRATION_LOCK_POLL_INTERVAL_MS = 1_000;
 
 export type RunMigrationsOptions = {
-  /** Override the preflight advisory-lock timeout. Default 15s. */
+  /** Override the advisory-lock acquisition timeout. Default 15s. */
   lockTimeoutMs?: number;
 };
 
 /**
- * Probe the advisory-lock key drizzle would use for migrations. If another
- * session holds it, fail with a clear message instead of letting drizzle's
- * `migrate()` hang forever. We don't keep the lock — see the key constant
- * above for the full rationale. See server-bootstrap-resilience-design.md §3 (T10).
+ * Acquire the migration advisory lock on `client`'s session, polling until
+ * `timeoutMs`. On success the lock is HELD — the caller keeps it for the
+ * whole migration and releases it implicitly when the session ends. On
+ * timeout, fail with a clear contention message instead of letting a
+ * concurrent replica run DDL in parallel. See the key constant above for
+ * the full rationale and server-bootstrap-resilience-design.md §3 (T10).
  */
-async function preflightMigrationLock(databaseUrl: string, timeoutMs: number): Promise<void> {
-  const ssl = sslOptions(databaseUrl);
-  const client = postgres(databaseUrl, { max: 1, ...ssl });
+async function acquireMigrationLock(client: postgres.Sql, timeoutMs: number): Promise<void> {
   const deadline = Date.now() + timeoutMs;
-  try {
-    while (true) {
-      const rows = (await client.unsafe(
-        `SELECT pg_try_advisory_lock(${MIGRATION_LOCK_KEY_SQL}) AS acquired`,
-      )) as Array<{
-        acquired: boolean;
-      }>;
-      if (rows[0]?.acquired) {
-        await client.unsafe(`SELECT pg_advisory_unlock(${MIGRATION_LOCK_KEY_SQL})`);
-        return;
-      }
-      if (Date.now() >= deadline) {
-        throw new Error(
-          `migration lock contention — another process holds drizzle migration lock (${MIGRATION_LOCK_KEY_SQL}) ` +
-            `after ${timeoutMs}ms`,
-        );
-      }
-      await new Promise((r) => setTimeout(r, MIGRATION_LOCK_POLL_INTERVAL_MS));
+  while (true) {
+    const rows = (await client.unsafe(`SELECT pg_try_advisory_lock(${MIGRATION_LOCK_KEY_SQL}) AS acquired`)) as Array<{
+      acquired: boolean;
+    }>;
+    if (rows[0]?.acquired) {
+      return;
     }
-  } finally {
-    await client.end();
+    if (Date.now() >= deadline) {
+      throw new Error(
+        `migration lock contention — another process holds drizzle migration lock (${MIGRATION_LOCK_KEY_SQL}) ` +
+          `after ${timeoutMs}ms`,
+      );
+    }
+    await new Promise((r) => setTimeout(r, MIGRATION_LOCK_POLL_INTERVAL_MS));
   }
 }
 
 /**
- * Run Drizzle database migrations. Returns the count of public tables after
- * migration, used as a rough indicator that the schema landed.
+ * Run Drizzle database migrations, serialized across replicas by a
+ * session-level advisory lock held on a single dedicated connection for the
+ * entire operation (journal read + DDL/backfills + journal insert). Returns
+ * the count of public tables after migration, used as a rough indicator
+ * that the schema landed.
  */
 export async function runMigrations(databaseUrl: string, options: RunMigrationsOptions = {}): Promise<number> {
   const migrationsFolder = resolveMigrationsFolder();
 
   validateJournalOrder(migrationsFolder);
 
-  await preflightMigrationLock(databaseUrl, options.lockTimeoutMs ?? DEFAULT_MIGRATION_LOCK_TIMEOUT_MS);
-
   const ssl = sslOptions(databaseUrl);
 
+  // One `max: 1` client = one physical session. The advisory lock taken on
+  // it below stays held across the drizzle `migrate()` call (session-level
+  // locks survive the transactions drizzle opens on this same connection)
+  // and is released when the session closes in `finally` — including on
+  // error paths, so a failed migration can't strand the lock.
   const client = postgres(databaseUrl, { max: 1, ...ssl });
-  const db = drizzle(client);
-
   try {
+    await acquireMigrationLock(client, options.lockTimeoutMs ?? DEFAULT_MIGRATION_LOCK_TIMEOUT_MS);
+
+    const db = drizzle(client);
     await migrate(db, { migrationsFolder });
-  } finally {
-    await client.end();
-  }
 
-  const countClient = postgres(databaseUrl, { max: 1, ...ssl });
-  try {
-    const result = await countClient`
+    const result = await client`
       SELECT count(*)::int AS count
       FROM information_schema.tables
       WHERE table_schema = 'public' AND table_type = 'BASE TABLE'
     `;
     return (result[0] as { count: number }).count;
   } finally {
-    await countClient.end();
+    // Closing the session releases the advisory lock server-side; no
+    // explicit pg_advisory_unlock needed (and none that could mask an
+    // in-flight migration error if the connection already died).
+    await client.end();
   }
 }