sql: a later close() waits for the pending close, for at most its timeout

docs/runtime/sql.mdx

@@ -980,6 +980,8 @@ await sql.close({ timeout: 0 }); // close all connections from the pool immediat
 
 `timeout` is in seconds. With `timeout: 0`, or when the timeout expires, queries that are still in flight reject with the adapter's `ERR_*_CONNECTION_CLOSED` error.
 
+You can call `close()` on the pool again while it is still closing. A later call waits for that close, for at most its `timeout`. When the timeout ends, and at once with `timeout: 0`, the pool rejects the queries that are still in flight and the call resolves.
+
 ---
 
 ## Reserved Connections

packages/bun-types/sql.d.ts

@@ -683,6 +683,10 @@ declare module "bun" {
      * Closes the database connection. With `timeout: 0` it closes
      * immediately; with no timeout it waits for all queries to finish first.
      *
+     * A pool can be closed again while it is still closing. A later call
+     * waits for that close, for at most its `timeout`. When the timeout
+     * ends, the pool rejects the queries that are still in flight.
+     *
      * @param options Optional `timeout` in seconds
      *
      * @example

src/js/internal/sql/shared.ts

@@ -943,6 +943,8 @@ abstract class BaseSQLAdapter<PooledConnection extends BasePooledConnection, Con
   public closed: boolean = false;
   public totalQueries: number = 0;
   public onAllQueriesFinished: (() => void) | null = null;
+  /// Settles when the close() in progress does.
+  #closing: Promise<any> | null = null;
   /// AsyncLocalStorage context the SQL instance was created in. onconnect/onclose run
   /// inside it rather than in whatever context the native callback happens to fire in
   /// (none for a socket event, the close() caller's when the socket closes synchronously).
@@ -1333,11 +1335,37 @@ abstract class BaseSQLAdapter<PooledConnection extends BasePooledConnection, Con
   /** Runs from close() after `closed` is set; overridden by Postgres for its LISTEN connection. */
   protected closeDedicatedConnections(): void {}
 
-  async close(options?: { timeout?: number }): Promise<void> {
-    if (this.closed) {
+  #stopWaitingForQueries() {
+    // #close() clears poolStarted
+    if (this.poolStarted) {
+      this.onAllQueriesFinished?.();
+    }
+  }
+
+  /// close({ timeout }) on a pool that is already closing.
+  #closeWithin(closing: Promise<any>, seconds: number): Promise<void> | undefined {
+    if (seconds === 0) {
+      this.#stopWaitingForQueries();
       return;
     }
+    const { promise, resolve } = Promise.withResolvers<void>();
+    // setTimeout turns a longer delay into 1 ms
+    const delay = Math.min(seconds * 1000, 2 ** 31 - 1);
+    const timer = setTimeout(() => {
+      // timeout is reached, lets close and probably fail some queries
+      this.#stopWaitingForQueries();
+      resolve();
+    }, delay);
+    timer.unref(); // dont block the event loop
+    const settle = () => {
+      clearTimeout(timer);
+      resolve();
+    };
+    closing.$then(settle, settle);
+    return promise;
+  }
 
+  async close(options?: { timeout?: number }): Promise<void> {
     let timeout = options?.timeout;
     // Presence, not truthiness: `timeout: 0` means close now, undefined/null mean drain with no timer.
     const hasTimeout = timeout != null;
@@ -1348,17 +1376,23 @@ abstract class BaseSQLAdapter<PooledConnection extends BasePooledConnection, Con
       }
     }
 
