sql: fire onclose only for pool connections that completed their handshake

src/js/internal/sql/shared.ts

@@ -591,6 +591,8 @@ const enum PooledConnectionFlags {
   reserved = 1 << 1,
   /// preReserved is used to indicate that the connection will be reserved in the future when queryCount drops to 0
   preReserved = 1 << 2,
+  /// onConnectFired is used to indicate that handleConnected ran for this slot, so the user's onconnect callback already fired (with null or an error)
+  onConnectFired = 1 << 3,
 }
 export type { PooledConnectionState };
 
@@ -657,6 +659,7 @@ abstract class BasePooledConnection<ConnectionHandle extends { close(): void; fl
     if (err) {
       err = this.wrapError(err);
     }
+    this.flags |= PooledConnectionFlags.onConnectFired;
     const connectionInfo = this.connectionInfo;
     try {
       // user code; a throw must not abort the pool bookkeeping below
@@ -759,9 +762,11 @@ abstract class BasePooledConnection<ConnectionHandle extends { close(): void; fl
 
   #finishClose(err: any) {
     const connectionInfo = this.connectionInfo;
+    const poolClosedSlotBeforeOnconnect =
+      this.onFinish !== null && !(this.flags & PooledConnectionFlags.onConnectFired);
     try {
       // user code; a throw must not abort the pool bookkeeping below
-      if (connectionInfo?.onclose) {
+      if (!poolClosedSlotBeforeOnconnect && connectionInfo?.onclose) {
         AsyncContextFrame.run(this.adapter.callbackAsyncContext, connectionInfo.onclose, connectionInfo, err);
       }
     } finally {
@@ -800,10 +805,11 @@ abstract class BasePooledConnection<ConnectionHandle extends { close(): void; fl
     if (this.adapter.closed) {
       return;
     }
-    // reset error and state
+    // reset error and state; the new cycle has not fired onconnect yet
     this.storedError = null;
     this.connectStartedAt = 0;
     this.state = PooledConnectionState.pending;
+    this.flags &= ~PooledConnectionFlags.onConnectFired;
     // retry connection
     this.#beginConnecting();
   }

test/js/sql/sql-close-pending-connection.test.ts

@@ -17,8 +17,8 @@
 // thing that can tear the connection down — without the fix these tests hang.
 
 import { SQL } from "bun";
-import { expect, test } from "bun:test";
-import { neverAnsweringServer } from "./wire-frames";
+import { expect, mock, test } from "bun:test";
+import { listeningServer, neverAnsweringServer, pgAuthenticationOk, pgReadyForQuery } from "./wire-frames";
 
 const drivers = [
   ["postgres", "postgres://postgres@", "ERR_POSTGRES_CONNECTION_CLOSED"],
@@ -41,6 +41,35 @@ for (const [name, scheme, closedCode] of drivers) {
     }
   });
 
+  // https://github.com/oven-sh/bun/issues/39940
+  //
+  // close() used to fire the user's onclose callback once per pool slot in
+  // the pending state, even when that slot's handshake never completed and
+  // onconnect never fired, so onconnect/onclose pairing drifted by up to
+  // `max` per pool close.
+  test(`${name}: close() does not fire onclose for slots that never connected`, async () => {
+    const { port, server, accepted } = await neverAnsweringServer();
+    try {
+      const onconnect = mock();
+      const onclose = mock();
+      const sql = new SQL({
+        url: `${scheme}127.0.0.1:${port}/db`,
+        max: 5,
+        connectionTimeout: 0,
+        onconnect,
+        onclose,
+      });
+      const queryError = sql`SELECT 1`.catch(e => e);
+      await accepted;
+      await sql.close({ timeout: "0" });
+      expect((await queryError).code).toBe(closedCode);
+      expect(onconnect).not.toHaveBeenCalled();
+      expect(onclose).not.toHaveBeenCalled();
+    } finally {
+      server.close();
+    }
+  });
+
   test(`${name}: forced close() resolves when called before the native handle is stored`, async () => {
     const { port, server } = await neverAnsweringServer();
     try {
@@ -56,6 +85,55 @@ for (const [name, scheme, closedCode] of drivers) {
   });
 }
 
+// https://github.com/oven-sh/bun/issues/39940
+//
+// The per-slot "fired onconnect" marker is per connect cycle. A slot that
+// connected once, closed, and is now redialing must not reuse the marker from
+// the previous cycle: a forced close() that lands mid-reconnect used to fire
+// a second onclose for a cycle whose onconnect never fired.
+test("postgres: close() mid-reconnect does not fire onclose for the unfinished cycle", async () => {
+  const firstClose = Promise.withResolvers<void>();
+  const onconnect = mock();
+  const onclose = mock(() => firstClose.resolve());
+  const secondAccepted = Promise.withResolvers<void>();
+  let firstSocket: import("node:net").Socket;
+  let connections = 0;
+  const { port, server } = await listeningServer(socket => {
+    if (++connections === 1) {
+      firstSocket = socket;
+      // complete the handshake so the slot fires onconnect
+      socket.once("data", () => socket.write(Buffer.concat([pgAuthenticationOk(), pgReadyForQuery()])));
+    } else {
+      // the reconnect stays mid-handshake
+      secondAccepted.resolve();
+    }
+  });
+  try {
+    const sql = new SQL({
+      url: `postgres://postgres@127.0.0.1:${port}/postgres`,
+      max: 1,
+      connectionTimeout: 0,
+      onconnect,
+      onclose,
+    });
+    await sql.connect();
+    expect(onconnect).toHaveBeenCalledTimes(1);
+    // drop the connection from the server side; onclose pairs with onconnect
+    firstSocket!.destroy();
+    await firstClose.promise;
+    expect(onclose).toHaveBeenCalledTimes(1);
+    // a new query redials the closed slot, then close() lands mid-handshake
+    const queryError = sql`SELECT 1`.catch(e => e);
+    await secondAccepted.promise;
+    await sql.close({ timeout: "0" });
+    expect((await queryError).code).toBe("ERR_POSTGRES_CONNECTION_CLOSED");
+    expect(onconnect).toHaveBeenCalledTimes(1);
+    expect(onclose).toHaveBeenCalledTimes(1);
+  } finally {
+    server.close();
+  }
+});
+
 // https://github.com/oven-sh/bun/issues/32198
 //
 // The pool's connection array is allocated as `new Array(max)` and filled one

test/js/sql/sql-onconnect-onclose-throw.test.ts

@@ -236,14 +236,17 @@ process.exit(0);
 // Buffer.alloc frame construction here.
 //
 // The forced-close path (#32095) and the throwing-callback path (#32037) meet
-// in the pool connection's close handler: the user's onclose runs first and
-// may throw, and the bookkeeping that follows it must still settle the
-// promise returned by close(). A server that accepts the TCP connection but
-// never answers keeps the connection mid-handshake, and connectionTimeout: 0
-// disables the connect timer, so close() is the only teardown path; if the
-// throw skipped the bookkeeping these fixtures would never print "closed".
-// The mock server lives in the fixture process (it must observe `accepted`
-// before forcing close) and is imported from ./wire-frames by absolute path.
+// in the pool connection's close handler: the bookkeeping that settles the
+// promise returned by close() must run even when a user callback throws. A
+// server that accepts the TCP connection but never answers keeps the
+// connection mid-handshake, and connectionTimeout: 0 disables the connect
+// timer, so close() is the only teardown path; if the bookkeeping were
+// skipped these fixtures would never print "closed". Since #39940 a slot
+// that never completed its handshake fired no onconnect, so the forced close
+// skips its onclose too; the throwing onclose stays installed to pin that it
+// is not invoked. The mock server lives in the fixture process (it must
+// observe `accepted` before forcing close) and is imported from ./wire-frames
+// by absolute path.
 function forcedCloseFixture(adapter: "postgres" | "mysql") {
   const url = adapter === "postgres" ? "postgres://postgres@127.0.0.1:" : "mysql://root@127.0.0.1:";
   const db = adapter === "postgres" ? "/postgres" : "/db";
@@ -275,12 +278,10 @@ for (const [adapter, closedCode] of [
   ["mysql", "ERR_MYSQL_CONNECTION_CLOSED"],
 ] as const) {
   test.concurrent(
-    `${adapter}: a throwing onclose does not prevent forced close() from resolving mid-handshake`,
+    `${adapter}: forced close() mid-handshake resolves and skips onclose for the never-connected slot`,
     async () => {
       const { stdout, exitCode } = await runFixture(forcedCloseFixture(adapter));
-      expect(stdout).toBe(
-        `onclose: ${closedCode}\nuncaught: boom from onclose\nclosed\nquery rejected: ${closedCode}\n`,
-      );
+      expect(stdout).toBe(`closed\nquery rejected: ${closedCode}\n`);
       expect(exitCode).toBe(0);
     },
   );

test/js/sql/sql.test.ts

@@ -862,6 +862,19 @@ if (isDockerEnabled()) {
       expect(error.code).toBe(`ERR_POSTGRES_LIFETIME_TIMEOUT`);
     });
 
+    // https://github.com/oven-sh/bun/issues/39940
+    // close() used to fire onclose once per pool slot, even for slots whose
+    // handshake never completed, so onconnect/onclose pairing drifted.
+    test("close() fires onclose only for connections that fired onconnect", async () => {
+      const onconnect = mock();
+      const onclose = mock();
+      const sql = postgres({ ...options, max: 10, onconnect, onclose });
+      await sql`select 1`;
+      await sql.close();
+      expect(onconnect).toHaveBeenCalled();
+      expect(onclose).toHaveBeenCalledTimes(onconnect.mock.calls.length);
+    });
+
     // Last one wins.
     test("Handles duplicate string column names", async () => {
       const result = await sql`select 1 as x, 2 as x, 3 as x`;