refactor(conn): Get/Put ref-counting with balancer quarantine#2235
Merged
Conversation
summaryInferred base version: v3.144.2 |
Contributor
There was a problem hiding this comment.
Pull request overview
Refactors the internal connection lifecycle to a reference-counted pool model, where balancer discovery rotates connection references through a one-interval quarantine to ensure dropped endpoints get their pooled gRPC connections closed promptly. This also removes the legacy “connection parking” feature and deprecates the associated TTL options.
Changes:
- Introduces
Pool.Get(ctx)/Put(ctx)ref-counting semantics and removes parking/last-usage tracking. - Adds balancer-side quarantine rotation (
rotateConnectionRefs) and E2E coverage for “node dropped → gRPC closes next discovery”. - Deprecates
WithConnectionTTLoptions and updates changelog/tests accordingly.
Reviewed changes
Copilot reviewed 22 out of 22 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/integration/discovery_test.go | Removes parking-based assertions/options from integration discovery test. |
| options.go | Deprecates public ydb.WithConnectionTTL (no-op). |
| internal/xsync/map.go | Removes size tracking; adds LoadOrStore. |
| internal/xsync/map_test.go | Updates tests for removed map length/clear behavior. |
| internal/xsync/last_usage_guard.go | Removes last-usage tracking implementation. |
| internal/xsync/last_usage_guard_test.go | Removes tests/bench for last-usage guard. |
| internal/xsync/last_usage_guard_start.go | Removes last-usage Start/stop implementation. |
| internal/mock/conn.go | Updates conn mock for removed LastUsage API. |
| internal/endpoint/endpoint.go | Adds endpoint comparator for stable diffing. |
| internal/conn/pool.go | Implements ref-counted Get/Put, removes parking, adjusts removal semantics. |
| internal/conn/pool_test.go | Reworks tests to cover refcount/quarantine behavior and keying. |
| internal/conn/grpc_client_stream.go | Removes last-usage updates from stream operations. |
| internal/conn/grpc_client_stream_test.go | Updates tests for removed last-usage field. |
| internal/conn/conn.go | Removes last-usage/parking APIs; adds useCount. |
| internal/conn/conn_test.go | Removes parking tests. |
| internal/balancer/connections_state.go | Changes All() to return []conn.Conn. |
| internal/balancer/balancer.go | Adds quarantine rotation and releases refs on Close(). |
| internal/balancer/balancer_test.go | Updates tests for new connection types and ref APIs; adds unit test for rotation. |
| internal/balancer/balancer_discovery_close_test.go | Adds E2E test verifying gRPC close after endpoint drop + quarantine cycle. |
| driver.go | Updates bootstrap discovery client to use new pool.Get(ctx, ...). |
| config/config.go | Deprecates ConnectionTTL and config.WithConnectionTTL (no-op). |
| CHANGELOG.md | Documents refactor and TTL deprecations. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![robot-vibe-db[bot]](https://avatars.githubusercontent.com/in/3319454?s=60&v=4){kind=small}
There was a problem hiding this comment.
AI Review Summary
Verdict: ❌ 1 critical issue(s) found
Critical issues
- Critical | Medium:
Pool.Getnon-atomicgetOrCreate+useCount.Add(1)allows a concurrentPutto close a conn between lookup and increment, returning a closed*connwhose useCount is resurrected from 0→1 —internal/conn/pool.go:46-54
Other findings
- Major | Medium:
Pool.Getreturns nil when pool is closed; nil propagates throughrotateConnectionRefsintoapplyDiscoveredEndpointswherecc.Endpoint().Touch()panics with nil dereference —internal/conn/pool.go:47,internal/balancer/balancer.go:214 - Minor | High:
connectionsState.All()now returns the internals.allslice directly instead of a defensive copy, removing a safety net against future callers mutating the slice —internal/balancer/connections_state.go:48-54 - Nit | High:
ctxparameter inPool.Getis unused in the method body (understood to be for API symmetry withPut) —internal/conn/pool.go:46
This review was generated automatically. Critical issues require attention; other findings are advisory.
If this comment was useful, please give it a 👍 — it helps us improve the review bot.
|
Analysis performed by claude, claude-opus-4-6. |
There was a problem hiding this comment.
AI Review Summary
Verdict: ✅ No critical issues found
Critical issues
No critical issues found.
Other findings
- Minor | Medium: Race window between
Pool.Put(removeRef → Delete → Close) and concurrentPool.Get(tryAddRef) — a concurrent Get can obtain a reference to a conn that Put is about to close —internal/conn/pool.go:78 - Minor | High:
conn.removeRef()allowsuseCountto go negative on mismatched Get/Put with no detection —internal/conn/conn.go:241 - Minor | Medium:
Pool.Getretry loop has no upper bound, no context cancellation check, and noisClosedrecheck —internal/conn/pool.go:53 - Nit | Medium:
Pool.Getcontext parameter is documented as "reserved" but prevents callers from cancelling the retry loop —internal/conn/pool.go:48 - Nit | Low:
Balancer.quarantinehas no synchronization primitive — safety depends onrepeater.Stop()being synchronous, which is non-obvious —internal/balancer/balancer.go:79
This review was generated automatically. Critical issues require attention; other findings are advisory.
If this comment was useful, please give it a 👍 — it helps us improve the review bot.
|
Analysis performed by claude, claude-opus-4-6. |
Stop discovery repeater after releasing closeMu to avoid deadlock with in-flight applyDiscoveredEndpoints. Fail fast when discovery bootstrap Get returns nil. Call internal discovery Client.Close from driver wrapper. Export conn.ErrClosedPool and add regression tests. Co-authored-by: Cursor <cursoragent@cursor.com>
Replace the flaky scenario that held closeMu externally (which blocked Close itself) with a stub repeater that verifies closeMu is released before Stop is invoked. Co-authored-by: Cursor <cursoragent@cursor.com>
asmyasnikov
commented
Jul 14, 2026
There was a problem hiding this comment.
AI Review Summary
Verdict: ✅ No critical issues found
Critical issues
No critical issues found.
Other findings
- Minor | Medium: Discovery gRPC connection leak during concurrent
Close()+discoveryConn()—internal/balancer/balancer.go:316 - Minor | Medium:
clearStatedouble-Put on connections in bothquarantineandallis correct but undocumented invariant —internal/balancer/balancer.go:234 - Minor | Medium:
Invoke/NewStreamno longer guard against closed balancer at entry —internal/balancer/balancer.go:433 - Nit | Medium:
Pool.closedusesatomic.Boolbut nearly all reads are underp.mu—internal/conn/pool.go:36
This review was generated automatically. Critical issues require attention; other findings are advisory.
If this comment was useful, please give it a 👍 — it helps us improve the review bot.
|
Analysis performed by claude, claude-opus-4-6. |
kprokopenko
approved these changes
Jul 15, 2026
2 tasks
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #2235 +/- ##
==========================================
+ Coverage 77.77% 77.94% +0.17%
==========================================
Files 453 453
Lines 46084 46204 +120
==========================================
+ Hits 35841 36013 +172
+ Misses 8150 8105 -45
+ Partials 2093 2086 -7
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
kprokopenko
approved these changes
Jul 15, 2026
🌋 SLO Test Results🔴 8 workload(s) tested — 3 workload(s) exceeded failure thresholds
Threshold violations: database-sql-query:
native-query:
native-table:
native-bulk-upsert:
database-sql-table:
native-table-over-query-service:
Generated by ydb-slo-action |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Refactors connection lifecycle to reference-counted
conn.Poolcache and balancer-side quarantine:conn.Pool— reference-counted cache:GetincrementsconnValue.useCount(no gRPC dial),Putdecrements; at zero the entry is removed from the map and gRPC is closed outsidep.muviatryPutBalancer— owns a quarantine for dropped endpoints: a removed node stays alive for one discovery interval; the next successfulListEndpointscallsPuton the quarantine slice and closes gRPC whenuseCountreaches zeronextState— the only place that callsPool.Get/Putfor discovery rotation (notnextConn/ RPC path)AcquireConn,UpdateEndpointUsage, and per-connusageMu;connectionsState.All()returns a defensive copy againBalancer.Close—closeMu.Lock(), setsb.closed, stops repeater,clearState+Puton active/quarantine, closes dedicated discovery gRPC (b.cc)Builds on #2238 (Conn interface / endpoint API cleanup extracted there).
Discovery lifecycle (
nextState)On each successful
ListEndpoints,applyDiscoveredEndpointsruns:Algorithm:
Getevery endpoint in the new discovery result →newActivePutevery connection in the currentquarantine(release previous quarantine)(active, newActive)— old active becomesnewQuarantine, freshly acquired refs becomenewActiveconnValue.useCountis incremented byPool.Getand decremented byPool.Put. When it reaches0, the map entry is deleted underp.muandconn.Closeruns outside the lock.A dropped node is closed on the second discovery after removal: the first discovery only moves it to
newQuarantinewithuseCount ≥ 1and gRPC still alive.Sequence diagram — pool ref-counts & quarantine
Five discovery phases with endpoints A, B (bootstrap), then C added and dropped.
useCountis per endpoint key inconn.Pool.sequenceDiagram autonumber participant LE as ListEndpoints participant B as Balancer participant P as Pool (connValue.useCount) participant cA as conn A participant cB as conn B participant cC as conn C Note over B,P: Phase 1 — bootstrap [A, B]<br/>quarantine=[], active=[] LE->>B: discovery [A, B] B->>P: Get(A), Get(B) P->>cA: new conn, useCount 0→1 P->>cB: new conn, useCount 0→1 B->>P: Put(quarantine ∅) Note over B: newQuarantine=[], newActive=[cA, cB] Note over B,P: Phase 2 — same [A, B]<br/>quarantine=[], active=[cA, cB] LE->>B: discovery [A, B] B->>P: Get(A), Get(B) P->>cA: useCount 1→2 P->>cB: useCount 1→2 B->>P: Put(quarantine ∅) Note over B: newQuarantine=[cA, cB], newActive=[cA, cB] Note over B,P: Phase 3 — same [A, B], release quarantine<br/>quarantine=[cA, cB], active=[cA, cB] LE->>B: discovery [A, B] B->>P: Get(A), Get(B) P->>cA: useCount 2→3 P->>cB: useCount 2→3 B->>P: Put(cA, cB from quarantine) P->>cA: useCount 3→2 P->>cB: useCount 3→2 Note over B: newQuarantine=[cA, cB], newActive=[cA, cB]<br/>net useCount unchanged (2) Note over B,P: Phase 4 — [A, B, C] node C added<br/>quarantine=[cA, cB], active=[cA, cB] LE->>B: discovery [A, B, C] B->>P: Get(A), Get(B), Get(C) P->>cA: useCount 2→3 P->>cB: useCount 2→3 P->>cC: new conn, useCount 0→1 B->>P: Put(cA, cB from quarantine) P->>cA: useCount 3→2 P->>cB: useCount 3→2 Note over B: newQuarantine=[cA, cB], newActive=[cA, cB, cC] Note over B,P: Phase 5 — [A, B] node C dropped<br/>quarantine=[cA, cB], active=[cA, cB, cC] LE->>B: discovery [A, B] B->>P: Get(A), Get(B) P->>cA: useCount 2→3 P->>cB: useCount 2→3 B->>P: Put(cA, cB, cC from quarantine) P->>cA: useCount 3→2 P->>cB: useCount 3→2 Note over cC: useCount stays 1 — gRPC alive<br/>cC in newQuarantine only Note over B: newQuarantine=[cA, cB, cC], newActive=[cA, cB] Note over B,P: Phase 6 — same [A, B], quarantine released, C closed<br/>quarantine=[cA, cB, cC], active=[cA, cB] LE->>B: discovery [A, B] B->>P: Get(A), Get(B) P->>cA: useCount 2→3 P->>cB: useCount 2→3 B->>P: Put(cA, cB, cC from quarantine) P->>cA: useCount 3→2 P->>cB: useCount 3→2 P->>cC: useCount 1→0 P->>cC: Close gRPC, remove from pool Note over B: newQuarantine=[cA, cB], newActive=[cA, cB] Note over B,P: Balancer.Close() — release all refs B->>B: stop discovery repeater B->>P: Put(cA, cB from active) P->>cA: useCount 2→1 P->>cB: useCount 2→1 B->>P: Put(cA, cB from quarantine) P->>cA: useCount 1→0 → Close gRPC P->>cB: useCount 1→0 → Close gRPC Note over B: also closes dedicated discovery gRPC (b.cc)Test plan
go test -race ./internal/conn/...go test -race ./internal/balancer/...TestNextState— ref-count / quarantine unit test fornextState(7 discovery rounds incl. empty cluster)TestBalancerDiscoveryDropClosesGRPC— E2E gRPC mock: node drop → quarantine → close on next discoveryTestPool_GetPut— deadlock, concurrent Get/Put, nil after Release, extra Put no-op, slow close outside mutexTestBalancer_Close— Invoke after close, concurrent close idempotency, discovery during closeTestApplyDiscoveredEndpointsClosedPool,TestNextStateClosedPool,TestConnectionsState_AllReturnsDefensiveCopy