fix: address code review findings for telemetry infrastructure

- Add workspace_id to telemetry log serialization (was silently dropped)
- Fix socket leak: consume HTTP response body on success and error paths
- Add 10s timeout to all telemetry/feature-flag fetch calls
- Fix thundering herd: deduplicate concurrent feature flag fetches
- Add Promise.resolve().catch() to flush() to prevent unhandled rejections
- Add HALF_OPEN inflight guard to CircuitBreaker (limit to 1 probe)
- Rename TelemetryEventType.ERROR to 'telemetry.error' (avoid EventEmitter collision)
- Extract shared buildTelemetryUrl utility enforcing HTTPS
- Clamp server-provided TTL to [60, 3600] seconds
- Skip export silently when auth headers are missing
- Log circuit breaker OPEN transitions at warn level
- Fix CircuitBreakerStub HALF_OPEN behavior to match real implementation
- Snapshot Map keys before iteration in close()
- Remove unnecessary 'as any' cast and '| string' type widening

Co-authored-by: Isaac

Author: samikshya-db

lib/telemetry/CircuitBreaker.ts

@@ -68,6 +68,9 @@ export class CircuitBreaker {
 
   private nextAttempt?: Date;
 
+  /** Number of in-flight requests in HALF_OPEN state (limits to 1 probe) */
+  private halfOpenInflight = 0;
+
   private readonly config: CircuitBreakerConfig;
 
   constructor(private context: IClientContext, config?: Partial<CircuitBreakerConfig>) {
@@ -95,16 +98,30 @@ export class CircuitBreaker {
       // Timeout expired, transition to HALF_OPEN
       this.state = CircuitBreakerState.HALF_OPEN;
       this.successCount = 0;
+      this.halfOpenInflight = 0;
       logger.log(LogLevel.debug, 'Circuit breaker transitioned to HALF_OPEN');
     }
 
+    // In HALF_OPEN state, allow only one probe request at a time
+    if (this.state === CircuitBreakerState.HALF_OPEN && this.halfOpenInflight > 0) {
+      throw new Error('Circuit breaker OPEN');
+    }
+
+    if (this.state === CircuitBreakerState.HALF_OPEN) {
+      this.halfOpenInflight += 1;
+    }
+
     try {
       const result = await operation();
       this.onSuccess();
       return result;
     } catch (error) {
       this.onFailure();
       throw error;
+    } finally {
+      if (this.halfOpenInflight > 0) {
+        this.halfOpenInflight -= 1;
+      }
     }
   }
 
@@ -171,7 +188,11 @@ export class CircuitBreaker {
     if (this.state === CircuitBreakerState.HALF_OPEN || this.failureCount >= this.config.failureThreshold) {
       this.state = CircuitBreakerState.OPEN;
       this.nextAttempt = new Date(Date.now() + this.config.timeout);
-      logger.log(LogLevel.debug, `Circuit breaker transitioned to OPEN (will retry after ${this.config.timeout}ms)`);
+      // Log at warn level for OPEN transitions — meaningful operational signal
+      logger.log(
+        LogLevel.warn,
+        `Telemetry circuit breaker OPEN after ${this.failureCount} failures (will retry after ${this.config.timeout}ms)`,
+      );
     }
   }
 }
@@ -202,6 +223,12 @@ export class CircuitBreakerRegistry {
       this.breakers.set(host, breaker);
       const logger = this.context.getLogger();
       logger.log(LogLevel.debug, `Created circuit breaker for host: ${host}`);
+    } else if (config) {
+      const logger = this.context.getLogger();
+      logger.log(
+        LogLevel.debug,
+        `Circuit breaker for host ${host} already exists; provided config will be ignored`,
+      );
     }
     return breaker;
   }

lib/telemetry/DatabricksTelemetryExporter.ts

@@ -19,8 +19,9 @@ import fetch, { Response, RequestInit } from 'node-fetch';
 import IClientContext from '../contracts/IClientContext';
 import { LogLevel } from '../contracts/IDBSQLLogger';
 import { TelemetryMetric, DEFAULT_TELEMETRY_CONFIG } from './types';
-import { CircuitBreakerRegistry } from './CircuitBreaker';
+import { CircuitBreaker, CircuitBreakerRegistry } from './CircuitBreaker';
 import ExceptionClassifier from './ExceptionClassifier';
