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..cd661e6b29 --- /dev/null +++ b/packages/db/src/client-teardown-registry.test.ts @@ -0,0 +1,93 @@ +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(); + + const order: string[] = []; + + // 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); + order.push("clients-closed"); + + // 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(); + order.push("cluster-stopped"); + + expect(order).toEqual(["clients-closed", "cluster-stopped"]); + + // Let any timer the driver's teardown scheduled run to completion. If a + // deferred write still fired against a null socket, it would surface here + // as an unhandled error and fail this test file. + await new Promise((resolve) => setImmediate(() => setImmediate(resolve))); + + // A query on the reserved connection rejects through the ordinary + // closed-connection path; it does not throw past this call. + const settled = await reserved`select 1`.catch((error: unknown) => error); + expect(settled).toBeInstanceOf(Error); + }); + + 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(); + }); +}); diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index 1a96d3183d..e5f5b4ed22 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -14,6 +14,65 @@ 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"}`; +} + +// 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 +177,7 @@ 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;