diff --git a/internal-packages/testcontainers/package.json b/internal-packages/testcontainers/package.json index b42f05863bc..35f29b33652 100644 --- a/internal-packages/testcontainers/package.json +++ b/internal-packages/testcontainers/package.json @@ -15,12 +15,15 @@ "dependencies": { "@clickhouse/client": "^1.11.1", "@trigger.dev/database": "workspace:*", - "ioredis": "~5.6.0" + "ioredis": "~5.6.0", + "pg": "8.15.6" }, "devDependencies": { "@internal/run-ops-database": "workspace:*", + "@prisma/adapter-pg": "6.14.0", "@testcontainers/postgresql": "^11.14.0", "@testcontainers/redis": "^11.14.0", + "@types/pg": "8.11.14", "std-env": "^3.9.0", "testcontainers": "^11.14.0", "tinyexec": "^0.3.0" diff --git a/internal-packages/testcontainers/src/dbBlip.test.ts b/internal-packages/testcontainers/src/dbBlip.test.ts new file mode 100644 index 00000000000..75474a7519b --- /dev/null +++ b/internal-packages/testcontainers/src/dbBlip.test.ts @@ -0,0 +1,177 @@ +import { describe, expect } from "vitest"; +import { Pool } from "pg"; +import { PrismaPg } from "@prisma/adapter-pg"; +import { PrismaClient } from "@trigger.dev/database"; +import { postgresBlipTest } from "./index"; + +// A minimal infra retry, standing in for the shared read-retry util so this +// file can demonstrate the harness end-to-end on its own. +async function withRetry(fn: () => Promise, maxAttempts = 8): Promise { + let lastError: unknown; + for (let attempt = 0; attempt < maxAttempts; attempt++) { + try { + return await fn(); + } catch (error) { + lastError = error; + await new Promise((r) => setTimeout(r, Math.min(50 * (attempt + 1), 250))); + } + } + throw lastError; +} + +// Production runs the pg driver adapter, so the client under test is adapter-backed. +async function adapterClient(connectionString: string) { + const pool = new Pool({ connectionString }); + // A severed idle connection makes the pg Pool emit 'error'; swallow it so an + // unhandled event can't crash the test worker before recovery is asserted. + pool.on("error", () => {}); + const client = new PrismaClient({ adapter: new PrismaPg(pool) }); + const dispose = async () => { + try { + await client.$disconnect(); + } finally { + await pool.end(); + } + }; + return { client, dispose }; +} + +async function createProbeTable(client: PrismaClient) { + await client.$executeRawUnsafe( + `CREATE TABLE IF NOT EXISTS blip_probe (id uuid PRIMARY KEY, tag text NOT NULL)` + ); +} + +async function countTag(client: PrismaClient, tag: string): Promise { + const rows = await client.$queryRawUnsafe<{ n: number }[]>( + `SELECT count(*)::int AS n FROM blip_probe WHERE tag = $1`, + tag + ); + return rows[0]?.n ?? 0; +} + +describe("DbBlipController", () => { + postgresBlipTest( + "a pooled adapter client transparently survives an idle-connection drop", + { timeout: 60_000 }, + async ({ postgresContainer, blip }) => { + const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri()); + try { + await client.user.count(); // warm the pool + const terminated = await blip.severIdle(); + expect(terminated).toBeGreaterThan(0); + // The pool evicts the dead idle connection; the next read just works. + await new Promise((r) => setTimeout(r, 200)); + const count = await client.user.count(); + expect(typeof count).toBe("number"); + } finally { + await dispose(); + } + } + ); + + postgresBlipTest( + "severDuringNextStatement fails an in-flight statement", + { timeout: 60_000 }, + async ({ postgresContainer, blip }) => { + const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri()); + try { + const slow = client.$queryRawUnsafe(`SELECT pg_sleep(3)`); + // PrismaPromise is lazy — form the assertion so the query actually starts. + const rejected = expect(slow).rejects.toThrow(); + await blip.severDuringNextStatement({ queryContains: "pg_sleep" }); + await rejected; + } finally { + await dispose(); + } + } + ); + + postgresBlipTest( + "a read recovers after a mid-flight blip", + { timeout: 60_000 }, + async ({ postgresContainer, blip }) => { + const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri()); + try { + const severed = client.$queryRawUnsafe(`SELECT pg_sleep(3)`).catch(() => undefined); + await blip.severDuringNextStatement({ queryContains: "pg_sleep" }); + await severed; + const count = await withRetry(() => client.user.count()); + expect(typeof count).toBe("number"); + } finally { + await dispose(); + } + } + ); + + postgresBlipTest( + "a non-idempotent write double-applies on retry after a post-commit blip; the idempotent form does not", + { timeout: 60_000 }, + async ({ postgresContainer, blip }) => { + const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri()); + try { + await createProbeTable(client); + + // Model the dangerous case: the write commits, then a later statement in + // the same op is severed mid-flight (ack lost), and the caller retries. + let nonIdempotentAttempts = 0; + const nonIdempotentWrite = async () => { + nonIdempotentAttempts++; + await client.$executeRawUnsafe( + `INSERT INTO blip_probe (id, tag) VALUES (gen_random_uuid(), 'non-idempotent')` + ); + if (nonIdempotentAttempts === 1) { + await client.$queryRawUnsafe(`SELECT pg_sleep(3)`); // severed → throws after the commit + } + }; + const nonIdempotentDone = withRetry(nonIdempotentWrite); + await blip.severDuringNextStatement({ queryContains: "pg_sleep" }); + await nonIdempotentDone; + expect(await countTag(client, "non-idempotent")).toBe(2); // the hazard, proven + + // The idempotent form: a fixed id + ON CONFLICT makes the replay a no-op. + let idempotentAttempts = 0; + const idempotentWrite = async () => { + idempotentAttempts++; + await client.$executeRawUnsafe( + `INSERT INTO blip_probe (id, tag) + VALUES ('00000000-0000-0000-0000-000000000001', 'idempotent') + ON CONFLICT (id) DO NOTHING` + ); + if (idempotentAttempts === 1) { + await client.$queryRawUnsafe(`SELECT pg_sleep(3)`); + } + }; + const idempotentDone = withRetry(idempotentWrite); + await blip.severDuringNextStatement({ queryContains: "pg_sleep" }); + await idempotentDone; + expect(await countTag(client, "idempotent")).toBe(1); // exactly once despite retry + } finally { + await dispose(); + } + } + ); + + // Regression: queryContains must match as literal text, not as an ILIKE pattern. + // The active query contains "fooXbar"; under ILIKE the pattern "foo_bar" (with the + // wildcard `_`) would wrongly match and terminate it. The literal matcher must not, + // so the sever times out instead of killing the wrong statement. + postgresBlipTest( + "severDuringNextStatement matches queryContains literally, not as an ILIKE pattern", + { timeout: 60_000 }, + async ({ postgresContainer, blip }) => { + const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri()); + const slow = client + .$queryRawUnsafe(`SELECT pg_sleep(3) /* marker fooXbar */`) + .catch(() => undefined); + try { + await expect( + blip.severDuringNextStatement({ queryContains: "foo_bar", timeoutMs: 1000, pollMs: 25 }) + ).rejects.toThrow(/no active statement/i); + } finally { + await dispose(); + await slow; + } + } + ); +}); diff --git a/internal-packages/testcontainers/src/dbBlip.ts b/internal-packages/testcontainers/src/dbBlip.ts new file mode 100644 index 00000000000..3042318ef8e --- /dev/null +++ b/internal-packages/testcontainers/src/dbBlip.ts @@ -0,0 +1,107 @@ +import { Client } from "pg"; + +/** + * Simulates a connection blip against a test Postgres (via a separate admin + * connection that terminates backends), so a vertical can prove its DB code + * survives a disconnect. Reproduces the mid-statement / stale-connection + * signatures (P1017, "Connection terminated unexpectedly"). + */ +export type DbBlipController = { + /** Terminate every idle client backend except this harness's own, so the + * next operation hits a dead connection. Returns the number terminated. */ + severIdle(): Promise; + + /** Poll for an active client statement (optionally matching `queryContains` + * literally), then terminate it mid-flight. Rejects if none appears within + * `timeoutMs`. Terminating by pid isn't atomic with statement completion, so + * target a statement with a real execution window (e.g. `pg_sleep`) — a query + * that finishes first leaves its connection idle and it is closed anyway. */ + severDuringNextStatement(opts?: { + queryContains?: string; + timeoutMs?: number; + pollMs?: number; + }): Promise; +}; + +/** A {@link DbBlipController} plus the teardown for its admin connection. */ +export type DbBlipHandle = DbBlipController & { close(): Promise }; + +// Reserved application_name for the harness's control connections. The severs +// exclude every connection using it (by name, plus their own pid), so multiple +// controllers on one database can't kill each other's admin. A client-under-test +// must not use this name. +const ADMIN_APPLICATION_NAME = "trigger-db-blip-admin"; + +const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +/** Opens an isolated admin connection and returns a handle that can sever the + * other connections on that database. `close()` in teardown. */ +export async function createDbBlipController(connectionUri: string): Promise { + // Raw pg (not Prisma): the control connection must be one identifiable backend we can exclude from the sever, independent of the client under test. + const admin = new Client({ + connectionString: connectionUri, + application_name: ADMIN_APPLICATION_NAME, + }); + await admin.connect(); + // Swallow async connection errors so a consumer that severs a DB the admin + // isn't excluded from (or drops it while open) can't crash the test worker. + admin.on("error", () => {}); + + async function severIdle(): Promise { + const result = await admin.query<{ terminated: boolean }>( + `SELECT pg_terminate_backend(pid) AS terminated + FROM pg_stat_activity + WHERE datname = current_database() + AND pid <> pg_backend_pid() + AND backend_type = 'client backend' + AND application_name IS DISTINCT FROM $1 + AND state = 'idle'`, + [ADMIN_APPLICATION_NAME] + ); + // Count only backends that were actually terminated (a backend that exits + // between selection and signalling returns false). + return result.rows.filter((row) => row.terminated === true).length; + } + + async function severDuringNextStatement(opts?: { + queryContains?: string; + timeoutMs?: number; + pollMs?: number; + }): Promise { + const queryContains = opts?.queryContains ?? null; + const timeoutMs = opts?.timeoutMs ?? 5000; + const pollMs = opts?.pollMs ?? 25; + const deadline = Date.now() + timeoutMs; + + while (Date.now() < deadline) { + // Select and terminate in one statement so the backend can't go idle + // between picking it and killing it; return only when it was terminated. + const terminated = await admin.query<{ ok: boolean }>( + `SELECT pg_terminate_backend(pid) AS ok + FROM pg_stat_activity + WHERE datname = current_database() + AND state = 'active' + AND pid <> pg_backend_pid() + AND backend_type = 'client backend' + AND application_name IS DISTINCT FROM $1 + AND ($2::text IS NULL OR strpos(lower(query), lower($2)) > 0) + LIMIT 1`, + [ADMIN_APPLICATION_NAME, queryContains] + ); + + if (terminated.rows[0]?.ok === true) { + return; + } + + await sleep(pollMs); + } + + throw new Error( + `severDuringNextStatement: no active statement${ + queryContains ? ` matching ${JSON.stringify(queryContains)}` : "" + } appeared within ${timeoutMs}ms` + ); + } + + return { severIdle, severDuringNextStatement, close: () => admin.end() }; +} diff --git a/internal-packages/testcontainers/src/index.ts b/internal-packages/testcontainers/src/index.ts index e72a63766e9..ba923d463e8 100644 --- a/internal-packages/testcontainers/src/index.ts +++ b/internal-packages/testcontainers/src/index.ts @@ -13,6 +13,7 @@ import { runClickhouseMigrations, truncateClickhouseTables, } from "./clickhouse"; +import { createDbBlipController, type DbBlipController } from "./dbBlip"; import { getTaskMetadata, logCleanup, logSetup } from "./logs"; import { type MinIOConnectionConfig, type StartedMinIOContainer, MinIOContainer } from "./minio"; import { @@ -35,6 +36,7 @@ export { } from "./utils"; export { OtelCollectorContainer, StartedOtelCollectorContainer } from "./otelCollector"; export { laggingReplica, type LaggingModel } from "./laggingReplica"; +export { createDbBlipController, type DbBlipController, type DbBlipHandle } from "./dbBlip"; export { logCleanup }; export type { MinIOConnectionConfig }; @@ -353,6 +355,32 @@ export const postgresTest = withWarmup( } ); +export type PostgresBlipTestContext = PostgresTestContext & { blip: DbBlipController }; + +const blipFromContainer = async ( + { postgresContainer }: { postgresContainer: StartedPostgreSqlContainer } & TestContext, + use: Use +) => { + const handle = await createDbBlipController(postgresContainer.getConnectionUri()); + try { + await use(handle); + } finally { + await handle.close(); + } +}; + +// postgresTest + a DbBlipController bound to the same per-test database. +export const postgresBlipTest = withWarmup( + test.extend({ + postgresContainer: clonedPostgresContainer, + prisma: prismaFromContainer, + blip: blipFromContainer, + }), + async () => { + await getWorkerPostgresContainer(); + } +); + type HeteroPostgresTestContext = { // PG14 (legacy / control-plane DB analog) postgresContainer14: StartedPostgreSqlContainer; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4c9442bfcb1..1b11e1627e3 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1389,16 +1389,25 @@ importers: ioredis: specifier: ~5.6.0 version: 5.6.1 + pg: + specifier: 8.15.6 + version: 8.15.6 devDependencies: '@internal/run-ops-database': specifier: workspace:* version: link:../run-ops-database + '@prisma/adapter-pg': + specifier: 6.14.0 + version: 6.14.0 '@testcontainers/postgresql': specifier: ^11.14.0 version: 11.14.0 '@testcontainers/redis': specifier: ^11.14.0 version: 11.14.0 + '@types/pg': + specifier: 8.11.14 + version: 8.11.14 std-env: specifier: ^3.9.0 version: 3.9.0