+import { buildTelemetryUrl } from './telemetryUtils';
 import driverVersion from '../version';
 
 /**
@@ -87,8 +88,8 @@ interface DatabricksTelemetryPayload {
  * Exports telemetry metrics to Databricks telemetry service.
  *
  * Endpoints:
- * - Authenticated: /api/2.0/sql/telemetry-ext
- * - Unauthenticated: /api/2.0/sql/telemetry-unauth
+ * - Authenticated: /telemetry-ext
+ * - Unauthenticated: /telemetry-unauth
  *
  * Features:
  * - Circuit breaker integration for endpoint protection
@@ -98,7 +99,7 @@ interface DatabricksTelemetryPayload {
  * - CRITICAL: All logging at LogLevel.debug ONLY
  */
 export default class DatabricksTelemetryExporter {
-  private circuitBreaker;
+  private readonly circuitBreaker: CircuitBreaker;
 
   private readonly userAgent: string;
 
@@ -207,8 +208,8 @@ export default class DatabricksTelemetryExporter {
     // Determine endpoint based on authentication mode
     const authenticatedExport = config.telemetryAuthenticatedExport ?? DEFAULT_TELEMETRY_CONFIG.authenticatedExport;
     const endpoint = authenticatedExport
-      ? this.buildUrl(this.host, '/telemetry-ext')
-      : this.buildUrl(this.host, '/telemetry-unauth');
+      ? buildTelemetryUrl(this.host, '/telemetry-ext')
+      : buildTelemetryUrl(this.host, '/telemetry-unauth');
 
     // Format payload - each log is JSON-stringified to match JDBC format
     const telemetryLogs = metrics.map((m) => this.toTelemetryLog(m));
@@ -230,6 +231,14 @@ export default class DatabricksTelemetryExporter {
     // Get authentication headers if using authenticated endpoint
     const authHeaders = authenticatedExport ? await this.context.getAuthHeaders() : {};
 
+    // Skip export if authenticated mode is requested but no auth headers available.
+    // Note: all auth providers in this codebase return plain objects (Record<string, string>).
+    const headersObj = authHeaders as Record<string, string>;
+    if (authenticatedExport && (!headersObj || !headersObj['Authorization'])) {
+      logger.log(LogLevel.debug, 'Skipping telemetry export: authenticated mode but no Authorization header');
+      return;
+    }
+
     // Make HTTP POST request with authentication and proxy support
     const response: Response = await this.sendRequest(endpoint, {
       method: 'POST',
@@ -239,14 +248,19 @@ export default class DatabricksTelemetryExporter {
         'User-Agent': this.userAgent,
       },
       body: JSON.stringify(payload),
+      timeout: 10000, // 10 second timeout to prevent indefinite hangs
     });
 
     if (!response.ok) {
+      // Consume response body to release socket back to connection pool
+      await response.text().catch(() => {});
       const error: any = new Error(`Telemetry export failed: ${response.status} ${response.statusText}`);
       error.statusCode = response.status;
       throw error;
     }
 
+    // Consume response body to release socket back to connection pool
+    await response.text().catch(() => {});
     logger.log(LogLevel.debug, `Successfully exported ${metrics.length} telemetry metrics`);
   }
 
@@ -265,6 +279,7 @@ export default class DatabricksTelemetryExporter {
    */
   private toTelemetryLog(metric: TelemetryMetric): DatabricksTelemetryLog {
     const log: DatabricksTelemetryLog = {
+      workspace_id: metric.workspaceId,
       frontend_log_event_id: this.generateUUID(),
       context: {
         client_context: {
@@ -321,16 +336,6 @@ export default class DatabricksTelemetryExporter {
     return log;
   }
 
-  /**
-   * Build full URL from host and path, handling protocol correctly.
-   */
-  private buildUrl(host: string, path: string): string {
-    if (host.startsWith('http://') || host.startsWith('https://')) {
-      return `${host}${path}`;
-    }
-    return `https://${host}${path}`;
-  }
-
   /**
    * Generate a UUID v4.
    */

lib/telemetry/FeatureFlagCache.ts

@@ -17,6 +17,7 @@
 import fetch from 'node-fetch';
 import IClientContext from '../contracts/IClientContext';
 import { LogLevel } from '../contracts/IDBSQLLogger';
+import { buildTelemetryUrl } from './telemetryUtils';
 import driverVersion from '../version';
 
 /**
@@ -37,8 +38,15 @@ export interface FeatureFlagContext {
 export default class FeatureFlagCache {
   private contexts: Map<string, FeatureFlagContext>;
 
+  /** In-flight fetch promises for deduplication (prevents thundering herd) */
+  private fetchPromises: Map<string, Promise<boolean>> = new Map();
+
   private readonly CACHE_DURATION_MS = 15 * 60 * 1000; // 15 minutes
 
+  private readonly MIN_CACHE_DURATION_S = 60; // 1 minute minimum TTL
+
+  private readonly MAX_CACHE_DURATION_S = 3600; // 1 hour maximum TTL
+
   private readonly FEATURE_FLAG_NAME = 'databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForNodeJs';
 
   constructor(private context: IClientContext) {
@@ -72,6 +80,7 @@ export default class FeatureFlagCache {
       ctx.refCount -= 1;
       if (ctx.refCount <= 0) {
         this.contexts.delete(host);
+        this.fetchPromises.delete(host); // Invalidate stale in-flight fetch
       }
     }
   }
@@ -91,22 +100,34 @@ export default class FeatureFlagCache {
     const isExpired = !ctx.lastFetched || Date.now() - ctx.lastFetched.getTime() > ctx.cacheDuration;
 
     if (isExpired) {
-      try {
-        // Fetch feature flag from server
-        ctx.telemetryEnabled = await this.fetchFeatureFlag(host);
-        ctx.lastFetched = new Date();
-      } catch (error: any) {
-        // Log at debug level only, never propagate exceptions
-        logger.log(LogLevel.debug, `Error fetching feature flag: ${error.message}`);
+      // Deduplicate concurrent fetches for the same host (prevents thundering herd)
+      if (!this.fetchPromises.has(host)) {
+        const fetchPromise = this.fetchFeatureFlag(host)
+          .then((enabled) => {
+            ctx.telemetryEnabled = enabled;
+            ctx.lastFetched = new Date();
+            return enabled;
+          })
+          .catch((error: any) => {
+            logger.log(LogLevel.debug, `Error fetching feature flag: ${error.message}`);
+            return ctx.telemetryEnabled ?? false;
+          })
+          .finally(() => {
+            this.fetchPromises.delete(host);
+          });
+        this.fetchPromises.set(host, fetchPromise);
       }
+
+      // Promise is guaranteed to resolve (never rejects) due to .catch() in the chain above
+      await this.fetchPromises.get(host);
     }
 
     return ctx.telemetryEnabled ?? false;
   }
 
   /**
    * Fetches feature flag from server using connector-service API.
-   * Calls GET /api/2.0/connector-service/feature-flags/OSS_NODEJS/{version}
+   * Calls GET /api/2.0/connector-service/feature-flags/NODEJS/{version}
    *
    * @param host The host to fetch feature flag for
    * @returns true if feature flag is enabled, false otherwise
@@ -119,7 +140,7 @@ export default class FeatureFlagCache {
       const version = this.getDriverVersion();
 
       // Build feature flags endpoint for Node.js driver
-      const endpoint = this.buildUrl(host, `/api/2.0/connector-service/feature-flags/NODEJS/${version}`);
+      const endpoint = buildTelemetryUrl(host, `/api/2.0/connector-service/feature-flags/NODEJS/${version}`);
 
       // Get authentication headers
       const authHeaders = await this.context.getAuthHeaders();
@@ -139,9 +160,12 @@ export default class FeatureFlagCache {
           'User-Agent': `databricks-sql-nodejs/${driverVersion}`,
         },
         agent, // Include agent for proxy support
+        timeout: 10000, // 10 second timeout to prevent indefinite hangs
       });
 
       if (!response.ok) {
+        // Consume response body to release socket back to connection pool
+        await response.text().catch(() => {});
         logger.log(LogLevel.debug, `Feature flag fetch failed: ${response.status} ${response.statusText}`);
         return false;
       }
@@ -151,11 +175,12 @@ export default class FeatureFlagCache {
 
       // Response format: { flags: [{ name: string, value: string }], ttl_seconds?: number }
       if (data && data.flags && Array.isArray(data.flags)) {
-        // Update cache duration if TTL provided
+        // Update cache duration if TTL provided, clamped to safe bounds
         const ctx = this.contexts.get(host);
-        if (ctx && data.ttl_seconds) {
-          ctx.cacheDuration = data.ttl_seconds * 1000; // Convert to milliseconds
-          logger.log(LogLevel.debug, `Updated cache duration to ${data.ttl_seconds} seconds`);
+        if (ctx && typeof data.ttl_seconds === 'number' && data.ttl_seconds > 0) {
+          const clampedTtl = Math.max(this.MIN_CACHE_DURATION_S, Math.min(this.MAX_CACHE_DURATION_S, data.ttl_seconds));
+          ctx.cacheDuration = clampedTtl * 1000; // Convert to milliseconds
+          logger.log(LogLevel.debug, `Updated cache duration to ${clampedTtl} seconds`);
         }
 
         // Look for our specific feature flag
@@ -180,16 +205,6 @@ export default class FeatureFlagCache {
     }
   }
 
-  /**
-   * Build full URL from host and path, handling protocol correctly.
-   */
-  private buildUrl(host: string, path: string): string {
-    if (host.startsWith('http://') || host.startsWith('https://')) {
-      return `${host}${path}`;
-    }
-    return `https://${host}${path}`;
-  }
-
   /**
    * Gets the driver version without -oss suffix for API calls.
    * Format: "1.12.0" from "1.12.0-oss"

lib/telemetry/MetricsAggregator.ts

@@ -329,8 +329,11 @@ export default class MetricsAggregator {
 
       logger.log(LogLevel.debug, `Flushing ${metricsToExport.length} telemetry metrics`);
 
-      // Export metrics (exporter.export never throws)
-      this.exporter.export(metricsToExport);
+      // Export metrics - exporter.export is designed to never throw, but add safety catch
+      // to prevent unhandled promise rejections from crashing the process (Node.js 15+)
+      Promise.resolve(this.exporter.export(metricsToExport)).catch((err: any) => {
+        logger.log(LogLevel.debug, `Unexpected export error: ${err?.message}`);
+      });
 
       // Reset timer to avoid rapid successive flushes (e.g., batch flush at 25s then timer flush at 30s)
       // This ensures consistent spacing between exports and helps avoid rate limiting
@@ -380,8 +383,9 @@ export default class MetricsAggregator {
         this.flushTimer = null;
       }
 
-      // Complete any remaining statements
-      for (const statementId of this.statementMetrics.keys()) {
+      // Complete any remaining statements (snapshot keys to avoid mutation during iteration)
+      const remainingStatements = [...this.statementMetrics.keys()];
+      for (const statementId of remainingStatements) {
         this.completeStatement(statementId);
       }
 

lib/telemetry/TelemetryEventEmitter.ts

@@ -36,7 +36,7 @@ export default class TelemetryEventEmitter extends EventEmitter {
     super();
     // Check if telemetry is enabled from config
     // Default to false for safe rollout
-    const config = context.getConfig() as any;
+    const config = context.getConfig();
     this.enabled = config.telemetryEnabled ?? false;
   }
 

lib/telemetry/telemetryUtils.ts

@@ -0,0 +1,24 @@
+/**
+ * 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, always using HTTPS.
+ * Strips any existing protocol prefix and enforces HTTPS.
+ */
+export function buildTelemetryUrl(host: string, path: string): string {
+  const cleanHost = host.replace(/^https?:\/\//, '').replace(/\/+$/, '');
+  return `https://${cleanHost}${path}`;
+}

lib/telemetry/types.ts

@@ -27,7 +27,7 @@ export enum TelemetryEventType {
   STATEMENT_START = 'statement.start',
   STATEMENT_COMPLETE = 'statement.complete',
   CLOUDFETCH_CHUNK = 'cloudfetch.chunk',
-  ERROR = 'error',
+  ERROR = 'telemetry.error',
 }
 
 /**
@@ -78,7 +78,7 @@ export const DEFAULT_TELEMETRY_CONFIG: Required<TelemetryConfiguration> = {
  */
 export interface TelemetryEvent {
   /** Type of the event */
-  eventType: TelemetryEventType | string;
+  eventType: TelemetryEventType;
 
   /** Timestamp when the event occurred (milliseconds since epoch) */
   timestamp: number;

tests/unit/.stubs/CircuitBreakerStub.ts

@@ -112,7 +112,8 @@ export default class CircuitBreakerStub {
   private onFailure(): void {
     this.failureCount++;
     this.successCount = 0;
-    if (this.failureCount >= 5) {
+    // In HALF_OPEN state, any single failure immediately reopens (matches real implementation)
+    if (this.state === CircuitBreakerState.HALF_OPEN || this.failureCount >= 5) {
       this.state = CircuitBreakerState.OPEN;
     }
   }

tests/unit/.stubs/ClientContextStub.ts

@@ -51,6 +51,6 @@ export default class ClientContextStub implements IClientContext {
   }
 
   public async getAuthHeaders(): Promise<HeadersInit> {
-    return {};
+    return { Authorization: 'Bearer test-token' };
   }
 }