feat: add Symbol.dispose and Symbol.asyncDispose support for Connections, Pools, and Pool Clusters (#4112)

* feat: support `using`, `dispose`, and `asyncDispose`

* chore: implement types

* ci: add tests

* chore: implement `dispose` for Pool Cluster

* refactor: improve test organization

* refactor: force `using` keyword

* docs: add docs for all types of connections (including how to end them)

Author: wellwelwel

lib/base/connection.js

@@ -926,6 +926,12 @@ class BaseConnection extends EventEmitter {
     return this.addCommand(new Commands.ServerHandshake(args));
   }
 
+  [Symbol.dispose]() {
+    if (!this._closing) {
+      this.end();
+    }
+  }
+
   // ===============================================================
   end(callback) {
     if (this.config.isServer) {

lib/base/pool.js

@@ -108,6 +108,12 @@ class BasePool extends EventEmitter {
     }
   }
 
+  [Symbol.dispose]() {
+    if (!this._closed) {
+      this.end();
+    }
+  }
+
   end(cb) {
     this._closed = true;
     clearTimeout(this._removeIdleTimeoutConnectionsTimer);

lib/base/pool_connection.js

@@ -29,6 +29,10 @@ class BasePoolConnection extends BaseConnection {
     this._pool.releaseConnection(this);
   }
 
+  [Symbol.dispose]() {
+    this.release();
+  }
+
   end(callback) {
     if (this.config.gracefulEnd) {
       this._removeFromPool();

lib/pool_cluster.js

@@ -229,6 +229,12 @@ class PoolCluster extends EventEmitter {
     namespace.getConnection(cb);
   }
 
+  [Symbol.dispose]() {
+    if (!this._closed) {
+      this.end();
+    }
+  }
+
   end(callback) {
     const cb =
       callback !== undefined

lib/promise/connection.js

@@ -66,6 +66,12 @@ class PromiseConnection extends EventEmitter {
     });
   }
 
+  async [Symbol.asyncDispose]() {
+    if (!this.connection._closing) {
+      await this.end();
+    }
+  }
+
   beginTransaction() {
     const c = this.connection;
     const localErr = new Error();

lib/promise/pool.js

@@ -85,6 +85,12 @@ class PromisePool extends EventEmitter {
       });
     });
   }
+
+  async [Symbol.asyncDispose]() {
+    if (!this.pool._closed) {
+      await this.end();
+    }
+  }
 }
 
 (function (functionsToWrap) {

lib/promise/pool_connection.js

@@ -14,6 +14,10 @@ class PromisePoolConnection extends PromiseConnection {
       arguments
     );
   }
+
+  async [Symbol.asyncDispose]() {
+    this.release();
+  }
 }
 
 module.exports = PromisePoolConnection;

promise.d.ts

@@ -65,6 +65,8 @@ export interface Connection extends QueryableAndExecutableBase {
 
   end(options?: any): Promise<void>;
 
+  [Symbol.asyncDispose](): Promise<void>;
+
   destroy(): void;
 
   pause(): void;
@@ -82,6 +84,7 @@ export interface Connection extends QueryableAndExecutableBase {
 export interface PoolConnection extends Connection {
   release(): void;
   connection: Connection;
+  [Symbol.asyncDispose](): Promise<void>;
 }
 
 export interface Pool extends Connection {
@@ -112,6 +115,8 @@ export interface PoolCluster extends EventEmitter {
 
   end(): Promise<void>;
 
+  [Symbol.asyncDispose](): Promise<void>;
+
   getConnection(): Promise<PoolConnection>;
   getConnection(group: string): Promise<PoolConnection>;
   getConnection(group: string, selector: string): Promise<PoolConnection>;

promise.js

@@ -134,6 +134,12 @@ class PromisePoolCluster extends EventEmitter {
       });
     });
   }
+
+  async [Symbol.asyncDispose]() {
+    if (!this.poolCluster._closed) {
+      await this.end();
+    }
+  }
 }
 
 /**

test/integration/dispose/async/connection.test.mts

@@ -0,0 +1,49 @@
+import type { RowDataPacket } from '../../../../index.js';
+import { assert, describe, it, skip } from 'poku';
+import { createConnection } from '../../../../promise.js';
+import { config } from '../../../common.test.mjs';
+
+if (!('asyncDispose' in Symbol)) {
+  skip('Symbol.asyncDispose is not supported in this runtime');
+}
+
+await describe('PromiseConnection should implement Symbol.asyncDispose', async () => {
+  await using conn = await createConnection(config);
+
+  it('should be a function', () => {
+    assert.strictEqual(typeof conn[Symbol.asyncDispose], 'function');
+  });
+});
+
+await describe('asyncDispose should end the connection', async () => {
+  await using conn = await createConnection(config);
+  const [rows] = await conn.query<RowDataPacket[]>('SELECT 1');
+
+  await conn[Symbol.asyncDispose]();
+
+  it('should have received the query result', () => {
+    assert.deepStrictEqual(rows, [{ 1: 1 }]);
+  });
+
+  it('should have closed the connection', () => {
+    // @ts-expect-error: internal access
+    assert.strictEqual(conn.connection._closing, true);
+  });
+});
+
+await describe('asyncDispose should handle manual destroy before asyncDispose on connection', async () => {
+  await using conn = await createConnection(config);
+  const [rows] = await conn.query<RowDataPacket[]>('SELECT 1');
+
+  conn.destroy();
+  await conn[Symbol.asyncDispose]();
+
+  it('should have received the query result', () => {
+    assert.deepStrictEqual(rows, [{ 1: 1 }]);
+  });
+
+  it('should have closed the connection', () => {
+    // @ts-expect-error: internal access
+    assert.strictEqual(conn.connection._closing, true);
+  });
+});