From ea5444f730b8f6572347b1323fa1ddcb320ac796 Mon Sep 17 00:00:00 2001 From: Serena <94026305+serenakeyitan@users.noreply.github.com> Date: Thu, 23 Jul 2026 15:45:38 -0700 Subject: [PATCH] fix(server): hold migration advisory lock across the whole migration The old preflight acquired pg_try_advisory_lock(hashtext('drizzle_migrations')) and released it immediately, then ran drizzle migrate() on a separate, unlocked connection. Under multi-replica rollout every replica could pass preflight simultaneously and execute DDL/backfills in parallel, causing duplicate-object errors, duplicated backfill work, or failed boot replicas. runMigrations now opens one max:1 connection, acquires the session-level advisory lock on it (same key, same 15s poll/timeout, same contention error message), and keeps it held across the drizzle migrate() call and the table count; closing the session in finally releases the lock on both success and error paths. Non-owner replicas keep polling; once the owner releases, they acquire the lock, see the advanced journal, and no-op. Closes #1690 --- .../bootstrap-migration-lock.test.ts | 121 +++++++++++++++--- packages/server/src/bootstrap-server.ts | 6 +- packages/server/src/db/migrate.ts | 107 ++++++++-------- 3 files changed, 163 insertions(+), 71 deletions(-) diff --git a/packages/server/src/__tests__/bootstrap-migration-lock.test.ts b/packages/server/src/__tests__/bootstrap-migration-lock.test.ts index 81ccdb8a5..26d286694 100644 --- a/packages/server/src/__tests__/bootstrap-migration-lock.test.ts +++ b/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); }); diff --git a/packages/server/src/bootstrap-server.ts b/packages/server/src/bootstrap-server.ts index b80212e09..4305eecc0 100644 --- a/packages/server/src/bootstrap-server.ts +++ b/packages/server/src/bootstrap-server.ts @@ -87,8 +87,10 @@ export async function startServer(deps: ServerBootstrapDeps = {}): Promise // 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 diff --git a/packages/server/src/db/migrate.ts b/packages/server/src/db/migrate.ts index 84ebe7aea..f656d87f4 100644 --- a/packages/server/src/db/migrate.ts +++ b/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 { - const ssl = sslOptions(databaseUrl); - const client = postgres(databaseUrl, { max: 1, ...ssl }); +async function acquireMigrationLock(client: postgres.Sql, timeoutMs: number): Promise { 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 { 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(); } }