fix: share one close across concurrent session and client close() calls

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Concurrent `close()` calls on a session or client now share one close. Previously each call sent its own Thrift `CloseSession` for the same session, and concurrent `DBSQLClient.close()` calls released the shared telemetry client once per call, which could close it for other clients on the same host.
 - Thrift CloudFetch downloads are now bounded end to end by `retriesTimeout` (default 15 minutes), including every attempt and retry backoff. Previously a stalled download, where the response headers or body never arrived, blocked the fetch forever. A download still running when the deadline expires is aborted and fails with `RetryError` (`TIMEOUT_EXCEEDED`, message `CloudFetch result download failed: Retry timeout exceeded`). (#520)
 - Thrift connections made through `proxy` now apply `checkServerCertificate`, `customCaCert` and `clientCert`/`clientKey`. Previously the proxied path ignored them and verified against Node's process-wide trust store, so a certificate outside the configured trust set could be accepted. **Compatibility:** proxied connections now verify unless `checkServerCertificate: false` is set; as on direct connections, `NODE_TLS_REJECT_UNAUTHORIZED=0` no longer turns verification off. For an `https` proxy these options also apply to the proxy's own certificate. (#485)
 - Honor the client `useCloudFetch: false` default on the kernel backend so complete results can be fetched without cloud-storage downloads. (#574)

lib/DBSQLClient.ts

@@ -125,6 +125,9 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
 
   private telemetryEmitter?: TelemetryEventEmitter;
 
+  // Shared by concurrent close() calls so sessions and telemetry are released once
+  private closing?: Promise<void>;
+
   // True once we've shipped the full DriverConfiguration on a CONNECTION_OPEN
   // event for this client. Subsequent openSession events for the same client
   // strip the (~1KB, static-for-the-process) blob — a long-running client
@@ -933,6 +936,15 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
    * method on shutdown so the aggregator drains its remaining metrics.
    */
   public async close(): Promise<void> {
+    if (!this.closing) {
+      this.closing = this.closeOnce().finally(() => {
+        this.closing = undefined;
+      });
+    }
+    return this.closing;
+  }
+
+  private async closeOnce(): Promise<void> {
     try {
       await this.sessions.closeAll();
       await this.backend?.close();

lib/DBSQLSession.ts

@@ -59,6 +59,9 @@ export default class DBSQLSession implements IDBSQLSession {
 
   private isOpen = true;
 
+  // Shared by concurrent close() calls so the session is closed with one request
+  private closing?: Promise<Status>;
+
   public onClose?: () => void;
 
   private operations = new CloseableCollection<DBSQLOperation>();
@@ -346,7 +349,17 @@ export default class DBSQLSession implements IDBSQLSession {
     if (!this.isOpen) {
       return Status.success();
     }
+    if (!this.closing) {
+      // Cleared once settled so a close that failed before CloseSession (e.g. an
+      // operation failed to close) can be retried
+      this.closing = this.closeOnce().finally(() => {
+        this.closing = undefined;
+      });
+    }
+    return this.closing;
+  }
 
+  private async closeOnce(): Promise<Status> {
     await this.operations.closeAll();
 
     // `latencyMs` on CONNECTION_CLOSE is the closeSession RPC duration, not the

tests/unit/DBSQLClient.test.ts

@@ -628,6 +628,30 @@ describe('DBSQLClient.close', () => {
     expect(client['sessions']['items'].size).to.eq(0);
     expect(thriftClient.CloseSession.called).to.be.true;
   });
+
+  it('should close sessions and release telemetry once when closed concurrently', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = sinon.spy(new ThriftClientStub());
+
+    client['client'] = thriftClient;
+    client['connectionProvider'] = new ConnectionProviderStub();
+    client['authProvider'] = new AuthProviderStub();
+    client['backend'] = new ThriftBackend({ context: client, onConnectionEvent: () => {} });
+    client['host'] = 'test.cloud.databricks.com';
+    client['telemetryClient'] = {} as any;
+    await client.openSession();
+    const backendClose = sinon.spy(client['backend'], 'close');
+    const releaseClient = sinon.stub(TelemetryClientProvider.getInstance(), 'releaseClient').resolves();
+
+    try {
+      await Promise.all([client.close(), client.close(), client.close()]);
+      expect(thriftClient.CloseSession.calledOnce).to.be.true;
+      expect(backendClose.calledOnce).to.be.true;
+      expect(releaseClient.calledOnce).to.be.true;
+    } finally {
+      releaseClient.restore();
+    }
+  });
 });
 
 describe('DBSQLClient[Symbol.asyncDispose]', () => {

tests/unit/DBSQLSession.test.ts

@@ -702,6 +702,55 @@ describe('DBSQLSession', () => {
       expect(driver.closeSession.callCount).to.eq(1); // second time it should not be called
     });
 
+    it('should send one close request for concurrent close calls', async () => {
+      const context = new ClientContextStub();
+      const driver = sinon.spy(context.driver);
+      const session = createSessionForTest({ handle: sessionHandleStub, context });
+
+      const results = await Promise.all([session.close(), session.close(), session.close()]);
+      results.forEach((result) => expect(result).instanceOf(Status));
+      expect(session['isOpen']).to.be.false;
+      expect(driver.closeSession.callCount).to.eq(1);
+    });
+
+    it('should reject every concurrent close call when the close request fails', async () => {
+      const context = new ClientContextStub();
+      const error = new Error('close failed');
+      const closeSession = sinon.stub(context.driver, 'closeSession').rejects(error);
+      const session = createSessionForTest({ handle: sessionHandleStub, context });
+
+      const results = await Promise.allSettled([session.close(), session.close()]);
+      expect(results).to.deep.equal([
+        { status: 'rejected', reason: error },
+        { status: 'rejected', reason: error },
+      ]);
+      expect(closeSession.callCount).to.eq(1);
+      expect((await session.close()).isSuccess).to.be.true;
+    });
+
+    it('should retry close after closing its operations failed', async () => {
+      const context = new ClientContextStub();
+      const driver = sinon.spy(context.driver);
+      const session = createSessionForTest({ handle: sessionHandleStub, context });
+      const error = new Error('operation close failed');
+      const closeAll = sinon.stub(session['operations'], 'closeAll');
+      closeAll.onFirstCall().rejects(error);
+      closeAll.callThrough();
+
+      try {
+        await session.close();
+        expect.fail('It should throw an error');
+      } catch (e) {
+        expect(e).to.equal(error);
+      }
+      expect(session['isOpen']).to.be.true;
+      expect(driver.closeSession.callCount).to.eq(0);
+
+      expect(await session.close()).instanceOf(Status);
+      expect(session['isOpen']).to.be.false;
+      expect(driver.closeSession.callCount).to.eq(1);
+    });
+
     it('should close operations that belong to it', async () => {
       const session = createSessionForTest({ handle: sessionHandleStub, context: new ClientContextStub() });
       const operation = await session.executeStatement('SELECT * FROM table');