Fix unbounded growth of closed-connection tracking in client factories

[sreekanth-db]

Summary

DatabricksHttpClientFactory.closeConnection() left a tombstone entry per (uuid, type) in its static map and never removed it (3 per closed connection), and TelemetryClientFactory tracked closed UUIDs in an unbounded set. Both grew without bound in high-churn / connection-pool deployments, eventually causing OutOfMemoryError. (The HTTP clients themselves were already closed — what leaked was the per-connection map/set entries.)

Changes

• DatabricksHttpClientFactory: closeConnection() now removes the live client entries instead of replacing them with tombstones. Closed-connection detection moved to a WeakHashMap-backed set keyed by the connection UUID, so entries are reclaimed once the connection is unreachable. getClient() returns the ClosedConnectionHttpClient sentinel for closed connections, with a guard for the close-during-create race.
• TelemetryClientFactory: same WeakHashMap-backed set for closedConnectionUuids (was an unbounded newKeySet()).
• Added a regression test asserting the HTTP client map does not accumulate entries across many open/close cycles, while use-after-close still returns the sentinel.

Behavior preserved

The use-after-close protection from #1325/#1333 is intact: after close(), getClient returns the sentinel and getTelemetryClient returns NoopTelemetryClient. The closed-marker is keyed by the connection UUID (unique per connection), so pooled connections sharing a URL are not affected.

Testing

• DatabricksHttpClientTest, TelemetryHttpClientLeakTest, telemetry suites — all pass.
• Verified against a live SQL warehouse: 5 open/close cycles previously left 15 residual entries; now leaves 0.

This pull request and its description were written by Isaac.

[vikrantpuppala]

The only new test asserts factory.liveClientCount() == 0, but liveClientCount() returns instances.size() — a different map from the closedConnections set this PR actually changes. closeConnection() already does instances.remove(...), so this assertion is guaranteed by construction and would pass identically for a strong-keyed unbounded set (i.e. the pre-fix code). Nothing measures closedConnections, and use-after-close returns the sentinel via the first check (getClient line 59), so the test never exercises the WeakHashMap reclamation either.

Why it matters: the specific memory leak this PR claims to fix — unbounded growth of the closed-connection set — could regress and every test would still pass green, shipping the OOM back to high-churn users.

• Suggested fix: add @VisibleForTesting int closedConnectionCount() (returns closedConnections.size()); after many open/close cycles with UUIDs that go out of scope, System.gc(), and a bounded wait/retry, assert the set shrinks well below N (or hold a WeakReference and assert it becomes null).

---

_{Posted by code-review-squad · /full-review · feedback: #code-review-squad-feedback}

[vikrantpuppala]

Both closed-connection sets became Collections.synchronizedSet(newSetFromMap(new WeakHashMap<>())). getClient calls closedConnections.contains(uuid) up to twice per invocation (lines 59 and 65); getTelemetryClient calls it on every telemetry export. The old telemetry set was lock-free ConcurrentHashMap.newKeySet(). Every contains() now acquires a single intrinsic monitor and runs WeakHashMap.expungeStaleEntries() under that lock — serializing exactly the high-churn / connection-pool workload this PR targets.

• Suggested fix: keep reads lock-free — a ConcurrentHashMap-backed structure with periodic pruning, or a lock-free fast path in front of the weak set. At minimum, collapse the two HTTP contains() calls (the post-computeIfAbsent recheck is only needed for a newly-created client).

---

_{Posted by code-review-squad · /full-review · feedback: #code-review-squad-feedback}

[vikrantpuppala]

The new post-computeIfAbsent undo branch (lines 65-71: instances.remove(key, client) + closeQuietly + return sentinel) is never reached by a deterministic test. The new test closes before the second getClient, so the sentinel is returned by the first check (line 59), never this branch. testConcurrentCloseAndGetClientDoesNotLeak is timing-dependent and does not assert liveClientCount()==0.

Why it matters: a bug here (wrong remove key, missing closeQuietly, inverted condition) would leak a live client + socket for a connection racing with close — the exact class from issue #1325 — with no failing test.

• Suggested fix: spy on the computeIfAbsent mapping function to call closeConnection mid-creation, then assert getClient returns ClosedConnectionHttpClient.INSTANCE AND liveClientCount()==0.

---

_{Posted by code-review-squad · /full-review · feedback: #code-review-squad-feedback}

[vikrantpuppala]

WeakHashMap reclaims by key-object reachability, but contains() matches by .equals(). The design works only because getConnectionUuid() returns the same final String connectionUuid field object on every call (DatabricksConnectionContext.java:59,948-949). If a future caller passes an .equals()-equal but distinct String (a pooled/rehydrated/decorator context, or new String(uuid)) after the original context is GC'd, contains() returns false and getClient/getTelemetryClient will silently recreate a live client for a closed connection — defeating issue #1325's use-after-close protection. The regression test reuses the same mock object, so it never exercises the equal-but-distinct-string path.

• Suggested fix: document the invariant on the field ("keys must be the identical String returned by getConnectionUuid()"), and add a test that closes with one context and queries with a second context returning an .equals()-equal-but-distinct string; or key the weak set on the context object itself.

---

_{Posted by code-review-squad · /full-review · feedback: #code-review-squad-feedback}

[vikrantpuppala]

Code Review Squad — Review

Score: 80/100 — MODERATE RISK

Sound, well-reasoned memory-leak fix — the WeakHashMap keyed on the context's own final UUID String reclaims closed-connection markers exactly when the connection can no longer be used (security & language reviewers independently confirmed no correctness/security regression). Two substantive concerns survive verification: the regression test measures the wrong map (instances, already cleared by close) and never the closedConnections set the fix changed, and the switch from lock-free newKeySet() to synchronizedSet(WeakHashMap) adds a global monitor to the getClient/getTelemetryClient hot paths.

---

Could not post inline comment for: F5 — see body below.

Medium Findings

F5 — src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java (line 89) — HTTP and telemetry factories diverge on the concurrent-close contract

The HTTP factory added a post-create recheck/undo for a closeConnection racing with client creation (DatabricksHttpClientFactory.getClient lines 65-71). The telemetry factory has no symmetric recheck after its compute(...) block — it relies only on the "mark closed FIRST" ordering in closeTelemetryClient (lines 167-171). A getTelemetryClient that wins the race against closeTelemetryClient can leave a live TelemetryClient (with a scheduled flush task) in telemetryClientHolders after the connection is marked closed. Sibling factories solving the same problem two different ways is a maintenance/drift hazard — the class comment even asserts symmetry.

Note: severity is Low in the report (marked Medium here only so it posts inline).

• Suggested fix: extract a shared ClosedConnectionRegistry / computeIfLive primitive so both factories inherit identical race semantics, or add the symmetric recheck to the telemetry path.