diff --git a/.prettierignore b/.prettierignore index 92e9533a..9a9ec6bc 100644 --- a/.prettierignore +++ b/.prettierignore @@ -5,6 +5,7 @@ node_modules .nyc_output coverage_e2e coverage_unit +coverage .clinic dist diff --git a/lib/telemetry/TelemetryClient.ts b/lib/telemetry/TelemetryClient.ts new file mode 100644 index 00000000..d74d1417 --- /dev/null +++ b/lib/telemetry/TelemetryClient.ts @@ -0,0 +1,65 @@ +/** + * Copyright (c) 2025 Databricks Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import IClientContext from '../contracts/IClientContext'; +import { LogLevel } from '../contracts/IDBSQLLogger'; + +/** + * Telemetry client for a specific host. + * Managed by TelemetryClientProvider with reference counting. + * One client instance is shared across all connections to the same host. + */ +class TelemetryClient { + private closed: boolean = false; + + constructor(private context: IClientContext, private host: string) { + const logger = context.getLogger(); + logger.log(LogLevel.debug, `Created TelemetryClient for host: ${host}`); + } + + /** + * Gets the host associated with this client. + */ + getHost(): string { + return this.host; + } + + /** + * Checks if the client has been closed. + */ + isClosed(): boolean { + return this.closed; + } + + /** + * Closes the telemetry client and releases resources. + * Should only be called by TelemetryClientProvider when reference count reaches zero. + */ + close(): void { + if (this.closed) { + return; + } + try { + this.context.getLogger().log(LogLevel.debug, `Closing TelemetryClient for host: ${this.host}`); + } catch { + // swallow + } finally { + this.closed = true; + } + } +} + +export default TelemetryClient; diff --git a/lib/telemetry/TelemetryClientProvider.ts b/lib/telemetry/TelemetryClientProvider.ts new file mode 100644 index 00000000..c0da29f0 --- /dev/null +++ b/lib/telemetry/TelemetryClientProvider.ts @@ -0,0 +1,173 @@ +/** + * Copyright (c) 2025 Databricks Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import IClientContext from '../contracts/IClientContext'; +import { LogLevel } from '../contracts/IDBSQLLogger'; +import TelemetryClient from './TelemetryClient'; + +/** + * Holds a telemetry client and its reference count. + * The reference count tracks how many connections are using this client. + */ +interface TelemetryClientHolder { + client: TelemetryClient; + refCount: number; +} + +// Soft cap on distinct host entries. Above this the provider warns once so a +// misconfigured caller (per-request hosts, unnormalized aliases) is visible in +// logs rather than silently growing the map. +const MAX_CLIENTS_SOFT_LIMIT = 128; + +/** + * Manages one telemetry client per host. + * Prevents rate limiting by sharing clients across connections to the same host. + * Instance-based (not singleton), stored in DBSQLClient. + * + * Reference counts are incremented and decremented synchronously, and + * `close()` is sync today, so there is no await between map mutation and + * client teardown. The map entry is removed before `close()` runs so a + * concurrent `getOrCreateClient` call for the same host gets a fresh + * instance rather than receiving this closing one. When `close()` becomes + * async (e.g. HTTP flush in [5/7]) the flow will need to `await` after the + * delete to preserve the same invariant. + */ +class TelemetryClientProvider { + private clients: Map; + + private softLimitWarned = false; + + constructor(private context: IClientContext) { + this.clients = new Map(); + const logger = context.getLogger(); + logger.log(LogLevel.debug, 'Created TelemetryClientProvider'); + } + + /** + * Canonicalize host so aliases (scheme, default port, trailing slash, case, + * trailing dot, surrounding whitespace) map to the same entry. Kept to a + * lightweight lexical normalization — `buildTelemetryUrl` still performs + * the strict security validation when a request is actually built. + */ + private static normalizeHostKey(host: string): string { + return host + .trim() + .toLowerCase() + .replace(/^https?:\/\//, '') + .replace(/\/+$/, '') + .replace(/\.$/, '') + .replace(/:443$/, ''); + } + + /** + * Gets or creates a telemetry client for the specified host. + * Increments the reference count for the client. + * + * @param host The host identifier (e.g., "workspace.cloud.databricks.com") + * @returns The telemetry client for the host + */ + getOrCreateClient(host: string): TelemetryClient { + const logger = this.context.getLogger(); + const key = TelemetryClientProvider.normalizeHostKey(host); + let holder = this.clients.get(key); + + if (!holder) { + const client = new TelemetryClient(this.context, key); + holder = { + client, + refCount: 0, + }; + this.clients.set(key, holder); + logger.log(LogLevel.debug, `Created new TelemetryClient for host: ${host}`); + + if (!this.softLimitWarned && this.clients.size > MAX_CLIENTS_SOFT_LIMIT) { + this.softLimitWarned = true; + logger.log( + LogLevel.warn, + `TelemetryClientProvider has ${this.clients.size} distinct hosts — possible alias or leak`, + ); + } + } + + holder.refCount += 1; + logger.log(LogLevel.debug, `TelemetryClient reference count for ${host}: ${holder.refCount}`); + + return holder.client; + } + + /** + * Releases a telemetry client for the specified host. + * Decrements the reference count and closes the client when it reaches zero. + * + * @param host The host identifier + */ + releaseClient(host: string): void { + const logger = this.context.getLogger(); + const key = TelemetryClientProvider.normalizeHostKey(host); + const holder = this.clients.get(key); + + if (!holder) { + logger.log(LogLevel.debug, `No TelemetryClient found for host: ${host}`); + return; + } + + // Guard against double-release: a caller releasing more times than it got + // would otherwise drive refCount negative and close a client another + // caller is still holding. Warn loudly and refuse to decrement further. + if (holder.refCount <= 0) { + logger.log(LogLevel.warn, `Unbalanced release for TelemetryClient host: ${host}`); + return; + } + + holder.refCount -= 1; + logger.log(LogLevel.debug, `TelemetryClient reference count for ${host}: ${holder.refCount}`); + + // Close and remove client when reference count reaches zero. + // Remove from map before calling close so a concurrent getOrCreateClient + // creates a fresh client rather than receiving this closing one. + if (holder.refCount <= 0) { + this.clients.delete(key); + try { + holder.client.close(); + logger.log(LogLevel.debug, `Closed and removed TelemetryClient for host: ${host}`); + } catch (error) { + const msg = error instanceof Error ? error.message : String(error); + logger.log(LogLevel.debug, `Error releasing TelemetryClient: ${msg}`); + } + } + } + + /** + * @internal Exposed for testing only. + */ + getRefCount(host: string): number { + const holder = this.clients.get(TelemetryClientProvider.normalizeHostKey(host)); + return holder ? holder.refCount : 0; + } + + /** + * @internal Exposed for testing only. + */ + getActiveClients(): Map { + const result = new Map(); + for (const [host, holder] of this.clients.entries()) { + result.set(host, holder.client); + } + return result; + } +} + +export default TelemetryClientProvider; diff --git a/lib/telemetry/urlUtils.ts b/lib/telemetry/urlUtils.ts new file mode 100644 index 00000000..4dd8535e --- /dev/null +++ b/lib/telemetry/urlUtils.ts @@ -0,0 +1,31 @@ +/** + * Copyright (c) 2025 Databricks Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * Build full URL from host and path, handling protocol correctly. + * @param host The hostname (with or without protocol) + * @param path The path to append (should start with /) + * @returns Full URL with protocol + */ +// eslint-disable-next-line import/prefer-default-export +export function buildUrl(host: string, path: string): string { + // Check if host already has protocol + if (host.startsWith('http://') || host.startsWith('https://')) { + return `${host}${path}`; + } + // Add https:// if no protocol present + return `https://${host}${path}`; +} diff --git a/tests/unit/telemetry/TelemetryClient.test.ts b/tests/unit/telemetry/TelemetryClient.test.ts new file mode 100644 index 00000000..4f4c3f1d --- /dev/null +++ b/tests/unit/telemetry/TelemetryClient.test.ts @@ -0,0 +1,152 @@ +/** + * Copyright (c) 2025 Databricks Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { expect } from 'chai'; +import sinon from 'sinon'; +import TelemetryClient from '../../../lib/telemetry/TelemetryClient'; +import ClientContextStub from '../.stubs/ClientContextStub'; +import { LogLevel } from '../../../lib/contracts/IDBSQLLogger'; + +describe('TelemetryClient', () => { + const HOST = 'workspace.cloud.databricks.com'; + + describe('Constructor', () => { + it('should create client with host', () => { + const context = new ClientContextStub(); + const client = new TelemetryClient(context, HOST); + + expect(client.getHost()).to.equal(HOST); + expect(client.isClosed()).to.be.false; + }); + + it('should log creation at debug level', () => { + const context = new ClientContextStub(); + const logSpy = sinon.spy(context.logger, 'log'); + + new TelemetryClient(context, HOST); + + expect(logSpy.calledWith(LogLevel.debug, `Created TelemetryClient for host: ${HOST}`)).to.be.true; + }); + }); + + describe('getHost', () => { + it('should return the host identifier', () => { + const context = new ClientContextStub(); + const client = new TelemetryClient(context, HOST); + + expect(client.getHost()).to.equal(HOST); + }); + }); + + describe('isClosed', () => { + it('should return false initially', () => { + const context = new ClientContextStub(); + const client = new TelemetryClient(context, HOST); + + expect(client.isClosed()).to.be.false; + }); + + it('should return true after close', async () => { + const context = new ClientContextStub(); + const client = new TelemetryClient(context, HOST); + + client.close(); + + expect(client.isClosed()).to.be.true; + }); + }); + + describe('close', () => { + it('should set closed flag', async () => { + const context = new ClientContextStub(); + const client = new TelemetryClient(context, HOST); + + client.close(); + + expect(client.isClosed()).to.be.true; + }); + + it('should log closure at debug level', async () => { + const context = new ClientContextStub(); + const logSpy = sinon.spy(context.logger, 'log'); + const client = new TelemetryClient(context, HOST); + + client.close(); + + expect(logSpy.calledWith(LogLevel.debug, `Closing TelemetryClient for host: ${HOST}`)).to.be.true; + }); + + it('should be idempotent', async () => { + const context = new ClientContextStub(); + const logSpy = sinon.spy(context.logger, 'log'); + const client = new TelemetryClient(context, HOST); + + client.close(); + const firstCallCount = logSpy.callCount; + + client.close(); + + // Should not log again on second close + expect(logSpy.callCount).to.equal(firstCallCount); + expect(client.isClosed()).to.be.true; + }); + + it('should swallow all exceptions', async () => { + const context = new ClientContextStub(); + const client = new TelemetryClient(context, HOST); + + // Force an error by stubbing the logger + const error = new Error('Logger error'); + sinon.stub(context.logger, 'log').throws(error); + + expect(() => client.close()).to.not.throw(); + }); + + it('should still set closed when logger throws', () => { + const context = new ClientContextStub(); + const client = new TelemetryClient(context, HOST); + + sinon.stub(context.logger, 'log').throws(new Error('Logger error')); + + client.close(); + + expect(client.isClosed()).to.be.true; + }); + }); + + describe('Context usage', () => { + it('should use logger from context', () => { + const context = new ClientContextStub(); + const logSpy = sinon.spy(context.logger, 'log'); + + new TelemetryClient(context, HOST); + + expect(logSpy.called).to.be.true; + }); + + it('should log all messages at debug level only', async () => { + const context = new ClientContextStub(); + const logSpy = sinon.spy(context.logger, 'log'); + const client = new TelemetryClient(context, HOST); + + client.close(); + + logSpy.getCalls().forEach((call) => { + expect(call.args[0]).to.equal(LogLevel.debug); + }); + }); + }); +}); diff --git a/tests/unit/telemetry/TelemetryClientProvider.test.ts b/tests/unit/telemetry/TelemetryClientProvider.test.ts new file mode 100644 index 00000000..59ae3b99 --- /dev/null +++ b/tests/unit/telemetry/TelemetryClientProvider.test.ts @@ -0,0 +1,474 @@ +/** + * Copyright (c) 2025 Databricks Contributors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { expect } from 'chai'; +import sinon from 'sinon'; +import TelemetryClientProvider from '../../../lib/telemetry/TelemetryClientProvider'; +import TelemetryClient from '../../../lib/telemetry/TelemetryClient'; +import ClientContextStub from '../.stubs/ClientContextStub'; +import { LogLevel } from '../../../lib/contracts/IDBSQLLogger'; + +describe('TelemetryClientProvider', () => { + const HOST1 = 'workspace1.cloud.databricks.com'; + const HOST2 = 'workspace2.cloud.databricks.com'; + + describe('Constructor', () => { + it('should create provider with empty client map', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + expect(provider.getActiveClients().size).to.equal(0); + }); + + it('should log creation at debug level', () => { + const context = new ClientContextStub(); + const logSpy = sinon.spy(context.logger, 'log'); + + new TelemetryClientProvider(context); + + expect(logSpy.calledWith(LogLevel.debug, sinon.match(/created.*telemetryclientprovider/i))).to.be.true; + }); + }); + + describe('getOrCreateClient', () => { + it('should create one client per host', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client1 = provider.getOrCreateClient(HOST1); + const client2 = provider.getOrCreateClient(HOST2); + + expect(client1).to.be.instanceOf(TelemetryClient); + expect(client2).to.be.instanceOf(TelemetryClient); + expect(client1).to.not.equal(client2); + expect(provider.getActiveClients().size).to.equal(2); + }); + + it('should share client across multiple connections to same host', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client1 = provider.getOrCreateClient(HOST1); + const client2 = provider.getOrCreateClient(HOST1); + const client3 = provider.getOrCreateClient(HOST1); + + expect(client1).to.equal(client2); + expect(client2).to.equal(client3); + expect(provider.getActiveClients().size).to.equal(1); + }); + + it('should increment reference count on each call', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + provider.getOrCreateClient(HOST1); + expect(provider.getRefCount(HOST1)).to.equal(1); + + provider.getOrCreateClient(HOST1); + expect(provider.getRefCount(HOST1)).to.equal(2); + + provider.getOrCreateClient(HOST1); + expect(provider.getRefCount(HOST1)).to.equal(3); + }); + + it('should log client creation at debug level', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + const logSpy = sinon.spy(context.logger, 'log'); + + provider.getOrCreateClient(HOST1); + + expect(logSpy.calledWith(LogLevel.debug, sinon.match(/created new telemetryclient/i))).to.be.true; + }); + + it('should log reference count at debug level', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + const logSpy = sinon.spy(context.logger, 'log'); + + provider.getOrCreateClient(HOST1); + + expect(logSpy.calledWith(LogLevel.debug, sinon.match(/reference count for/i).and(sinon.match(/: 1$/)))).to.be + .true; + }); + + it('should pass context to TelemetryClient', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client = provider.getOrCreateClient(HOST1); + + expect(client.getHost()).to.equal(HOST1); + }); + }); + + describe('releaseClient', () => { + it('should decrement reference count on release', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + provider.getOrCreateClient(HOST1); + provider.getOrCreateClient(HOST1); + provider.getOrCreateClient(HOST1); + expect(provider.getRefCount(HOST1)).to.equal(3); + + provider.releaseClient(HOST1); + expect(provider.getRefCount(HOST1)).to.equal(2); + + provider.releaseClient(HOST1); + expect(provider.getRefCount(HOST1)).to.equal(1); + }); + + it('should close client when reference count reaches zero', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client = provider.getOrCreateClient(HOST1); + const closeSpy = sinon.spy(client, 'close'); + + provider.releaseClient(HOST1); + + expect(closeSpy.calledOnce).to.be.true; + expect(client.isClosed()).to.be.true; + }); + + it('should remove client from map when reference count reaches zero', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + provider.getOrCreateClient(HOST1); + expect(provider.getActiveClients().size).to.equal(1); + + provider.releaseClient(HOST1); + + expect(provider.getActiveClients().size).to.equal(0); + expect(provider.getRefCount(HOST1)).to.equal(0); + }); + + it('should NOT close client while other connections exist', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client = provider.getOrCreateClient(HOST1); + provider.getOrCreateClient(HOST1); + provider.getOrCreateClient(HOST1); + const closeSpy = sinon.spy(client, 'close'); + + provider.releaseClient(HOST1); + + expect(closeSpy.called).to.be.false; + expect(client.isClosed()).to.be.false; + expect(provider.getActiveClients().size).to.equal(1); + }); + + it('should handle releasing non-existent client gracefully', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + const logSpy = sinon.spy(context.logger, 'log'); + + provider.releaseClient(HOST1); + + expect(logSpy.calledWith(LogLevel.debug, sinon.match(/no telemetryclient found/i))).to.be.true; + }); + + it('should log reference count decrease at debug level', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + const logSpy = sinon.spy(context.logger, 'log'); + + provider.getOrCreateClient(HOST1); + provider.getOrCreateClient(HOST1); + + provider.releaseClient(HOST1); + + expect(logSpy.calledWith(LogLevel.debug, sinon.match(/reference count for/i).and(sinon.match(/: 1$/)))).to.be + .true; + }); + + it('should log client closure at debug level', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + const logSpy = sinon.spy(context.logger, 'log'); + + provider.getOrCreateClient(HOST1); + provider.releaseClient(HOST1); + + expect(logSpy.calledWith(LogLevel.debug, sinon.match(/closed and removed telemetryclient/i))).to.be.true; + }); + + it('should swallow Error during client closure', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client = provider.getOrCreateClient(HOST1); + const error = new Error('Close error'); + sinon.stub(client, 'close').throws(error); + const logSpy = sinon.spy(context.logger, 'log'); + + provider.releaseClient(HOST1); + + expect( + logSpy.calledWith( + LogLevel.debug, + sinon.match(/error releasing telemetryclient/i).and(sinon.match(/close error/i)), + ), + ).to.be.true; + }); + + it('should swallow non-Error throws during client closure', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client = provider.getOrCreateClient(HOST1); + // Non-Error throws — string, null, undefined — must not escape the catch. + sinon.stub(client, 'close').callsFake(() => { + // eslint-disable-next-line no-throw-literal + throw 'stringy-error'; + }); + const logSpy = sinon.spy(context.logger, 'log'); + + expect(() => provider.releaseClient(HOST1)).to.not.throw(); + expect( + logSpy.calledWith( + LogLevel.debug, + sinon.match(/error releasing telemetryclient/i).and(sinon.match(/stringy-error/)), + ), + ).to.be.true; + }); + + it('should not throw or corrupt state on double-release', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + // One get, two releases — the second must not throw and must not + // leave the provider in a state where refCount is negative. + provider.getOrCreateClient(HOST1); + provider.releaseClient(HOST1); + + expect(() => provider.releaseClient(HOST1)).to.not.throw(); + expect(provider.getRefCount(HOST1)).to.equal(0); + expect(provider.getActiveClients().size).to.equal(0); + }); + + it('should return a fresh non-closed client after full release', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const first = provider.getOrCreateClient(HOST1); + provider.releaseClient(HOST1); + expect(first.isClosed()).to.be.true; + + const second = provider.getOrCreateClient(HOST1); + expect(second).to.not.equal(first); + expect(second.isClosed()).to.be.false; + expect(provider.getRefCount(HOST1)).to.equal(1); + }); + }); + + describe('Host normalization', () => { + it('should treat scheme, case, port and trailing slash as the same host', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const a = provider.getOrCreateClient('workspace.cloud.databricks.com'); + const b = provider.getOrCreateClient('https://workspace.cloud.databricks.com'); + const c = provider.getOrCreateClient('https://WorkSpace.CLOUD.databricks.com/'); + const d = provider.getOrCreateClient('workspace.cloud.databricks.com:443'); + const e = provider.getOrCreateClient(' workspace.cloud.databricks.com. '); + + expect(a).to.equal(b); + expect(a).to.equal(c); + expect(a).to.equal(d); + expect(a).to.equal(e); + expect(provider.getActiveClients().size).to.equal(1); + expect(provider.getRefCount('workspace.cloud.databricks.com')).to.equal(5); + }); + + it('should release under an alias correctly', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + provider.getOrCreateClient('example.com'); + provider.releaseClient('HTTPS://Example.COM/'); + + expect(provider.getRefCount('example.com')).to.equal(0); + expect(provider.getActiveClients().size).to.equal(0); + }); + }); + + describe('Reference counting', () => { + it('should track reference counts independently per host', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + provider.getOrCreateClient(HOST1); + provider.getOrCreateClient(HOST1); + provider.getOrCreateClient(HOST2); + provider.getOrCreateClient(HOST2); + provider.getOrCreateClient(HOST2); + + expect(provider.getRefCount(HOST1)).to.equal(2); + expect(provider.getRefCount(HOST2)).to.equal(3); + + provider.releaseClient(HOST1); + + expect(provider.getRefCount(HOST1)).to.equal(1); + expect(provider.getRefCount(HOST2)).to.equal(3); + }); + + it('should close only last connection for each host', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client1 = provider.getOrCreateClient(HOST1); + provider.getOrCreateClient(HOST1); + const client2 = provider.getOrCreateClient(HOST2); + + provider.releaseClient(HOST1); + expect(client1.isClosed()).to.be.false; + expect(provider.getActiveClients().size).to.equal(2); + + provider.releaseClient(HOST1); + expect(client1.isClosed()).to.be.true; + expect(provider.getActiveClients().size).to.equal(1); + + provider.releaseClient(HOST2); + expect(client2.isClosed()).to.be.true; + expect(provider.getActiveClients().size).to.equal(0); + }); + }); + + describe('Per-host isolation', () => { + it('should isolate clients by host', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client1 = provider.getOrCreateClient(HOST1); + const client2 = provider.getOrCreateClient(HOST2); + + expect(client1.getHost()).to.equal(HOST1); + expect(client2.getHost()).to.equal(HOST2); + expect(client1).to.not.equal(client2); + }); + + it('should allow closing one host without affecting others', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client1 = provider.getOrCreateClient(HOST1); + const client2 = provider.getOrCreateClient(HOST2); + + provider.releaseClient(HOST1); + + expect(client1.isClosed()).to.be.true; + expect(client2.isClosed()).to.be.false; + expect(provider.getActiveClients().size).to.equal(1); + }); + }); + + describe('getRefCount', () => { + it('should return 0 for non-existent host', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + expect(provider.getRefCount(HOST1)).to.equal(0); + }); + + it('should return current reference count for existing host', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + provider.getOrCreateClient(HOST1); + expect(provider.getRefCount(HOST1)).to.equal(1); + + provider.getOrCreateClient(HOST1); + expect(provider.getRefCount(HOST1)).to.equal(2); + }); + }); + + describe('getActiveClients', () => { + it('should return empty map initially', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const clients = provider.getActiveClients(); + + expect(clients.size).to.equal(0); + }); + + it('should return all active clients', () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + const client1 = provider.getOrCreateClient(HOST1); + const client2 = provider.getOrCreateClient(HOST2); + + const clients = provider.getActiveClients(); + + expect(clients.size).to.equal(2); + expect(clients.get(HOST1)).to.equal(client1); + expect(clients.get(HOST2)).to.equal(client2); + }); + + it('should not include closed clients', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + + provider.getOrCreateClient(HOST1); + provider.getOrCreateClient(HOST2); + + provider.releaseClient(HOST1); + + const clients = provider.getActiveClients(); + + expect(clients.size).to.equal(1); + expect(clients.has(HOST1)).to.be.false; + expect(clients.has(HOST2)).to.be.true; + }); + }); + + describe('Context usage', () => { + it('should use logger from context for all logging', () => { + const context = new ClientContextStub(); + const logSpy = sinon.spy(context.logger, 'log'); + const provider = new TelemetryClientProvider(context); + + provider.getOrCreateClient(HOST1); + + expect(logSpy.called).to.be.true; + logSpy.getCalls().forEach((call) => { + expect(call.args[0]).to.equal(LogLevel.debug); + }); + }); + + it('should log close errors at debug level', async () => { + const context = new ClientContextStub(); + const provider = new TelemetryClientProvider(context); + const logSpy = sinon.spy(context.logger, 'log'); + + const client = provider.getOrCreateClient(HOST1); + sinon.stub(client, 'close').throws(new Error('Test error')); + + provider.releaseClient(HOST1); + + const errorLogs = logSpy.getCalls().filter((call) => /error releasing/i.test(String(call.args[1]))); + expect(errorLogs.length).to.be.greaterThan(0); + errorLogs.forEach((call) => { + expect(call.args[0]).to.equal(LogLevel.debug); + }); + }); + }); +});