From b730ae388623ed187dcac5b09e54911b8a23a69b Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 08:22:53 +0000 Subject: [PATCH 1/2] fix: share one close across concurrent session and client close() calls DBSQLSession.close() marked the session closed only after the CloseSession RPC returned, so concurrent callers each sent CloseSession for the same handle. DBSQLClient.close() had the same race and released the shared telemetry client once per caller. Both now return the in-flight close to concurrent callers. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 1 + lib/DBSQLClient.ts | 12 ++++++++++++ lib/DBSQLSession.ts | 9 +++++++++ tests/unit/DBSQLClient.test.ts | 24 ++++++++++++++++++++++++ tests/unit/DBSQLSession.test.ts | 26 ++++++++++++++++++++++++++ 5 files changed, 72 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 34532b62..a4e4b917 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index 4256873a..00e2b68e 100644 --- a/lib/DBSQLClient.ts +++ b/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; + // 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 { + if (!this.closing) { + this.closing = this.closeOnce().finally(() => { + this.closing = undefined; + }); + } + return this.closing; + } + + private async closeOnce(): Promise { try { await this.sessions.closeAll(); await this.backend?.close(); diff --git a/lib/DBSQLSession.ts b/lib/DBSQLSession.ts index 5bca61d6..a1b20d6c 100644 --- a/lib/DBSQLSession.ts +++ b/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; + public onClose?: () => void; private operations = new CloseableCollection(); @@ -346,7 +349,13 @@ export default class DBSQLSession implements IDBSQLSession { if (!this.isOpen) { return Status.success(); } + if (!this.closing) { + this.closing = this.closeOnce(); + } + return this.closing; + } + private async closeOnce(): Promise { await this.operations.closeAll(); // `latencyMs` on CONNECTION_CLOSE is the closeSession RPC duration, not the diff --git a/tests/unit/DBSQLClient.test.ts b/tests/unit/DBSQLClient.test.ts index 386f268f..a7d594e7 100644 --- a/tests/unit/DBSQLClient.test.ts +++ b/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]', () => { diff --git a/tests/unit/DBSQLSession.test.ts b/tests/unit/DBSQLSession.test.ts index 00498c54..82298c1c 100644 --- a/tests/unit/DBSQLSession.test.ts +++ b/tests/unit/DBSQLSession.test.ts @@ -702,6 +702,32 @@ 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 close operations that belong to it', async () => { const session = createSessionForTest({ handle: sessionHandleStub, context: new ClientContextStub() }); const operation = await session.executeStatement('SELECT * FROM table'); From defd9b9458c9cd91f8c4a9993a1fb48740ae8973 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 08:34:36 +0000 Subject: [PATCH 2/2] fix: allow a session close to be retried after it failed early The shared session close was never cleared, so if closing the session's operations rejected before CloseSession was sent, every later close() returned the same rejection and the session could not be closed. Clear the shared close once it settles, as DBSQLClient.close() already does. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- lib/DBSQLSession.ts | 6 +++++- tests/unit/DBSQLSession.test.ts | 23 +++++++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/lib/DBSQLSession.ts b/lib/DBSQLSession.ts index a1b20d6c..a7c2790a 100644 --- a/lib/DBSQLSession.ts +++ b/lib/DBSQLSession.ts @@ -350,7 +350,11 @@ export default class DBSQLSession implements IDBSQLSession { return Status.success(); } if (!this.closing) { - this.closing = this.closeOnce(); + // 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; } diff --git a/tests/unit/DBSQLSession.test.ts b/tests/unit/DBSQLSession.test.ts index 82298c1c..a667c8b8 100644 --- a/tests/unit/DBSQLSession.test.ts +++ b/tests/unit/DBSQLSession.test.ts @@ -728,6 +728,29 @@ describe('DBSQLSession', () => { 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');