diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index ed7de6716f5..70030b051d3 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -5,6 +5,7 @@ import { readFile, readdir } from "node:fs/promises"; import { fileURLToPath } from "node:url"; import postgres from "postgres"; import * as schema from "./schema/index.js"; +import { withTransientWriteRetry } from "./transient-write-retry.js"; const MIGRATIONS_FOLDER = fileURLToPath(new URL("./migrations", import.meta.url)); const DRIZZLE_MIGRATIONS_TABLE = "__drizzle_migrations"; @@ -259,7 +260,10 @@ export function createDb(url: string, options?: DatabaseClientOptions) { const sql = postgres(url, postgresJsOptions(resolved)); const key = hostPortKeyOrNull(url); if (key) registerClient(key, sql); - return drizzlePg(sql, { schema }); + // The registry keeps the real client (teardown must end the actual pool); + // drizzle gets the retrying face so a pooler-recycled socket replays the + // query instead of failing the request that happened to draw it. + return drizzlePg(withTransientWriteRetry(sql), { schema }); } export async function getPostgresDataDirectory(url: string): Promise { diff --git a/packages/db/src/transient-write-retry.test.ts b/packages/db/src/transient-write-retry.test.ts new file mode 100644 index 00000000000..b144cb2c779 --- /dev/null +++ b/packages/db/src/transient-write-retry.test.ts @@ -0,0 +1,191 @@ +import net from "node:net"; +import { sql as drizzleSql } from "drizzle-orm"; +import { afterEach, describe, expect, it } from "vitest"; +import type { Sql } from "postgres"; +import { closeRegisteredClients, createDb } from "./client.js"; +import { isTransientWritePhaseError, withTransientWriteRetry } from "./transient-write-retry.js"; + +/** The exact shape postgres.js raises when the socket write fails. */ +function writeClosedError(): Error { + const error = new Error("write CONNECTION_CLOSED db.example.internal:5432"); + (error as Error & { code: string }).code = "CONNECTION_CLOSED"; + return error; +} + +/** A stub standing in for the root postgres.js client's `unsafe` surface. */ +function stubSql(behavior: { failures: number; error?: () => Error }) { + let calls = 0; + const state = { + get calls() { + return calls; + }, + }; + const unsafe = (_query: string, _parameters?: unknown[]) => { + calls += 1; + const attempt = calls; + const shouldFail = attempt <= behavior.failures; + const outcome = () => + shouldFail + ? Promise.reject((behavior.error ?? writeClosedError)()) + : Promise.resolve([{ attempt }]); + return { + then: (onFulfilled?: ((value: unknown) => unknown) | null, onRejected?: ((reason: unknown) => unknown) | null) => + outcome().then(onFulfilled, onRejected), + values: () => outcome().then(() => [[attempt]]), + }; + }; + const sql = { unsafe } as unknown as Sql; + return { sql, state }; +} + +describe("isTransientWritePhaseError", () => { + it("matches only the write-phase connection-closed shape", () => { + expect(isTransientWritePhaseError(writeClosedError())).toBe(true); + + const midQuery = new Error("read CONNECTION_CLOSED db.example.internal:5432"); + (midQuery as Error & { code: string }).code = "CONNECTION_CLOSED"; + expect(isTransientWritePhaseError(midQuery)).toBe(false); + + const ended = new Error("write CONNECTION_ENDED db.example.internal:5432"); + (ended as Error & { code: string }).code = "CONNECTION_ENDED"; + expect(isTransientWritePhaseError(ended)).toBe(false); + + expect(isTransientWritePhaseError(new Error("write CONNECTION_CLOSED x:1"))).toBe(false); // no code + expect(isTransientWritePhaseError(null)).toBe(false); + }); +}); + +describe("withTransientWriteRetry", () => { + it("replays a query whose socket write failed, and returns the replay's rows", async () => { + const { sql, state } = stubSql({ failures: 1 }); + const rows = await withTransientWriteRetry(sql).unsafe("select 1", []); + expect(rows).toEqual([{ attempt: 2 }]); + expect(state.calls).toBe(2); + }); + + it("replays the .values() form drizzle uses", async () => { + const { sql, state } = stubSql({ failures: 2 }); + const values = await withTransientWriteRetry(sql).unsafe("select 1", []).values(); + expect(values).toEqual([[3]]); + expect(state.calls).toBe(3); + }); + + it("gives up after the attempt budget and surfaces the driver error", async () => { + const { sql, state } = stubSql({ failures: Number.POSITIVE_INFINITY }); + await expect(withTransientWriteRetry(sql).unsafe("select 1", [])).rejects.toThrow( + "write CONNECTION_CLOSED", + ); + expect(state.calls).toBe(3); + }); + + it("does not replay an error that may have reached the server", async () => { + const midQuery = () => { + const error = new Error("read CONNECTION_CLOSED db.example.internal:5432"); + (error as Error & { code: string }).code = "CONNECTION_CLOSED"; + return error; + }; + const { sql, state } = stubSql({ failures: Number.POSITIVE_INFINITY, error: midQuery }); + await expect(withTransientWriteRetry(sql).unsafe("select 1", [])).rejects.toThrow( + "read CONNECTION_CLOSED", + ); + expect(state.calls).toBe(1); + }); + + it("runs one execution no matter how many handlers attach to the pending query", async () => { + const { sql, state } = stubSql({ failures: 0 }); + const pending = withTransientWriteRetry(sql).unsafe("select 1", []); + await Promise.all([pending, pending.catch(() => undefined), pending.then((rows) => rows)]); + expect(state.calls).toBe(1); + }); + + it("never executes a query twice when a caller both awaits it and takes its values", async () => { + // `.values()` selects the row shape of the one execution, like the driver. + // Starting a second execution here would repeat a mutation's effect. + const { sql, state } = stubSql({ failures: 0 }); + const pending = withTransientWriteRetry(sql).unsafe("insert into t values (1)", []); + const rows = await pending; + const values = await pending.values(); + expect(state.calls).toBe(1); + expect(values).toBe(rows); + }); + + it("takes the values shape when it is chosen before the query runs", async () => { + const { sql, state } = stubSql({ failures: 0 }); + const pending = withTransientWriteRetry(sql).unsafe("select 1", []); + expect(await pending.values()).toEqual([[1]]); + expect(state.calls).toBe(1); + }); +}); + +describe("createDb with the retrying client", () => { + // The same minimal wire-protocol fake as client-teardown-registry.test.ts: + // enough of the startup and query flow for drizzle to run a real query + // through the proxied client, proving the retry face is transparent. + function startFakePostgresServer(): Promise<{ server: net.Server; port: number }> { + const authOk = Buffer.from([0x52, 0, 0, 0, 8, 0, 0, 0, 0]); + const readyForQuery = Buffer.from([0x5a, 0, 0, 0, 5, 0x49]); + const emptyQueryReply = Buffer.concat([ + Buffer.from([0x31, 0, 0, 0, 4]), + Buffer.from([0x32, 0, 0, 0, 4]), + Buffer.from([0x54, 0, 0, 0, 6, 0, 0]), + Buffer.from([0x43, 0, 0, 0, 0x0d, 0x53, 0x45, 0x4c, 0x45, 0x43, 0x54, 0x20, 0x30, 0]), + ]); + const server = net.createServer((socket) => { + let greeted = false; + socket.on("data", () => { + if (!greeted) { + greeted = true; + socket.write(Buffer.concat([authOk, readyForQuery])); + return; + } + socket.write(Buffer.concat([emptyQueryReply, readyForQuery])); + }); + socket.on("error", () => {}); + }); + return new Promise((resolve) => { + server.listen(0, "127.0.0.1", () => { + resolve({ server, port: (server.address() as net.AddressInfo).port }); + }); + }); + } + + let server: net.Server | null = null; + let url: string | null = null; + + afterEach(async () => { + if (url) await closeRegisteredClients(url); + if (server) await new Promise((resolve) => server!.close(resolve)); + server = null; + url = null; + }); + + it("still answers ordinary drizzle queries through the proxy", async () => { + const started = await startFakePostgresServer(); + server = started.server; + url = `postgres://test:test@127.0.0.1:${started.port}/test`; + const db = createDb(url, { connectTimeoutSeconds: 5, prepare: false }); + await expect(db.execute(drizzleSql`select 0`)).resolves.toBeDefined(); + }); + + it("leaves every other client surface reachable through the proxy", async () => { + // Drizzle itself only awaits `unsafe()` or takes its `.values()`, but the + // client is reachable as `db.$client`, and callers use it as a tagged + // template, open transactions on it, and end it. A wrapper that broke any + // of those would fail far from here, so pin them against the real driver. + const started = await startFakePostgresServer(); + server = started.server; + url = `postgres://test:test@127.0.0.1:${started.port}/test`; + const db = createDb(url, { connectTimeoutSeconds: 5, prepare: false }); + const client = (db as unknown as { $client: Sql }).$client; + + await expect(client`select 1`).resolves.toBeDefined(); + await expect(client.unsafe("select 1", [])).resolves.toBeDefined(); + await expect( + client.begin(async (tx) => { + await tx.unsafe("select 1", []); + return "transaction result"; + }), + ).resolves.toBe("transaction result"); + await expect(client.end({ timeout: 1 })).resolves.toBeUndefined(); + }); +}); diff --git a/packages/db/src/transient-write-retry.ts b/packages/db/src/transient-write-retry.ts new file mode 100644 index 00000000000..6d65751345f --- /dev/null +++ b/packages/db/src/transient-write-retry.ts @@ -0,0 +1,90 @@ +import type { Sql } from "postgres"; + +/** + * Replays a query whose bytes never reached the server. + * + * Behind a connection pooler (Neon's PgBouncer), the server side of an idle + * connection can be recycled while the client still holds the socket. The + * next query then fails at the socket write — postgres.js reports it as + * `write CONNECTION_CLOSED :` with `code: "CONNECTION_CLOSED"`. + * Because the write itself failed, the server never saw the query, so + * replaying it on a fresh connection cannot double-execute anything — the + * replay is safe for reads and writes alike. Every other error, including a + * connection lost mid-query (where the server may have acted), propagates + * untouched. + * + * Drizzle routes every non-transactional query through the root client's + * `unsafe`; queries inside `db.transaction()` run on the scoped client + * `sql.begin()` hands out, which this wrapper deliberately does not touch — + * a transaction that loses its connection must abort, not replay. + */ +const TRANSIENT_WRITE_ATTEMPTS = 3; +const TRANSIENT_WRITE_BACKOFF_MS = 50; + +export function isTransientWritePhaseError(error: unknown): boolean { + return ( + error instanceof Error && + (error as { code?: unknown }).code === "CONNECTION_CLOSED" && + error.message.startsWith("write CONNECTION_CLOSED") + ); +} + +type UnsafeFn = (query: string, parameters?: unknown[]) => { + values: () => Promise; +} & PromiseLike; + +export function withTransientWriteRetry(sql: T): T { + const runWithRetry = async (execute: () => Promise): Promise => { + for (let attempt = 0; ; attempt++) { + try { + return await execute(); + } catch (error) { + if (attempt >= TRANSIENT_WRITE_ATTEMPTS - 1 || !isTransientWritePhaseError(error)) { + throw error; + } + await new Promise((resolve) => setTimeout(resolve, TRANSIENT_WRITE_BACKOFF_MS * (attempt + 1))); + } + } + }; + + return new Proxy(sql, { + get(target, property, receiver) { + if (property !== "unsafe") return Reflect.get(target, property, receiver); + const unsafe = target.unsafe.bind(target) as UnsafeFn; + const retryingUnsafe = (query: string, parameters?: unknown[], ...rest: unknown[]) => { + if (rest.length > 0) { + // An options argument selects driver surfaces this wrapper does not + // model; leave those calls exactly as they were. + return (target.unsafe as (...args: unknown[]) => unknown)(query, parameters, ...rest); + } + // postgres.js queries execute lazily, exactly once, and `.values()` + // selects the row shape of that one execution rather than starting a + // second one. Mirror both halves: the execution starts when a consumer + // first settles the query, `.values()` marks the shape and returns the + // same pending query, and every later handler joins the same run. A + // shape chosen after the run started cannot change it — the same as + // the driver, and the reason a mutation can never execute twice here. + let started: Promise | undefined; + let wantsValues = false; + const runOnce = () => + (started ??= runWithRetry(() => + wantsValues + ? unsafe(query, parameters).values() + : Promise.resolve(unsafe(query, parameters)), + )); + const pending = { + then: (onFulfilled?: ((value: unknown) => unknown) | null, onRejected?: ((reason: unknown) => unknown) | null) => + runOnce().then(onFulfilled, onRejected), + catch: (onRejected?: ((reason: unknown) => unknown) | null) => runOnce().catch(onRejected), + finally: (onFinally?: (() => void) | null) => runOnce().finally(onFinally), + values: () => { + if (!started) wantsValues = true; + return pending; + }, + }; + return pending; + }; + return retryingUnsafe; + }, + }) as T; +}