fix(db): close test database clients before the embedded cluster stops

An embedded-Postgres test fixture stopped its cluster while test clients
still held open connections. The postgres.js driver could schedule a
deferred write on a connection whose backend had just died, and that
write threw past a bare timer callback and crashed the process. CI saw
this as a green shard that still exited 1 with an unhandled error.

Track every client the database package hands out, keyed by host and
port. The embedded-Postgres test fixture now ends every live client for
its cluster before it stops that cluster, so no client can outlive the
server it points at.

Co-authored-by: Paperclip <noreply@paperclip.ing>

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();
+  });
+});

packages/db/src/client.ts

@@ -14,6 +14,65 @@ function createUtilitySql(url: string) {
   return postgres(url, { max: 1, onnotice: () => {} });
 }
 
+type RegisteredPostgresClient = ReturnType<typeof postgres>;
+
+/**
+ * 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<string, Set<WeakRef<RegisteredPostgresClient>>>();
+const clientFinalizer = new FinalizationRegistry<{ hostPortKey: string; ref: WeakRef<RegisteredPostgresClient> }>(
+  ({ 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<void> {
+  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<string
 export function createDb(url: string, options?: DatabaseClientOptions) {
   const resolved = options ?? databaseClientOptionsFromEnv();
   const sql = postgres(url, postgresJsOptions(resolved));
+  registerClient(hostPortKey(url), sql);
   return drizzlePg(sql, { schema });
 }
 

packages/db/src/index.ts

@@ -1,5 +1,6 @@
 export {
   createDb,
+  closeRegisteredClients,
   getPostgresDataDirectory,
   ensurePostgresDatabase,
   resetPostgresDatabase,

packages/db/src/test-embedded-postgres.ts

@@ -2,7 +2,7 @@ import fs from "node:fs";
 import net from "node:net";
 import os from "node:os";
 import path from "node:path";
-import { applyPendingMigrations, ensurePostgresDatabase } from "./client.js";
+import { applyPendingMigrations, closeRegisteredClients, ensurePostgresDatabase } from "./client.js";
 import {
   createEmbeddedPostgresLogBuffer,
   formatEmbeddedPostgresError,
@@ -279,6 +279,12 @@ export async function startEmbeddedPostgresTestDatabase(
     return {
       connectionString,
       cleanup: async () => {
+        // 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));
       },
     };

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;