+    if (this.closed) {
+      const closing = this.#closing;
+      return hasTimeout && closing !== null ? this.#closeWithin(closing, timeout!) : closing;
+    }
+
     this.closed = true;
     this.closeDedicatedConnections();
 
     if (hasTimeout) {
       if (timeout === 0 || !this.hasPendingQueries()) {
         // close immediately
-        await this.#close();
+        await (this.#closing = this.#close());
         return;
       }
 
       const { promise, resolve } = Promise.withResolvers<void>();
+      this.#closing = promise;
       const timer = setTimeout(() => {
         // timeout is reached, lets close and probably fail some queries
         this.#close().finally(resolve);
@@ -1375,12 +1409,13 @@ abstract class BaseSQLAdapter<PooledConnection extends BasePooledConnection, Con
     } else {
       if (!this.hasPendingQueries()) {
         // close immediately
-        await this.#close();
+        await (this.#closing = this.#close());
         return;
       }
 
       // gracefully close the pool
       const { promise, resolve } = Promise.withResolvers<void>();
+      this.#closing = promise;
 
       this.onAllQueriesFinished = () => {
         // everything is closed, lets close the pool

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

@@ -18,6 +18,7 @@
 
 import { SQL } from "bun";
 import { expect, mock, test } from "bun:test";
+import { bunEnv, bunExe } from "harness";
 import type { Server, Socket } from "node:net";
 import {
   listeningServer,
@@ -296,3 +297,194 @@ for (const [name, scheme, closedCode] of drivers) {
     });
   }
 }
+
+// close() on a pool that an earlier close() call is still closing. A later call waits for that close, for at
+// most its timeout. When the timeout ends, the pool stops waiting for the queries in flight.
+type CloseOptions = { timeout?: number } | undefined;
+const show = (options: CloseOptions) => `close(${options ? `{ timeout: ${options.timeout} }` : ""})`;
+
+function heldQuery(sql: SQL) {
+  return sql`select 1 as x`.simple().then(
+    rows => ({ rows: [...rows] }),
+    e => ({ code: e.code }),
+  );
+}
+
+for (const [name, scheme, closedCode] of drivers) {
+  const rows = { rows: [{ x: "1" }] };
+
+  const forcing: [first: CloseOptions, later: CloseOptions][] = [
+    [undefined, { timeout: 0 }],
+    [{ timeout: 30 }, { timeout: 0 }],
+    [undefined, { timeout: 0.01 }],
+    [{ timeout: 30 }, { timeout: 0.01 }],
+  ];
+  for (const [first, later] of forcing) {
+    test(`${name}: ${show(later)} closes a pool that waits in ${show(first)}`, async () => {
+      const { port, server, commandReceived, respond } = await heldQueryMocks[name]();
+      try {
+        const sql = new SQL({ url: `${scheme}127.0.0.1:${port}/db`, max: 1 });
+        const query = heldQuery(sql);
+        await commandReceived;
+        const firstClosed = sql.close(first);
+        await sql.close(later);
+        // a later call that closed nothing leaves the query in flight, and this answer resolves it
+        respond();
+        expect(await query).toEqual({ code: closedCode });
+        await firstClosed;
+      } finally {
+        server.close();
+      }
+    });
+  }
+
+  const joining: [first: CloseOptions, later: CloseOptions][] = [
+    [undefined, undefined],
+    [{ timeout: 30 }, undefined],
+    [undefined, { timeout: 30 }],
+    // more milliseconds than a timer can hold
+    [undefined, { timeout: 3_000_000 }],
+  ];
+  for (const [first, later] of joining) {
+    test(`${name}: ${show(later)} settles with a pool that waits in ${show(first)}`, async () => {
+      const { port, server, commandReceived, respond } = await heldQueryMocks[name]();
+      try {
+        const sql = new SQL({ url: `${scheme}127.0.0.1:${port}/db`, max: 1 });
+        const order: string[] = [];
+        const query = heldQuery(sql).finally(() => order.push("query"));
+        await commandReceived;
+        const firstClosed = sql.close(first).then(() => order.push("first"));
+        const laterClosed = sql.close(later).then(() => order.push("later"));
+        respond();
+        expect(await query).toEqual(rows);
+        await Promise.all([firstClosed, laterClosed]);
+        expect(order).toEqual(["query", "first", "later"]);
+      } finally {
+        server.close();
+      }
+    });
+  }
+
+  test(`${name}: close({ timeout }) closes the pool when the timeout expires, also after a later close()`, async () => {
+    const { port, server, commandReceived } = await heldQueryMocks[name]();
+    try {
+      const sql = new SQL({ url: `${scheme}127.0.0.1:${port}/db`, max: 1 });
+      const query = heldQuery(sql);
+      await commandReceived;
+      // the server never answers, so only the timeout can settle these
+      await Promise.all([sql.close({ timeout: 0.01 }), sql.close(), sql.end({ timeout: 30 })]);
+      expect(await query).toEqual({ code: closedCode });
+      expect(await Promise.all([sql.close(), sql.end({ timeout: 0 })])).toEqual([undefined, undefined]);
+    } finally {
+      server.close();
+    }
+  });
+
+  test(`${name}: a later close() with an invalid timeout rejects and changes nothing`, async () => {
+    const { port, server, commandReceived, respond } = await heldQueryMocks[name]();
+    try {
+      const sql = new SQL({ url: `${scheme}127.0.0.1:${port}/db`, max: 1 });
+      const query = heldQuery(sql);
+      await commandReceived;
+      const firstClosed = sql.close();
+      const codes = [-1, 2 ** 32, NaN].map(timeout => sql.close({ timeout }).catch(e => e.code));
+      expect(await Promise.all(codes)).toEqual(Array(3).fill("ERR_INVALID_ARG_VALUE"));
+      respond();
+      expect(await query).toEqual(rows);
+      await firstClosed;
+      expect(await sql.close({ timeout: -1 }).catch(e => e.code)).toBe("ERR_INVALID_ARG_VALUE");
+    } finally {
+      server.close();
+    }
+  });
+
+  test(`${name}: end() and asyncDispose reach the same close as close()`, async () => {
+    const { port, server, commandReceived, respond } = await heldQueryMocks[name]();
+    try {
+      const sql = new SQL({ url: `${scheme}127.0.0.1:${port}/db`, max: 1 });
+      const order: string[] = [];
+      const query = heldQuery(sql).finally(() => order.push("query"));
+      await commandReceived;
+      const disposed = sql[Symbol.asyncDispose]().then(() => order.push("asyncDispose"));
+      const joined = sql.end().then(() => order.push("end"));
+      await sql.end({ timeout: 0 });
+      respond();
+      expect(await query).toEqual({ code: closedCode });
+      await Promise.all([disposed, joined]);
+      expect(order).toEqual(["query", "asyncDispose", "end"]);
+    } finally {
+      server.close();
+    }
+  });
+
+  test(`${name}: only a later close() without a timeout waits for a forced close that still closes a slot`, async () => {
+    const { port, server, commandReceived } = await heldQueryMocks[name]();
+    const password = Promise.withResolvers<string>();
+    try {
+      let slots = 0;
+      const order: string[] = [];
+      const sql = new SQL({
+        url: `${scheme}127.0.0.1:${port}/db`,
+        max: 2,
+        // the second slot has no native connection until the test resolves its password
+        password: () => (++slots === 1 ? "" : password.promise),
+      });
+      const query = heldQuery(sql);
+      await commandReceived;
+      const forced = sql.close({ timeout: 0 }).then(() => order.push("forced"));
+      const later = sql.close().then(() => order.push("later"));
+      const laterForced = sql.close({ timeout: 0 }).then(() => order.push("later forced"));
+      expect(await query).toEqual({ code: closedCode });
+      // one turn of the event loop: a later close() that did not wait with the first one has settled by now
+      await new Promise(resolve => setImmediate(resolve));
+      order.push("password");
+      password.resolve("");
+      await Promise.all([forced, later, laterForced]);
+      expect(order).toEqual(["later forced", "password", "forced", "later"]);
+    } finally {
+      password.resolve("");
+      server.close();
+    }
+  });
+}
+
+test("Bun.sql.close() and Bun.sql.end() reach the same close", async () => {
+  const { port, server, commandReceived, respond } = await heldQueryMocks.postgres();
+  try {
+    await using proc = Bun.spawn({
+      cmd: [
+        bunExe(),
+        "-e",
+        `
+          const { sql } = Bun;
+          const query = sql\`select 1 as x\`.simple().then(() => "rows", e => e.code);
+          process.stdin.once("data", async () => {
+            const first = sql.close();
+            await sql.end({ timeout: 0 });
+            console.log("forced");
+            console.log("query: " + (await query));
+            await first;
+            console.log("closed");
+          });
+        `,
+      ],
+      env: { ...bunEnv, DATABASE_URL: `postgres://postgres@127.0.0.1:${port}/db` },
+      stdin: "pipe",
+      stdout: "pipe",
+      stderr: "inherit",
+    });
+    await commandReceived;
+    proc.stdin.write("close\n");
+    proc.stdin.end();
+    let stdout = "";
+    for await (const chunk of proc.stdout) {
+      stdout += Buffer.from(chunk).toString();
+      // an end() that closed nothing leaves the query in flight, and this answer resolves it
+      if (stdout.includes("forced\n")) respond();
+    }
+    expect(stdout).toBe("forced\nquery: ERR_POSTGRES_CONNECTION_CLOSED\nclosed\n");
+    expect(await proc.exited).toBe(0);
+  } finally {
+    server.close();
+  }
+});

test/js/sql/sql-mysql.test.ts

@@ -1017,6 +1017,26 @@ if (isDockerEnabled()) {
           return expect(await error).toBe("ERR_MYSQL_CONNECTION_CLOSED");
         });
 
+        test("Connection end again waits for the pending end", async () => {
+          const sql = new SQL(getOptions());
+          const order: string[] = [];
+          const promise = sql`select SLEEP(0.2) as x`.execute().finally(() => order.push("query"));
+          const first = sql.end().then(() => order.push("first end"));
+          const later = sql.end().then(() => order.push("later end"));
+          expect(await promise).toEqual([{ x: 0 }]);
+          await Promise.all([first, later]);
+          expect(order).toEqual(["query", "first end", "later end"]);
+        });
+
+        test("Connection destroyed after end", async () => {
+          const sql = new SQL(getOptions());
+          const error = sql`select SLEEP(2)`.execute().catch(err => err.code);
+          const ended = sql.end();
+          await sql.end({ timeout: 0 });
+          expect(await error).toBe("ERR_MYSQL_CONNECTION_CLOSED");
+          await ended;
+        });
+
         test("unsafe", async () => {
           await using sql = new SQL({ ...getOptions(), max: 1 });
           await sql`create temporary table test_unsafe (x int)`;

test/js/sql/sql.test.ts

@@ -1845,6 +1845,26 @@ if (isDockerEnabled()) {
       return expect(await error).toBe("ERR_POSTGRES_CONNECTION_CLOSED");
     });
 
+    test("Connection end again waits for the pending end", async () => {
+      const sql = postgres(options);
+      const order: string[] = [];
+      const promise = sql`select pg_sleep(0.2) as x`.execute().finally(() => order.push("query"));
+      const first = sql.end().then(() => order.push("first end"));
+      const later = sql.end().then(() => order.push("later end"));
+      expect(await promise).toEqual([{ x: "" }]);
+      await Promise.all([first, later]);
+      expect(order).toEqual(["query", "first end", "later end"]);
+    });
+
+    test("Connection destroyed after end", async () => {
+      const sql = postgres(options);
+      const error = sql`select pg_sleep(2)`.execute().catch(err => err.code);
+      const ended = sql.end();
+      await sql.end({ timeout: 0 });
+      expect(await error).toBe("ERR_POSTGRES_CONNECTION_CLOSED");
+      await ended;
+    });
+
     // t('transform column', async() => {
     //   const sql = postgres({
     //     ...options,