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..a7c2790a 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,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 { 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..a667c8b8 100644 --- a/tests/unit/DBSQLSession.test.ts +++ b/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');