diff --git a/packages/db/src/client-teardown-registry.test.ts b/packages/db/src/client-teardown-registry.test.ts new file mode 100644 index 0000000000..7e3414e01e --- /dev/null +++ b/packages/db/src/client-teardown-registry.test.ts @@ -0,0 +1,112 @@ +import net from "node:net"; +import { afterEach, describe, expect, it } from "vitest"; +import { closeRegisteredClients, createDb } from "./client.js"; + +/** + * A fake wire-protocol server. It speaks just enough of the startup and + * query flow to hand a postgres.js client an open connection: it replies to + * the startup message with `AuthenticationOk` plus `ReadyForQuery`, then + * replies to any later message with an empty result set plus + * `ReadyForQuery`. No real query needs to succeed for this test. + */ +function startFakePostgresServer(): Promise<{ server: net.Server; port: number; backendSockets: net.Socket[] }> { + const backendSockets: net.Socket[] = []; + 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]), // ParseComplete + Buffer.from([0x32, 0, 0, 0, 4]), // BindComplete + Buffer.from([0x54, 0, 0, 0, 6, 0, 0]), // RowDescription, zero fields + Buffer.from([0x43, 0, 0, 0, 0x0d, 0x53, 0x45, 0x4c, 0x45, 0x43, 0x54, 0x20, 0x30, 0]), // CommandComplete "SELECT 0" + ]); + + const server = net.createServer((socket) => { + backendSockets.push(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", () => { + const port = (server.address() as net.AddressInfo).port; + resolve({ server, port, backendSockets }); + }); + }); +} + +describe("closeRegisteredClients", () => { + let server: net.Server | null = null; + + afterEach(async () => { + if (server) await new Promise((resolve) => server!.close(resolve)); + server = null; + }); + + it("ends a reserved connection before its backend dies, so no query can reach a null socket", async () => { + const started = await startFakePostgresServer(); + server = started.server; + const url = `postgres://test:test@127.0.0.1:${started.port}/test`; + + const db = createDb(url, { connectTimeoutSeconds: 5 }); + // `sql.reserve()` pins one physical connection. Drizzle `db.transaction()` + // reaches the same surface through `sql.begin()`, so this stands in for a + // suite that left a transaction connection open. + const reserved = await db.$client.reserve(); + + // The driver calls this only after it has fully processed a connection + // close: its socket reference cleared and any in-flight query failed. + // Waiting for it, instead of a fixed number of ticks, is what the + // historical crash reproduction does — it is real observed state from + // the driver, not a guess at timing. + const driverProcessedClose = new Promise((resolve) => { + db.$client.options.onclose = () => resolve(); + }); + + // This is the order our fixture owns: end every registered client for + // this host and port before a caller stops the cluster it points at. + await closeRegisteredClients(url); + + // Simulate the cluster stop that follows in the real fixture. Before the + // fix, killing the backend here while a client still held the reserved + // connection open crashed the process on a later deferred write. + for (const socket of started.backendSockets) socket.destroy(); + await driverProcessedClose; + + // A query sent only after the driver finished processing the close still + // buffers its frame for a deferred flush one tick later. If the fix let + // the reserved connection outlive the backend, that flush reaches a + // cleared socket reference and throws from inside the timer callback — + // this specific promise then never settles, because nothing on that + // path ever calls its resolve or reject. This test's own timeout (the + // suite default) is what turns that hang into a reported failure, + // alongside the unhandled exception the crash raises separately. + const settled = await reserved`select 1`.catch((error: unknown) => error); + expect(settled).toBeInstanceOf(Error); + + // Let the deferred flush actually run. If it still fires against a null + // socket, it surfaces here as an unhandled error and fails this test + // file — the exact signature the fix protects against. + await new Promise((resolve) => setImmediate(() => setImmediate(resolve))); + }); + + it("does nothing when no client is registered for a host and port", async () => { + await expect(closeRegisteredClients("postgres://test:test@127.0.0.1:1/test")).resolves.toBeUndefined(); + }); + + it("does not throw when createDb receives a URL that new URL() cannot parse", async () => { + let db: ReturnType | undefined; + expect(() => { + db = createDb("", { connectTimeoutSeconds: 1 }); + }).not.toThrow(); + + await db?.$client.end({ timeout: 0 }).catch(() => {}); + }); +}); diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index 1a96d3183d..7e75155dc2 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -14,6 +14,81 @@ function createUtilitySql(url: string) { return postgres(url, { max: 1, onnotice: () => {} }); } +type RegisteredPostgresClient = ReturnType; + +/** + * Derives a registry key from a connection URL's host and port only. We must + * not retain or log the full URL, because it carries credentials. + */ +function hostPortKey(url: string): string { + const parsed = new URL(url); + return `${parsed.hostname}:${parsed.port || "5432"}`; +} + +/** + * Same as `hostPortKey`, but returns `null` instead of throwing when the URL + * does not parse. `postgres(url)` tolerates a value `new URL()` rejects (an + * empty string falls back to the `PG*` environment variables), so `createDb` + * must tolerate it too: skip the registry entry and let the driver decide + * the outcome, instead of throwing an error the driver itself would not. + */ +function hostPortKeyOrNull(url: string): string | null { + try { + return hostPortKey(url); + } catch (error) { + if (error instanceof TypeError && (error as NodeJS.ErrnoException).code === "ERR_INVALID_URL") return null; + throw error; + } +} + +// Tracks every client `createDb` hands out, keyed by host and port, so a test +// fixture can end them before it stops the Postgres cluster they point at. A +// `WeakRef` plus `FinalizationRegistry` means a long-lived process (a real +// server) retains nothing extra: an unreferenced client is pruned on its own. +const clientsByHostPort = new Map>>(); +const clientFinalizer = new FinalizationRegistry<{ hostPortKey: string; ref: WeakRef }>( + ({ hostPortKey, ref }) => { + const refs = clientsByHostPort.get(hostPortKey); + if (!refs) return; + refs.delete(ref); + if (refs.size === 0) clientsByHostPort.delete(hostPortKey); + }, +); + +function registerClient(key: string, client: RegisteredPostgresClient): void { + const ref = new WeakRef(client); + let refs = clientsByHostPort.get(key); + if (!refs) { + refs = new Set(); + clientsByHostPort.set(key, refs); + } + refs.add(ref); + clientFinalizer.register(client, { hostPortKey: key, ref }, ref); +} + +/** + * Ends every live client `createDb` handed out for the given URL's host and + * port, then forgets them. Call this before stopping a Postgres cluster: a + * client that outlives the cluster it points at can crash the process (a + * reserved connection's deferred write firing after the socket is gone). + * Swallows individual `end()` errors so one bad client cannot block the rest. + */ +export async function closeRegisteredClients(url: string): Promise { + const key = hostPortKey(url); + const refs = clientsByHostPort.get(key); + if (!refs) return; + + clientsByHostPort.delete(key); + const clients: RegisteredPostgresClient[] = []; + for (const ref of refs) { + clientFinalizer.unregister(ref); + const client = ref.deref(); + if (client) clients.push(client); + } + + await Promise.all(clients.map((client) => client.end({ timeout: 1 }).catch(() => {}))); +} + function isSafeIdentifier(value: string): boolean { return /^[A-Za-z_][A-Za-z0-9_]*$/.test(value); } @@ -118,6 +193,8 @@ export function postgresJsOptions(options: DatabaseClientOptions): Record { + // End every client a caller created against this cluster first. A + // client that still holds a reserved connection when the cluster + // stops can crash the process: the stop kills the backend socket, + // but a queued write on that connection still fires later and finds + // a null socket. + await closeRegisteredClients(connectionString); await stopEmbeddedPostgresBounded(instance, () => cleanupEmbeddedPostgresTestDirs(dataDir)); }, }; diff --git a/server/src/__tests__/helpers/route-test-harness.ts b/server/src/__tests__/helpers/route-test-harness.ts index 6b704caf6a..013a23b945 100644 --- a/server/src/__tests__/helpers/route-test-harness.ts +++ b/server/src/__tests__/helpers/route-test-harness.ts @@ -49,6 +49,11 @@ export function useEmbeddedPostgres( } afterAll(async () => { + // End this suite's client before the cluster stops. `tempDb.cleanup()` + // also ends every client registered against this cluster, so this call + // is redundant defense in depth — it keeps the shutdown order visible + // here, not just inside the fixture. + await db?.$client.end({ timeout: 1 }).catch(() => {}); await tempDb?.cleanup(); tempDb = null; db = null;