Address PR review feedback: fix critical/high telemetry issues

- Fix double-close panic: add sync.Once to aggregator.close()
- Fix Interceptor.Close() closing shared aggregator: only flush, not close
- Bound export goroutines with semaphore (cap 8) in flushUnlocked
- Use cancellable context in flushLoop so exports stop on shutdown
- Change UserConfig.EnableTelemetry from bool to ConfigValue[bool] so
  explicit false in DSN properly disables telemetry (not treated as unset)
- Fix classifyError to use ordered slice for deterministic classification

Co-authored-by: samikshya-chand_data

Author: samikshya-db

connector.go

@@ -76,18 +76,12 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	}
 	log := logger.WithContext(conn.id, driverctx.CorrelationIdFromContext(ctx), "")
 
-	// Initialize telemetry: pass user opt-in flag; if unset, feature flags decide
-	var enableTelemetry *bool
-	if c.cfg.EnableTelemetry {
-		trueVal := true
-		enableTelemetry = &trueVal
-	}
-
+	// Initialize telemetry: client config overlay decides; if unset, feature flags decide
 	conn.telemetry = telemetry.InitializeForConnection(
 		ctx,
 		c.cfg.Host,
 		c.client,
-		enableTelemetry,
+		c.cfg.EnableTelemetry,
 	)
 	if conn.telemetry != nil {
 		log.Debug().Msg("telemetry initialized for connection")

internal/config/config.go

@@ -99,8 +99,10 @@ type UserConfig struct {
 	RetryWaitMax   time.Duration
 	RetryMax       int
 	// Telemetry configuration
-	EnableTelemetry          bool // Opt-in for telemetry; follows client > server > default priority
-	Transport                http.RoundTripper
+	// Uses config overlay pattern: client > server > default.
+	// Unset = check server feature flag; explicitly true/false overrides the server.
+	EnableTelemetry ConfigValue[bool]
+	Transport       http.RoundTripper
 	UseLz4Compression        bool
 	EnableMetricViewMetadata bool
 	CloudFetchConfig
@@ -290,7 +292,7 @@ func ParseDSN(dsn string) (UserConfig, error) {
 		if err != nil {
 			return UserConfig{}, err
 		}
-		ucfg.EnableTelemetry = enableTelemetry
+		ucfg.EnableTelemetry = NewConfigValue(enableTelemetry)
 	}
 
 	// for timezone we do a case insensitive key match.

telemetry/aggregator.go

@@ -20,6 +20,11 @@ type metricsAggregator struct {
 	flushInterval time.Duration
 	stopCh        chan struct{}
 	flushTimer    *time.Ticker
+
+	closeOnce sync.Once
+	ctx       context.Context    // Cancellable context for in-flight exports
+	cancel    context.CancelFunc // Cancels ctx on close
+	exportSem chan struct{}       // Bounds concurrent export goroutines
 }
 
 // statementMetrics holds aggregated metrics for a statement.
@@ -38,13 +43,17 @@ type statementMetrics struct {
 
 // newMetricsAggregator creates a new metrics aggregator.
 func newMetricsAggregator(exporter *telemetryExporter, cfg *Config) *metricsAggregator {
+	ctx, cancel := context.WithCancel(context.Background())
 	agg := &metricsAggregator{
 		statements:    make(map[string]*statementMetrics),
 		batch:         make([]*telemetryMetric, 0, cfg.BatchSize),
 		exporter:      exporter,
 		batchSize:     cfg.BatchSize,
 		flushInterval: cfg.FlushInterval,
 		stopCh:        make(chan struct{}),
+		ctx:           ctx,
+		cancel:        cancel,
+		exportSem:     make(chan struct{}, 8), // Bound to 8 concurrent exports
 	}
 
 	// Start background flush timer
@@ -182,7 +191,7 @@ func (agg *metricsAggregator) flushLoop() {
 	for {
 		select {
 		case <-agg.flushTimer.C:
-			agg.flush(context.Background())
+			agg.flush(agg.ctx) // Use cancellable context so exports stop on shutdown
 		case <-agg.stopCh:
 			return
 		}
@@ -207,9 +216,18 @@ func (agg *metricsAggregator) flushUnlocked(ctx context.Context) {
 	copy(metrics, agg.batch)
 	agg.batch = agg.batch[:0]
 
+	// Acquire semaphore slot; skip export if already at capacity to prevent goroutine leaks
+	select {
+	case agg.exportSem <- struct{}{}:
+	default:
+		logger.Debug().Msg("telemetry: export semaphore full, dropping metrics batch")
+		return
+	}
+
 	// Export asynchronously
 	go func() {
 		defer func() {
+			<-agg.exportSem
 			if r := recover(); r != nil {
 				logger.Debug().Msgf("telemetry: async export panic: %v", r)
 			}
@@ -219,8 +237,12 @@ func (agg *metricsAggregator) flushUnlocked(ctx context.Context) {
 }
 
 // close stops the aggregator and flushes pending metrics.
+// Safe to call multiple times — subsequent calls are no-ops for the stop/cancel step.
 func (agg *metricsAggregator) close(ctx context.Context) error {
-	close(agg.stopCh)
+	agg.closeOnce.Do(func() {
+		close(agg.stopCh)
+		agg.cancel() // Cancel in-flight periodic export goroutines
+	})
 	agg.flush(ctx)
 	return nil
 }

telemetry/driver_integration.go

@@ -15,24 +15,19 @@ import (
 //   - ctx: Context for the initialization
 //   - host: Databricks host
 //   - httpClient: HTTP client for making requests
-//   - enableTelemetry: User opt-in flag (nil = unset, true = enable, false = disable)
+//   - enableTelemetry: Client config overlay (unset = check server flag, true/false = override server)
 //
 // Returns:
 //   - *Interceptor: Telemetry interceptor if enabled, nil otherwise
 func InitializeForConnection(
 	ctx context.Context,
 	host string,
 	httpClient *http.Client,
-	enableTelemetry *bool,
+	enableTelemetry config.ConfigValue[bool],
 ) *Interceptor {
-	// Create telemetry config
+	// Create telemetry config and apply client overlay
 	cfg := DefaultConfig()
-
-	// Set EnableTelemetry based on user preference
-	if enableTelemetry != nil {
-		cfg.EnableTelemetry = config.NewConfigValue(*enableTelemetry)
-	}
-	// else: leave unset (will check server feature flag)
+	cfg.EnableTelemetry = enableTelemetry
 
 	// Check if telemetry should be enabled
 	if !isTelemetryEnabled(ctx, cfg, host, httpClient) {

telemetry/errors.go

@@ -54,22 +54,25 @@ func classifyError(err error) string {
 
 	errMsg := strings.ToLower(err.Error())
 
-	// Check for common error patterns
-	patterns := map[string]string{
-		"timeout":        "timeout",
-		"context cancel": "cancelled",
-		"connection":     "connection_error",
-		"authentication": "auth_error",
-		"unauthorized":   "auth_error",
-		"forbidden":      "permission_error",
-		"not found":      "not_found",
-		"syntax":         "syntax_error",
-		"invalid":        "invalid_request",
+	// Ordered patterns — first match wins, ensuring deterministic classification.
+	patterns := []struct {
+		pattern   string
+		errorType string
+	}{
+		{"timeout", "timeout"},
+		{"context cancel", "cancelled"},
+		{"connection", "connection_error"},
+		{"authentication", "auth_error"},
+		{"unauthorized", "auth_error"},
+		{"forbidden", "permission_error"},
+		{"not found", "not_found"},
+		{"syntax", "syntax_error"},
+		{"invalid", "invalid_request"},
 	}
 
-	for pattern, errorType := range patterns {
-		if strings.Contains(errMsg, pattern) {
-			return errorType
+	for _, p := range patterns {
+		if strings.Contains(errMsg, p.pattern) {
+			return p.errorType
 		}
 	}
 

telemetry/interceptor.go

@@ -157,12 +157,16 @@ func (i *Interceptor) completeStatement(ctx context.Context, statementID string,
 	i.aggregator.completeStatement(ctx, statementID, failed)
 }
 
-// Close shuts down the interceptor and flushes pending metrics.
+// Close flushes any pending per-connection metrics.
+// Does NOT close the shared aggregator — its lifecycle is managed via
+// ReleaseForConnection, which uses reference counting across all connections
+// to the same host.
 // Exported for use by the driver package.
 func (i *Interceptor) Close(ctx context.Context) error {
 	if !i.enabled {
 		return nil
 	}
 
-	return i.aggregator.close(ctx)
+	i.aggregator.flush(ctx)
+	return nil
 }