Fix telemetry HTTP client socket leak preventing CRaC checkpoint (#1333)

## Summary

Fixes #1325 (follow-up to #1233). After `Connection.close()`, delayed
telemetry flush tasks could re-create TELEMETRY HTTP clients that were
never closed, leaking TCP sockets and preventing CRaC checkpoint.

**Two cross-thread race conditions fixed:**

1. **TelemetryClient re-creation**: `getTelemetryClient()` after
`closeTelemetryClient()` created an orphaned `TelemetryClient` with a
periodic flush scheduler that nobody closes. Fixed by tracking closed
connection UUIDs and returning `NoopTelemetryClient`, and by reordering
`closeTelemetryClient()` to export pending collector events before
closing the client.

2. **HTTP client re-creation**: `getClient(ctx, TELEMETRY)` after
`removeClient(ctx)` re-created a `DatabricksHttpClient` via
`computeIfAbsent` that nobody closes. Fixed by adding
`closeConnection()` which permanently marks a connection as closed,
causing `getClient()` to return `null`.

### Files changed

| File | Change |
|------|--------|
| `TelemetryClientFactory.java` | Added `closedConnectionUuids` guard,
reordered close sequence |
| `DatabricksHttpClientFactory.java` | Added `closedConnections` guard,
new `closeConnection()` method |
| `DatabricksConnection.java` | Use `closeConnection()` instead of
`removeClient()` |
| `TelemetryPushClient.java` | Null guard for `getClient()` return value
|
| `TelemetryHttpClientLeakTest.java` | 3 reproduction/verification tests
|
| `RCA_SOCKET_LEAK_TELEMETRY_HTTP_CLIENT.md` | Full RCA with
reproduction and verification plan |

### Test results

- 3085 unit tests pass, 0 failures, 0 errors
- 3 new tests specifically reproduce and verify the fix for both race
conditions

## Test plan

- [x]
`TelemetryHttpClientLeakTest#testGetTelemetryClientAfterCloseReCreatesClient`
— verifies `NoopTelemetryClient` returned after close
- [x]
`TelemetryHttpClientLeakTest#testGetClientReturnsNullAfterCloseConnection`
— verifies `null` returned from HTTP client factory after close
- [x]
`TelemetryHttpClientLeakTest#testCloseTelemetryClientWithPendingCollectorEventsReCreatesClient`
— verifies pending collector events don't cause re-creation
- [x] All 152 existing telemetry tests pass
- [x] Full unit test suite (3085 tests) passes
- [ ] Manual verification with CRaC-enabled JDK (0 sockets after
`Connection.close()`)

This pull request was AI-assisted by Isaac.

---------

Signed-off-by: Gopal Lal <gopal.lal@databricks.com>

Author: gopalldb

NEXT_CHANGELOG.md

@@ -68,6 +68,7 @@ upgrading. These changes do not affect metadata on All-Purpose Clusters.
 - Fixed `?` characters inside SQL comments, string literals, and quoted identifiers being incorrectly counted as parameter placeholders when `supportManyParameters=1`. `SQLInterpolator` now uses `SqlCommentParser` to locate only real placeholders. Fixes #1331.
 - Fixed `MetadataOperationTimeout` not being applied when metadata operations use SHOW commands. Operations like `getTables`, `getSchemas`, and `getColumns` now respect the `MetadataOperationTimeout` connection property instead of hanging indefinitely with no timeout.
 - Reclassify transient server errors to standard SQL states (08S01, 40001) across all Thrift error sites. This ensures UC unavailability and concurrent modification errors surface consistently for better retry handling. Note: Dashboards and branching logic keyed on legacy XXUCC or 42000 must be updated.
+- Fixed telemetry HTTP client socket leak that prevented CRaC checkpoint. After `Connection.close()`, delayed telemetry flush tasks could re-create HTTP clients that were never closed, leaking TCP sockets. Fixes #1325.
 - Fixed client-side enforcement of `maxRows` limit. When `statement.setMaxRows()` is set, `ResultSet.next()` now returns false once the row limit is reached, even if the server returns more rows. Applies to all result types (Thrift, SEA, inline, CloudFetch).
 
 ---

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java

@@ -424,7 +424,7 @@ public void close() throws SQLException {
     TelemetryClientFactory.getInstance().closeTelemetryClient(connectionContext);
     DatabricksClientConfiguratorManager.getInstance().removeInstance(connectionContext);
     DatabricksDriverFeatureFlagsContextFactory.removeInstance(connectionContext);
-    DatabricksHttpClientFactory.getInstance().removeClient(connectionContext);
+    DatabricksHttpClientFactory.getInstance().closeConnection(connectionContext);
     DatabricksThreadContextHolder.clearAllContext();
   }
 

src/main/java/com/databricks/jdbc/dbclient/impl/http/ClosedConnectionHttpClient.java

@@ -0,0 +1,49 @@
+package com.databricks.jdbc.dbclient.impl.http;
+
+import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
+import com.databricks.jdbc.exception.DatabricksHttpException;
+import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
+import java.util.concurrent.Future;
+import org.apache.hc.core5.concurrent.FutureCallback;
+import org.apache.hc.core5.http.nio.AsyncRequestProducer;
+import org.apache.hc.core5.http.nio.AsyncResponseConsumer;
+import org.apache.http.client.methods.CloseableHttpResponse;
+import org.apache.http.client.methods.HttpUriRequest;
+
+/**
+ * Sentinel HTTP client returned for connections that have been closed. All operations throw {@link
+ * DatabricksHttpException} with a clear message, so callers that accidentally store and use a
+ * post-close client get an immediate, diagnosable failure instead of a silent null.
+ *
+ * <p>Symmetric with {@link com.databricks.jdbc.telemetry.NoopTelemetryClient} which serves the same
+ * role for telemetry clients.
+ */
+public final class ClosedConnectionHttpClient implements IDatabricksHttpClient {
+
+  static final ClosedConnectionHttpClient INSTANCE = new ClosedConnectionHttpClient();
+
+  private ClosedConnectionHttpClient() {}
+
+  @Override
+  public CloseableHttpResponse execute(HttpUriRequest request) throws DatabricksHttpException {
+    throw new DatabricksHttpException(
+        "Connection has been closed; HTTP client is no longer usable",
+        DatabricksDriverErrorCode.INVALID_STATE);
+  }
+
+  @Override
+  public CloseableHttpResponse execute(HttpUriRequest request, boolean supportGzipEncoding)
+      throws DatabricksHttpException {
+    throw new DatabricksHttpException(
+        "Connection has been closed; HTTP client is no longer usable",
+        DatabricksDriverErrorCode.INVALID_STATE);
+  }
+
+  @Override
+  public <T> Future<T> executeAsync(
+      AsyncRequestProducer requestProducer,
+      AsyncResponseConsumer<T> responseConsumer,
+      FutureCallback<T> callback) {
+    throw new IllegalStateException("Connection has been closed; HTTP client is no longer usable");
+  }
+}

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClientFactory.java

@@ -7,14 +7,23 @@
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
+import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class DatabricksHttpClientFactory {
   private static final JdbcLogger LOGGER =
       JdbcLoggerFactory.getLogger(DatabricksHttpClientFactory.class);
   private static final DatabricksHttpClientFactory INSTANCE = new DatabricksHttpClientFactory();
-  private final ConcurrentHashMap<SimpleEntry<String, HttpClientType>, DatabricksHttpClient>
+
+  /**
+   * Maps (connectionUuid, type) → HTTP client. On {@link #closeConnection}, real clients are
+   * replaced with {@link ClosedConnectionHttpClient#INSTANCE} tombstones. {@link #getClient}'s
+   * {@code computeIfAbsent} returns the tombstone for closed connections (key already exists) and
+   * creates a real client for new keys. No parallel sets needed — the closed marker lives in the
+   * map itself, bounded by live (uuid, type) pairs. See issue #1325.
+   */
+  private final ConcurrentHashMap<SimpleEntry<String, HttpClientType>, IDatabricksHttpClient>
       instances = new ConcurrentHashMap<>();
 
   private DatabricksHttpClientFactory() {
@@ -29,25 +38,67 @@ public IDatabricksHttpClient getClient(IDatabricksConnectionContext context) {
     return getClient(context, HttpClientType.COMMON);
   }
 
+  /**
+   * Returns an HTTP client for the given connection and type, creating one if needed. For closed
+   * connections, returns the {@link ClosedConnectionHttpClient} sentinel — callers that attempt to
+   * use it get an immediate {@link com.databricks.jdbc.exception.DatabricksHttpException} with a
+   * clear message. Never returns null.
+   */
   public IDatabricksHttpClient getClient(
       IDatabricksConnectionContext context, HttpClientType type) {
     return instances.computeIfAbsent(
         getClientKey(context.getConnectionUuid(), type),
         k -> new DatabricksHttpClient(context, type));
   }
 
+  /**
+   * Permanently closes all HTTP clients for the given connection and replaces them with tombstone
+   * sentinels that reject further use. Called from {@link
+   * com.databricks.jdbc.api.impl.DatabricksConnection#close()}.
+   */
+  public void closeConnection(IDatabricksConnectionContext context) {
+    String uuid = context.getConnectionUuid();
+    for (HttpClientType type : HttpClientType.values()) {
+      SimpleEntry<String, HttpClientType> key = getClientKey(uuid, type);
+      IDatabricksHttpClient old = instances.put(key, ClosedConnectionHttpClient.INSTANCE);
+      if (old != null && !(old instanceof ClosedConnectionHttpClient)) {
+        closeQuietly(old);
+      }
+    }
+  }
+
+  @VisibleForTesting
   public void removeClient(IDatabricksConnectionContext context) {
     for (HttpClientType type : HttpClientType.values()) {
       removeClient(context, type);
     }
   }
 
+  @VisibleForTesting
   public void removeClient(IDatabricksConnectionContext context, HttpClientType type) {
-    DatabricksHttpClient instance =
+    IDatabricksHttpClient instance =
         instances.remove(getClientKey(context.getConnectionUuid(), type));
-    if (instance != null) {
+    if (instance != null && !(instance instanceof ClosedConnectionHttpClient)) {
+      closeQuietly(instance);
+    }
+  }
+
+  /** Resets all state. For test cleanup only. */
+  @VisibleForTesting
+  public void reset() {
+    instances.forEach(
+        (key, client) -> {
+          if (!(client instanceof ClosedConnectionHttpClient)) {
+            closeQuietly(client);
+          }
+        });
+    instances.clear();
+  }
+
+  private static void closeQuietly(IDatabricksHttpClient client) {
+    if (client instanceof DatabricksHttpClient) {
       try {
-        instance.close();
+        ((DatabricksHttpClient) client).close();
       } catch (IOException e) {
         LOGGER.debug("Caught error while closing http client. Error {}", e);
       }

src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java

@@ -34,6 +34,14 @@ public class TelemetryClientFactory {
   @VisibleForTesting
   final Map<String, TelemetryClientHolder> noauthTelemetryClientHolders = new ConcurrentHashMap<>();
 
+  /**
+   * Tracks connection UUIDs that have been closed. {@link #getTelemetryClient} checks this set and
+   * returns {@link NoopTelemetryClient} for closed connections, preventing re-creation of orphaned
+   * TelemetryClients (issue #1325). Growth is bounded by total connections closed over JVM lifetime
+   * (~80 bytes per UUID) — negligible for typical connection pool sizes.
+   */
+  @VisibleForTesting final Set<String> closedConnectionUuids = ConcurrentHashMap.newKeySet();
+
   private final ExecutorService telemetryExecutorService;
   private ScheduledExecutorService sharedSchedulerService;
 
@@ -75,6 +83,11 @@ public static TelemetryClientFactory getInstance() {
   }
 
   public ITelemetryClient getTelemetryClient(IDatabricksConnectionContext connectionContext) {
+    // Reject closed connections to prevent re-creation of orphaned TelemetryClients (issue #1325).
+    String uuid = connectionContext.getConnectionUuid();
+    if (uuid != null && closedConnectionUuids.contains(uuid)) {
+      return NoopTelemetryClient.getInstance();
+    }
     if (!isTelemetryAllowedForConnection(connectionContext)) {
       return NoopTelemetryClient.getInstance();
     }
@@ -137,43 +150,54 @@ public ITelemetryClient getTelemetryClient(IDatabricksConnectionContext connecti
   /**
    * Closes telemetry client for a connection. Thread-safe: computeIfPresent ensures atomic locking,
    * preventing race conditions between connection removal and addition.
+   *
+   * <p>The connection UUID is removed from the open set FIRST to prevent getTelemetryClient() from
+   * re-creating a TelemetryClient during or after the close sequence. Pending TelemetryCollector
+   * events are exported BEFORE the TelemetryClient is closed (inside try-finally), so they are
+   * flushed through the existing client. See GitHub issue #1325.
    */
   public void closeTelemetryClient(IDatabricksConnectionContext connectionContext) {
     String key = TelemetryHelper.keyOf(connectionContext);
     String connectionUuid = connectionContext.getConnectionUuid();
-    // Atomically remove connection and close client if no connections remain for this key
-    telemetryClientHolders.computeIfPresent(
-        key,
-        (k, holder) -> {
-          holder.connectionUuids.remove(connectionUuid);
-          if (holder.connectionUuids.isEmpty()) {
-            closeTelemetryClient(holder.client, "telemetry client");
-            return null;
-          }
-          return holder;
-        });
-    // Atomically remove connection and close client if no connections remain for this key
-    noauthTelemetryClientHolders.computeIfPresent(
-        key,
-        (k, holder) -> {
-          holder.connectionUuids.remove(connectionUuid);
-          if (holder.connectionUuids.isEmpty()) {
-            closeTelemetryClient(holder.client, "unauthenticated telemetry client");
-            return null;
-          }
-          return holder;
-        });
-
-    // Export and remove the TelemetryCollector for this connection
-    TelemetryCollector collector =
-        TelemetryCollectorManager.getInstance().removeCollector(connectionContext);
-    if (collector != null) {
-      // Export any remaining telemetry before removing
-      collector.exportAllPendingTelemetryDetails();
+
+    // Mark connection closed FIRST to prevent getTelemetryClient() from re-creating a
+    // TelemetryClient during or after this close sequence (issue #1325).
+    if (connectionUuid != null) {
+      closedConnectionUuids.add(connectionUuid);
     }
 
-    // Clean up cached connection parameters to prevent memory leaks
-    TelemetryHelper.removeConnectionParameters(connectionContext.getConnectionUuid());
+    // Export pending events inside try-finally so holder cleanup always runs,
+    // even if export throws (F6).
+    try {
+      TelemetryCollector collector =
+          TelemetryCollectorManager.getInstance().removeCollector(connectionContext);
+      if (collector != null) {
+        collector.exportAllPendingTelemetryDetails();
+      }
+    } finally {
+      telemetryClientHolders.computeIfPresent(
+          key,
+          (k, holder) -> {
+            holder.connectionUuids.remove(connectionUuid);
+            if (holder.connectionUuids.isEmpty()) {
+              closeTelemetryClient(holder.client, "telemetry client");
+              return null;
+            }
+            return holder;
+          });
+      noauthTelemetryClientHolders.computeIfPresent(
+          key,
+          (k, holder) -> {
+            holder.connectionUuids.remove(connectionUuid);
+            if (holder.connectionUuids.isEmpty()) {
+              closeTelemetryClient(holder.client, "unauthenticated telemetry client");
+              return null;
+            }
+            return holder;
+          });
+
+      TelemetryHelper.removeConnectionParameters(connectionUuid);
+    }
   }
 
   public ExecutorService getTelemetryExecutorService() {
@@ -216,6 +240,7 @@ public void reset() {
     // Clear the maps
     telemetryClientHolders.clear();
     noauthTelemetryClientHolders.clear();
+    closedConnectionUuids.clear();
 
     // Clear cached connection parameters
     TelemetryHelper.clearConnectionParameterCache();

src/main/java/com/databricks/jdbc/telemetry/TelemetryPushClient.java

@@ -47,6 +47,11 @@ public void pushEvent(TelemetryRequest request) throws Exception {
     IDatabricksHttpClient httpClient =
         DatabricksHttpClientFactory.getInstance()
             .getClient(connectionContext, HttpClientType.TELEMETRY);
+    if (httpClient instanceof com.databricks.jdbc.dbclient.impl.http.ClosedConnectionHttpClient) {
+      // Connection was closed — sentinel returned to prevent socket leaks (issue #1325).
+      LOGGER.debug("Skipping telemetry push: connection has been closed");
+      return;
+    }
     String path =
         isAuthenticated
             ? PathConstants.TELEMETRY_PATH

src/test/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClientTest.java

@@ -25,6 +25,7 @@
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -46,6 +47,12 @@ public class DatabricksHttpClientTest {
   @BeforeEach
   public void setUp() {
     databricksHttpClient = new DatabricksHttpClient(mockHttpClient, mockConnectionManager);
+    DatabricksHttpClientFactory.getInstance().reset();
+  }
+
+  @AfterEach
+  public void tearDown() {
+    DatabricksHttpClientFactory.getInstance().reset();
   }
 
   @Test
@@ -269,8 +276,8 @@ public void testConcurrentClientCreation() throws InterruptedException, Executio
               () -> {
                 IDatabricksConnectionContext connectionContext =
                     mock(IDatabricksConnectionContext.class);
-                when(connectionContext.getConnectionUuid())
-                    .thenReturn(UUID.randomUUID().toString());
+                String uuid = UUID.randomUUID().toString();
+                when(connectionContext.getConnectionUuid()).thenReturn(uuid);
                 when(connectionContext.getHttpMaxConnectionsPerRoute()).thenReturn(100);
                 IDatabricksHttpClient client =
                     DatabricksHttpClientFactory.getInstance().getClient(connectionContext);

src/test/java/com/databricks/jdbc/telemetry/TelemetryClientFactoryTest.java

@@ -308,6 +308,8 @@ private void setupTelemetryHelperMock(MockedStatic<TelemetryHelper> mockedStatic
 
   private void setupMocksForTelemetryClient(IDatabricksConnectionContext context) {
     TelemetryClientFactory.getInstance().closeTelemetryClient(context);
+    // Remove from closed set so subsequent getTelemetryClient() calls succeed in tests
+    TelemetryClientFactory.getInstance().closedConnectionUuids.remove(context.getConnectionUuid());
     TelemetryAuthHelper.setupAuthMocks(context, clientConfigurator);
     Map<String, String> featureFlagMap = new HashMap<>();
     featureFlagMap.put(TELEMETRY_FEATURE_FLAG_NAME, "true");