From e1d2e279a3528dc98a0a21e176a9e3d0ae65fe83 Mon Sep 17 00:00:00 2001 From: Devin Foley Date: Mon, 14 Sep 2026 09:33:48 -0700 Subject: [PATCH] fix(db): replay a query whose socket write failed on a recycled pooled connection (#13417) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work > - A hosted instance keeps its data in Postgres, and it reaches that database through a connection pooler > - A pooler recycles the server side of a connection that sat idle. The client still holds the socket and believes it is open > - The next query fails while the driver writes it to that dead socket, with `write CONNECTION_CLOSED host:5432` > - The request that happened to draw the recycled connection fails, although nothing is wrong with the query or the database > - Only one code path guards against this today, so the same error keeps reaching users and error reporting from every other path > - This pull request replays a query when the write itself failed, because those bytes never reached the server > - The benefit is that a recycled connection costs one retry instead of one failed request ## Linked Issues or Issue Description No existing issue. The problem, in the bug report format: **What happened** Requests fail with `write CONNECTION_CLOSED :5432` (driver code `CONNECTION_CLOSED`). It is most visible on the first queries after an idle period, when the pool holds connections the pooler already recycled. **Expected behavior** A connection that the pooler recycled while it was idle is not a user-visible failure. The client notices the dead socket and uses a live connection. **Steps to reproduce** 1. Run the server against a pooled Postgres endpoint (for example a Neon `-pooler` host). 2. Leave the instance idle until the pooler recycles the server side of the pooled connections. 3. Issue any request that queries the database. **Paperclip version or commit** Present on master. ## What Changed - `packages/db/src/transient-write-retry.ts` (new) wraps the root `postgres.js` client. When a query fails while the driver writes it to the socket, the wrapper runs it again on a fresh connection. Three attempts, 50 ms then 100 ms backoff. - `packages/db/src/client.ts` gives Drizzle the wrapped client. The teardown registry keeps the real client, because shutdown must end the actual pool. - The wrapper is deliberately narrow: - It matches only the write phase (`code === "CONNECTION_CLOSED"` and a message that starts with `write CONNECTION_CLOSED`). The write failed, so the server never saw the query, and a replay cannot run anything twice. That makes it safe for reads and writes alike. - An error after the write propagates untouched, because the server may have acted on the query. - `CONNECTION_ENDED` and `CONNECTION_DESTROYED` propagate untouched, because they mean a deliberate shutdown. - Queries inside `db.transaction()` run on the scoped client that `sql.begin()` returns, which the wrapper does not touch. A transaction that loses its connection must abort, not replay. - A query runs once however many handlers attach to it, and it still executes lazily, like `postgres.js` itself. ## Verification ``` npx vitest run packages/db/src/transient-write-retry.test.ts # 7 passed npx vitest run packages/db/src/client-teardown-registry.test.ts \ packages/db/src/client-options.test.ts # existing db suites pass npx vitest run server/src/__tests__/cloud-tenant-transient-db-retry.test.ts # 6 passed cd packages/db && npx tsc --noEmit # clean ``` New tests cover the replay, the `.values()` form Drizzle uses, the attempt budget, an error that must not replay, and one execution per pending query. A Drizzle round trip runs through a fake wire-protocol server, which shows the wrapper is transparent to ordinary queries. ## Risks Low risk. - The retry only fires for a failure during the socket write, where the server never received the query. A replay therefore cannot duplicate an effect. - A permanently unreachable database costs two extra attempts and 150 ms before the same error surfaces. - Transactions keep exactly their current behavior. - Existing `retryOnTransientDbConnectionError` in the auth middleware stays. It wraps a broader set of codes for one path, and it is unaffected. - No migration. No configuration change. No API change. ## Model Used - Claude Fable 5 (`claude-fable-5`), 1M context, extended thinking, run through Claude Code with tool use and code execution. ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with `Fixes: #` / `Closes #` / `Refs #` OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip` URLs) - [x] My branch name describes the change (e.g. `docs/...`, `fix/...`) and contains no internal Paperclip ticket id or instance-derived details - [x] I have run tests locally and they pass - [x] I have added or updated tests where applicable - [ ] I have updated relevant documentation to reflect my changes — no documented behavior or configuration changes - [x] I have considered and documented any risks above - [ ] All Paperclip CI gates are green — pending first run - [ ] Greptile is 5/5 with no open P2s, recommendations, or follow-ups — pending first review - [x] I will address all Greptile and reviewer comments before requesting merge --- packages/db/src/client.ts | 6 +- packages/db/src/transient-write-retry.test.ts | 191 ++++++++++++++++++ packages/db/src/transient-write-retry.ts | 90 +++++++++ 3 files changed, 286 insertions(+), 1 deletion(-) create mode 100644 packages/db/src/transient-write-retry.test.ts create mode 100644 packages/db/src/transient-write-retry.ts 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; +}