feat(bigtable): add Session lifecycle (Start, Close, ForceClose, readLoop, heartBeatLoop) - #20215
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
Code Review
This pull request implements session-level lifecycle management, heartbeat monitoring, and virtual RPC (vRPC) dispatching for the Bigtable transport layer, including slot-claiming mechanisms and extensive unit tests. The review feedback highlights two key improvements: first, in ForceClose, passing the constructed unavailable error instead of nil to notifyClosed to preserve critical diagnostic information for registered hooks and tracers; second, in Invoke's Send failure path, gating the calls to onSlotDrained and signalQuiescent on a successful drainSlot to prevent redundant invocations during concurrent teardowns.
Comment on lines
+93
to
+99
| desc := "session force closed" | ||
| if req != nil && req.Description != "" { | ||
| desc = "session force closed: " + req.Description | ||
| } | ||
| s.cancelActiveRPCs(unavailable(closeReasonToCause(req), "%s", desc)) | ||
| s.signalQuiescent() | ||
| s.notifyClosed(nil) |
Contributor
There was a problem hiding this comment.
In ForceClose, the error context is lost when calling s.notifyClosed(nil). This means any registered OnClose hooks and the tracer will receive a nil error, making it appear as a clean shutdown even when the session was force-closed due to an error (such as a missed heartbeat or a session error). Passing the constructed unavailable error to notifyClosed preserves this critical diagnostic information.
desc := "session force closed"
if req != nil && req.Description != "" {
desc = "session force closed: " + req.Description
}
closeErr := unavailable(closeReasonToCause(req), "%s", desc)
s.cancelActiveRPCs(closeErr)
s.signalQuiescent()
s.notifyClosed(closeErr)
Comment on lines
+153
to
+162
| s.drainSlot(rpc) | ||
| // v3: drain is the sole "session became free" signal. Fire the | ||
| // pool wake here as well — ReleaseToPool's inExpectedCount guard | ||
| // (see SESSION_POOL_SPEC #6, I5) drops the re-enqueue if | ||
| // OnSessionClosing already dropped this handle in the teardown | ||
| // that the Send failure typically kicks off. | ||
| s.hooks.onSlotDrained() | ||
| if State(s.state.Load()) == StateClosing { | ||
| s.signalQuiescent() | ||
| } |
Contributor
There was a problem hiding this comment.
On Send failure, s.hooks.onSlotDrained() and s.signalQuiescent() are called unconditionally. However, if a concurrent teardown (like ForceClose or handleClose) already reclaimed the slot, s.drainSlot(rpc) will return ok=false. In that case, the teardown path has already handled the session's lifecycle, and calling onSlotDrained or signalQuiescent again is redundant and can lead to inconsistent pool state. These should only be invoked if drainSlot successfully reclaimed the slot.
Suggested change
| s.drainSlot(rpc) | |
| // v3: drain is the sole "session became free" signal. Fire the | |
| // pool wake here as well — ReleaseToPool's inExpectedCount guard | |
| // (see SESSION_POOL_SPEC #6, I5) drops the re-enqueue if | |
| // OnSessionClosing already dropped this handle in the teardown | |
| // that the Send failure typically kicks off. | |
| s.hooks.onSlotDrained() | |
| if State(s.state.Load()) == StateClosing { | |
| s.signalQuiescent() | |
| } | |
| if _, _, ok := s.drainSlot(rpc); ok { | |
| // v3: drain is the sole "session became free" signal. Fire the | |
| // pool wake here as well — ReleaseToPool's inExpectedCount guard | |
| // (see SESSION_POOL_SPEC #6, I5) drops the re-enqueue if | |
| // OnSessionClosing already dropped this handle in the teardown | |
| // that the Send failure typically kicks off. | |
| s.hooks.onSlotDrained() | |
| if State(s.state.Load()) == StateClosing { | |
| s.signalQuiescent() | |
| } | |
| } |
sushanb
force-pushed
the
bigtable-session-vrpc
branch
2 times, most recently
from
July 23, 2026 23:07
16047d6 to
af844b2
Compare
3 tasks
sushanb
commented
Jul 23, 2026
sushanb
left a comment
Contributor
Author
There was a problem hiding this comment.
igor-reviewer-agent
@ 5f7c59b. read the whole thing, looks solid overall. a handful of things:
session_lifecycle.go:716, 730 — Close on ctx-cancel and Send-failure both call s.ForceClose(nil). ForceClose then does setCloseReason(closeReasonLabel(nil)) which is setCloseReason("") — that hits the empty-string guard and bumps tagSessionCloseNoReason every time. by that point Close already stamped the real reason (line 701) so the CAS-once protects the value, but the tag fires spuriously on every legitimate ctx-cancelled or send-failed Close. please pass a real req through, something like:
s.ForceClose(&spb.CloseSessionRequest{
Reason: spb.CloseSessionRequest_CLOSE_SESSION_REASON_ERROR,
Description: "close aborted: " + ctxOrSendErr.Error(),
})session_lifecycle.go:685 — Close accepts StateStarting in its predicate. If Close fires while we're Starting (OpenSession sent, response not yet received), OnActive never fired but OnClosing/OnClose both do. SESSION_SPEC #4's fixed order is OnStart→OnActive→OnClosing→OnClose — is skipping OnActive on this path intentional (spec means "if OnActive fires, ordering holds") or a real gap? if intentional pls add a one-liner to notifyClosing spelling that out, downstream hooks may assume OnActive always precedes OnClosing.
session_lifecycle.go:915 — the go func() in handleGoAway with context.WithTimeout(context.Background(), 30*time.Second) is a per-GOAWAY goroutine detached from any parent lifecycle. If the pool tears down while we're in the drain wait on s.quiescent, this goroutine lives for up to 30s. I think you want the readLoop ctx here so pool shutdown cancels the drain wait cleanly. not blocking but nicer.
session_lifecycle.go:744-755 — the readLoop supervisor goroutine's ForceClose can't unblock the sibling Recv() (docstring says so), so it only helps state-observation. That's an extra goroutine per Session that mostly just watches ctx.Done. Cant you fold this into defer s.ForceClose(...) at the top of readLoop and have the caller cancel the stream ctx to break out of Recv? or is the state-observation-before-Recv-returns window load-bearing somewhere?
session_lifecycle.go:806 — handleOpenSession(_ *spb.OpenSessionResponse) — response body is discarded. when will we start consuming fields off it? add a one-line comment saying "server doesn't populate anything actionable today" or "TODO: plumb foo when server X ships", otherwise the _ reads like an oversight.
session_lifecycle.go:1100 — s/peerInfoExtracter/extractPeerInfo. its a mutator not a noun, and "extracter" is a non-standard spelling.
session_debug.go:288-310 — SessionEvent kinds doc lists close/hb-missed/hb-alive/ctx-done but session_vrpc.go's noteRetryAttempt emits kind "retry". please add it to the doc.
session_lifecycle.go:757 — no recover in readLoop. if handleSessionResponse panics on a malformed frame the goroutine dies silent, session state is stale, no cleanup. i dont have strong feelings about it but a defer that recovers and ForceCloses with a debug tag feels cheap.
test coverage is really thorough btw — TestHandleGoAway_PreservesInFlightRPC + TestHeartBeatLoop_IdleSessionIsNotTornDown are exactly the ones i'd have asked for.
sushanb
force-pushed
the
bigtable-session-lifecycle
branch
from
July 23, 2026 23:15
5f7c59b to
b790a6c
Compare
sushanb
force-pushed
the
bigtable-session-vrpc
branch
from
July 23, 2026 23:44
af844b2 to
c51caae
Compare
sushanb
force-pushed
the
bigtable-session-lifecycle
branch
from
July 23, 2026 23:53
35d3b16 to
1f8a958
Compare
sushanb
added a commit
to sushanb/google-cloud-go
that referenced
this pull request
Jul 23, 2026
Five nits from the igor-reviewer-agent review of PR googleapis#20215: - handleGoAway's per-GOAWAY drain goroutine now derives its 30s timeout from s.stream.Context() instead of context.Background(). Pool teardown cancels the stream ctx, which now unblocks the drain wait cleanly instead of hanging up to 30s on a detached background ctx. - handleOpenSession: doc comment now spells out that the response body is intentionally not consumed (PeerInfo comes from the stream Header, not the reply payload). Documents the '_ *spb.OpenSessionResponse' signature so a future reader doesn't wonder if fields were missed. - peerInfoExtracter → extractPeerInfo. Mutator, not a noun; and "extracter" is a non-standard spelling. Rename in impl + all test callers. - notifyClosing docstring adds the OnActive-skip caveat: if a session tears down while still Starting (Close/ForceClose before OpenSession completes), OnActive never runs but OnClosing/OnClose still do. SESSION_SPEC #4 orders hooks that fire, not hooks that don't. - readLoop gains a deferred panic-recover that records a debug tag and ForceCloses with REASON_ERROR + the panic value in the description. Avoids losing the readLoop goroutine silently if handleSessionResponse panics on a malformed frame. New tag constant tagSessionReadLoopPanic in debug_tracer.go.
sushanb
added a commit
to sushanb/google-cloud-go
that referenced
this pull request
Jul 24, 2026
Five nits from the igor-reviewer-agent review of PR googleapis#20215: - handleGoAway's per-GOAWAY drain goroutine now derives its 30s timeout from s.stream.Context() instead of context.Background(). Pool teardown cancels the stream ctx, which now unblocks the drain wait cleanly instead of hanging up to 30s on a detached background ctx. - handleOpenSession: doc comment now spells out that the response body is intentionally not consumed (PeerInfo comes from the stream Header, not the reply payload). Documents the '_ *spb.OpenSessionResponse' signature so a future reader doesn't wonder if fields were missed. - peerInfoExtracter → extractPeerInfo. Mutator, not a noun; and "extracter" is a non-standard spelling. Rename in impl + all test callers. - notifyClosing docstring adds the OnActive-skip caveat: if a session tears down while still Starting (Close/ForceClose before OpenSession completes), OnActive never runs but OnClosing/OnClose still do. SESSION_SPEC #4 orders hooks that fire, not hooks that don't. - readLoop gains a deferred panic-recover that records a debug tag and ForceCloses with REASON_ERROR + the panic value in the description. Avoids losing the readLoop goroutine silently if handleSessionResponse panics on a malformed frame. New tag constant tagSessionReadLoopPanic in debug_tracer.go.
sushanb
force-pushed
the
bigtable-session-lifecycle
branch
from
July 24, 2026 00:11
1f8a958 to
a0e1708
Compare
sushanb
added a commit
to sushanb/google-cloud-go
that referenced
this pull request
Jul 24, 2026
Applies the same nits accepted on PRs googleapis#20211/googleapis#20213/googleapis#20215 to the sessionz working branch. See per-PR commits cf200ae / c51caae / 1f8a958 for rationale. Plus the ForceClose(nil) spurious-tag fix from 35d3b16.
sushanb
force-pushed
the
bigtable-session-lifecycle
branch
from
July 24, 2026 00:26
a0e1708 to
922cf2a
Compare
sushanb
added a commit
to sushanb/google-cloud-go
that referenced
this pull request
Jul 24, 2026
Five nits from the igor-reviewer-agent review of PR googleapis#20215: - handleGoAway's per-GOAWAY drain goroutine now derives its 30s timeout from s.stream.Context() instead of context.Background(). Pool teardown cancels the stream ctx, which now unblocks the drain wait cleanly instead of hanging up to 30s on a detached background ctx. - handleOpenSession: doc comment now spells out that the response body is intentionally not consumed (PeerInfo comes from the stream Header, not the reply payload). Documents the '_ *spb.OpenSessionResponse' signature so a future reader doesn't wonder if fields were missed. - peerInfoExtracter → extractPeerInfo. Mutator, not a noun; and "extracter" is a non-standard spelling. Rename in impl + all test callers. - notifyClosing docstring adds the OnActive-skip caveat: if a session tears down while still Starting (Close/ForceClose before OpenSession completes), OnActive never runs but OnClosing/OnClose still do. SESSION_SPEC #4 orders hooks that fire, not hooks that don't. - readLoop gains a deferred panic-recover that records a debug tag and ForceCloses with REASON_ERROR + the panic value in the description. Avoids losing the readLoop goroutine silently if handleSessionResponse panics on a malformed frame. New tag constant tagSessionReadLoopPanic in debug_tracer.go.
Contributor
Author
|
@igor-reviewer-agent thanks — real
Skipped: OnActive-skip spec question (turned into the doc-note above); OpenSessionResponse discard clarification (also doc-note); readLoop-supervisor-fold — left as-is since the state-observation-before-Recv-returns window is load-bearing for the pool's stuck-session monitor. |
Introduces Session.Invoke and the one-in-flight vRPC dispatch that
routes server frames back to callers. Stacks on the sessionDebug bundle
(bigtable-session-debug); the counter/tracer/event writers land here.
New file: session_vrpc.go (~420 LOC)
- Session.Invoke — public entry point. Serializes claimSlot → Send →
awaitInvokeResult; deliver-vs-cancel branching lives in routeVRPCFrame.
- Slot lifecycle: claimSlot / drainSlot / markCancelled / activeVRPC —
slotMu-guarded, one vRPC in flight per session.
- handleVRPCResponse / handleVRPCErrorResponse + shared routeVRPCFrame —
Java-parity drain path with nil-active-VRPC + id-mismatch guards.
- deliver — writes to the caller's resultChan.
- cancelActiveRPCs — session-teardown path used by ForceClose /
handleGoAway / heartbeat watchdog.
- ForceClose — minimal port (transitionTo(Closed) → cancelActiveRPCs →
signalQuiescent). The full lifecycle-flavored ForceClose (setCloseReason
+ notifyClosing + notifyClosed + stream cleanup) lands with the
session_lifecycle follow-up.
Edits to session.go
- Add heartbeatWake chan struct{} field (cap-1 non-blocking) +
make(...) in NewSession — reactive watchdog wake channel.
- Add OnSlotDrained func() to SessionHooks + onSlotDrained() dispatcher.
Sole "session became free" signal; consumers wake pool waiters.
- Add Send(req) — mutex-guarded stream.Send wrapper; used by both this
PR's Invoke and the follow-up lifecycle Start.
- Add resetHeartbeatDeadline + wakeHeartbeatLoop — atomic-based reactive
wake pair; consumed by the follow-up lifecycle heartBeatLoop.
Edits to session_test.go
- Test fixtures: fakeStream / fakeDesc / newFakeStream / newRoundTripDesc /
makeActive / waitFor. Consumed by both this PR's TestInvoke*/TestHandleVRPC*
and the follow-up lifecycle tests.
- newTestSession signature change (stream, hooks); pre-existing 1-arg
call sites now use newDefaultTestSession shortcut.
Edits to debug_tracer.go
- tagSessionVRPCCancelledDrained — bookkeeping tag when a server response
drains a slot whose caller already ctx.Done'd.
Stack:
- Prereq: PR googleapis#20211 (bigtable-session-debug) — Session debug surface.
- Follow-up: session_lifecycle.go (Start, Close, readLoop, heartBeatLoop,
handleSessionResponse / handleOpenSession / handleGoAway / handleClose /
handleErrorResponse / peerInfoExtracter). Will extend ForceClose to its
full lifecycle-shaped body.
Address gemini-code-assist review on session_vrpc.go:153-162: the Send- failure branch was firing hooks.onSlotDrained() and signalQuiescent() unconditionally, ignoring drainSlot's ok return. When concurrent teardown (cancelActiveRPCs / ForceClose) beats the Invoke path to freeing the slot, the teardown path is responsible for pool notification via OnClosing/OnClose — firing onSlotDrained on the losing path violates the SessionHooks contract that reserves it for the drain-succeeded path (SESSION_SPEC #5/#10). Wrap both calls in `if _, _, ok := s.drainSlot(rpc); ok { ... }`. Idempotence of the underlying ops isn't the point — the design says teardown paths do NOT fire onSlotDrained, and this brings the Send-failure branch back in line.
…ents Twelve doc-comment references to Java SessionImpl line-number citations and "Java parity" phrasing rewritten to describe Go behavior directly. No code change.
Five nits from the igor-reviewer-agent review of PR googleapis#20213: - Return stillActive bool from markCancelled and use it in awaitInvokeResult so the ctx.Done branch does one slotMu take instead of two (markCancelled + recordCtxDone's activeVRPC). - Move ctx to the first arg of buildInvokeRequest per Go convention. - Drop the dead `_ = rpc` in processResult; remove rpc from the signature (only awaitInvokeResult calls it, no rpc reads today). - Delete the deliver() one-line wrapper — inline the buffered send at both call sites (routeVRPCFrame's non-cancelled branch, cancelActiveRPCs's non-cancelled branch). Two-word "cap-1, drainSlot serialized this write" comment replaces the paragraph docstring. - Replace time.Sleep(20ms) in TestInvoke_ForceCloseWhileSending_ BoundedReturn with a deterministic waitFor on activeVRPC() == nil, so the test doesn't rely on scheduler cooperation. Also converts the two remaining string-literal call sites in noteRetryAttempt and recordCtxDone to the typed SessionEventKind constants (SessionEventRetry, SessionEventCtxDone) landed in the prior PR-1 nit pass.
Java-bigtable's SessionImpl.java sets nextHeartbeat = now + heartbeatInterval at every reset (SessionImpl.java:440 in startRpc, :621 in handleHeartBeatResponse). Our 3x multiplier was a Go-only invention — tolerated 2 additional missed heartbeats before force-closing, but that extra grace comes at the cost of stale-session detection latency and diverges from the reference implementation. Under the new 1x deadline, a healthy session with regular server frames never trips the watchdog (each frame's resetHeartbeatDeadline pushes the deadline forward). Only sessions whose server has genuinely stopped sending frames for one full interval fire ForceClose — the actual "heartbeat missed" signal the watchdog exists to detect. Follow-up in session_lifecycle.go's heartBeatLoop will adjust its lastFrameAge = 3*interval - remaining arithmetic to match.
Three golint vet errors reported by CI on PR googleapis#20213: - session_vrpc.go:202: var virtRpc should be virtRPC (initialisms uppercase per Go convention). - session_vrpc_test.go:37,138: runVRpcCapture and runInvoke had (t *testing.T, ctx context.Context, ...) — vet requires ctx as the first parameter. Rename to (ctx, t, ...) and update all callers (6 total across the file). No behavior change.
Reverts commit 971e4ec. After tracing the race gemini flagged against the parent sessionz branch where the full pool machinery is in scope, the guard is redundant: - notifyClosing runs synchronously before cancelActiveRPCs on every teardown path (ForceClose, handleClose, handleGoAway, heartbeat miss). By the time cancelActiveRPCs wins drainSlot, the pool has already flipped inExpectedCount=false and removed the handle from afe.sessions via OnSessionClosing. - ReleaseToPool early-returns on !sh.inExpectedCount, so an unguarded onSlotDrained on the ok=false iteration is a no-op re-enqueue. - signalFree wakes a waiter which re-picks via ReadyAfes — the closing session is already out of that set, so the waiter never receives the dying handle. Emit stays dumb here; the pool absorbs dupes at ReleaseToPool.
mutianf
reviewed
Jul 27, 2026
| // failures) treats a failed OpenSession send the same as any | ||
| // other transport-side loss. errors.Is still resolves the | ||
| // underlying send error via sessionErr.Unwrap. | ||
| return unavailable(err, "session OpenSession request failed: %v", err) |
Contributor
There was a problem hiding this comment.
should we always wrap this in unavailable error? for example if the server reject open session because of unimplemented, we should fallback to unary path.
Contributor
Author
There was a problem hiding this comment.
Applied in c23501d — added TODO(sushanb) at this call site. For the scope of this PR we keep the blanket Unavailable wrap; the fallback to unary path when the server rejects OpenSession with codes.Unimplemented is real but requires plumbing on the client-side classic-vs-session router, so it lands separately.
| case <-ctx.Done(): | ||
| s.ForceClose(&spb.CloseSessionRequest{ | ||
| Reason: spb.CloseSessionRequest_CLOSE_SESSION_REASON_USER, | ||
| Description: "close context deadline exceeded during drain", |
Contributor
There was a problem hiding this comment.
nit: ctx.Done() could also be user cancel? so it's not necessarily because of deadline exceeded during drain?
Contributor
Author
There was a problem hiding this comment.
Applied in c23501d — swapped the hard-coded 'close context deadline exceeded during drain' for fmt.Sprintf("close context done during drain: %v", ctx.Err()) so operators see the concrete cause (Canceled vs DeadlineExceeded) in the sessionz close-reasons stream.
| default: | ||
| recordDebugTag(tagSessionUnknownResponse) | ||
| s.debugf("received SessionResponse with unknown payload type %T", p) | ||
| return |
Contributor
There was a problem hiding this comment.
should we reset the heartbeat deadline if response fall under this branch?
Contributor
Author
There was a problem hiding this comment.
No — the current default-branch behavior (no reset on unknown-payload) is intentional. See the doc comment above the switch: 'unknown frames do NOT reset the heartbeat watchdog, so a misbehaving server cannot keep the watchdog satisfied with junk payloads.' If we reset on unknown, a server that trickles unrecognized frames indefinitely keeps the watchdog fresh forever and stalled vRPCs are never detected. Keeping as-is.
Contributor
Author
There was a problem hiding this comment.
Looped back on this after weighing your original ask against the trade-off. Keeping current behavior (unknown frames do NOT reset). SESSION_SPEC #7 explicitly enshrines the rule with the same rationale I gave: 'a rogue future frame type would mask a broken stream.' Java parity agrees — SessionImpl.handleUnknownResponseMessage (SessionImpl.java:718-722) also does not reset the heartbeat.
The forward-compat rationale I was tempted to lean on is also weaker than it looks: the watchdog is only armed while a vRPC is in-flight, and during an in-flight vRPC the server MUST be sending Heartbeat frames on its keepalive cadence regardless of what other traffic is going out. A new-variant frame arriving instead of a heartbeat within the 100ms interval is itself a signal worth surfacing — that's a server-side keepalive bug, not a client-side forward-compat gap.
Added a Java-parity note + rationale to the docstring so a future reader doesn't attempt the same flip I did. TestHandleSessionResponse_UnknownDoesNotResetDeadline pins the behavior against silent regression.
| defer cancel() | ||
| _ = s.Close(ctx, &spb.CloseSessionRequest{ | ||
| Reason: spb.CloseSessionRequest_CLOSE_SESSION_REASON_GOAWAY, | ||
| Description: "client teardown after server GOAWAY", |
Contributor
There was a problem hiding this comment.
nit: I don't think we're tearing down the client? maybe just
Suggested change
| Description: "client teardown after server GOAWAY", | |
| Description: "session closing after server GOAWAY", |
Contributor
Author
There was a problem hiding this comment.
Applied in c23501d — swapped to your suggested text.
| // and worth flagging. | ||
| func isAbnormalCloseReason(reason string) bool { | ||
| switch reason { | ||
| case "StreamEnd:EOF", "StreamEnd:Canceled", |
Contributor
There was a problem hiding this comment.
should this be based on the state of the session? e.g. if it didn't transition from start -> ready -> closing -> wait server close -> close, then it's abnormal close?
Contributor
Author
There was a problem hiding this comment.
For Claude, please add a TODO to fix this. We will remove this and add the logic as @mutianf stated.
Contributor
Author
There was a problem hiding this comment.
Considered a state-based classifier but kept the reason-string approach for now. The current impl already captures state via CloseReason() being CAS-once-stamped at the moment of transition — every path that transitions out of Ready stamps a specific reason (Close→'User', ForceClose→closeReasonLabel(req), handleGoAway→'GoAway', heartbeat miss→'MissedHeartbeat', handleErrorResponse→'Error', handleClose→streamEndReason(err)). So 'the state transitions were clean' and 'the reason string is in the whitelist' encode the same signal — the whitelist just gives us finer-grained per-reason attribution downstream (sessionz's close-reasons breakdown treats StreamEnd:EOF vs StreamEnd:Unavailable as distinct buckets).
State-based would collapse those to one signal and force downstream consumers to re-derive the specific reason from state history. Would be happy to refactor as a follow-up if downstream code ever needs it — flagging in the memory as a design consideration.
| // only enforced while at least one VRPC is in flight: the server emits | ||
| // Heartbeats during long-running VRPCs, so an idle session legitimately | ||
| // receives no heartbeats and must not be torn down. | ||
| func (s *Session) heartBeatLoop(ctx context.Context) { |
Contributor
There was a problem hiding this comment.
nit:
Suggested change
| func (s *Session) heartBeatLoop(ctx context.Context) { | |
| func (s *Session) heartbeatLoop(ctx context.Context) { |
Contributor
Author
There was a problem hiding this comment.
Applied in c23501d — renamed heartBeatLoop → heartbeatLoop across all seven files (session.go, session_debug.go, session_lifecycle.go, session_lifecycle_test.go, session_pool_lifecycle.go, session_pool_scaling.go, session_tracer_test.go on the parent sessionz branch).
| addr := p.Addr.String() | ||
| s.remoteAddr.Store(&addr) | ||
| } | ||
| s.tracer.recordOpen(context.Background(), s.peerInfo.Load(), nil) |
Contributor
There was a problem hiding this comment.
if extractPeerInfo failed we won't call s.peerInfo.Store(), will s.peerInfo.Load() cause issue?
Contributor
Author
There was a problem hiding this comment.
Checked — all downstream call sites are nil-safe:
session_tracer.go:recordOpen— its docstring explicitly says 'may be nil if the handshake failed before headers arrived', and itspeerInfoLabels(peerInfo)helper handles nil (returns unknown/unspecified for both labels).session.go:PeerInfo()— just returns the atomic Load result, callers guard.session.go:AfeID()— explicitly nil-guardedif p := s.peerInfo.Load(); p != nil.session_lifecycle.go:129(notifyClosed) — passes into recordClose which is also nil-safe per its docstring.
So the pattern is 'store when we have it, nil-check at every read.' Adding a defensive zero-store on failure would prevent downstream from distinguishing 'PeerInfo unknown (server didn't send it)' from 'PeerInfo intentionally zero (invalid ApplicationFrontendId=0)'. Keeping the nil-tolerant read pattern.
…calation Three tightly-scoped changes on session_vrpc.go per mutianf's review on googleapis#20213: 1. Decode failure classified as StateTransportFailure (was StateServerResult). A malformed payload means we couldn't consume the wire bytes, so no server-application-level result exists to reason about; a retry (potentially on another session/AFE) may succeed if the bad payload was a one-off from a specific backend. 2. Broadened the vRPC-frame state whitelist to accept WaitServerClose. The drain window between our CloseSession send and the server's EOF legitimately delivers late VirtualRpc responses; treating them as wrong-state produced false protocol errors during graceful teardown. 3. Replaced silent debug-tag drops on the three routeVRPCFrame guard paths with a new SessionEventProtocolError event kind + escalation: - wrong state (New/Starting/Closed): tear down via cancelActiveRPCs - nil active vRPC (post-cancel race): log, drop, keep session - rpc_id mismatch (server desync): tear down via cancelActiveRPCs Silent drops on protocol violations left the current caller waiting for a response that would never arrive; escalating to cancelActiveRPCs delivers a terminal error so the retry oracle picks another session/AFE instead of hitting ctx timeout. Three new tests cover the routing behaviors: - TestHandleVRPCResponse_IDMismatch_TearsDownSession - TestHandleVRPCResponse_WrongState_TearsDownSession - TestHandleVRPCResponse_NilRPC_DoesNotTearDown Legacy behavior (silent drop into debug tag) preserved only for the nil-active-vRPC case because that path is a documented race between local ctx.Done cancel and a legitimate late server response.
…rr classification Follow-up to f4ab4fa addressing session-reviewer's two spec violations + Igor's nits. Same code as feat/bigtable-sessionz-debug 0a51296. Behavior deltas from f4ab4fa: 1. decodeErr classification REVERTED from StateTransportFailure back to StateServerResult. The prior flip widened the retry surface: an idempotent write whose response bytes were unreadable would have retried under transport-failure rules, risking double-apply. Server sent a response → request likely committed even if we can't parse the reply. Callers with a genuine idempotency token can override at the retry-oracle level. 2. Wrong-state and rpc_id-mismatch paths now call ForceClose( CLOSE_SESSION_REASON_ERROR) instead of cancelActiveRPCs alone. cancelActiveRPCs is a teardown-companion: it drains the slot but does NOT fire OnClosing/OnClose, so a session escalated via cancelActiveRPCs-alone stays Ready and orphaned from the pool's AFE routing set (once that pool-side machinery lands in a later PR — on this PR it's a no-op distinction but keeps the discipline aligned with the future pool wiring). ForceClose additionally marks the session Closed so subsequent Invokes on the same session-handle see state!=Ready and return Uncommitted for the retry oracle. 3. Split SessionEventProtocolError vs SessionEventLateFrame: - ProtocolError: wrong-state (New/Starting/Closed) + id-mismatch (genuine server desync). Both escalated via ForceClose. - LateFrame: nil-active-vRPC — legit ctx.Done race, session stays Ready. Separate kind so operators grep-filtering "protocol-error" don't see benign late-frame drops swamping desync alerts. Test coverage expanded to 5 tests: * TestHandleVRPCResponse_IDMismatch_TearsDownSession — asserts status.Code == Unavailable, state == Closed, event ring contains ProtocolError. * TestHandleVRPCResponse_WrongState_TearsDownSession — uses StateStarting (Closed early-returns from ForceClose CAS; the caller in Closed state is by definition already unblocked). * TestHandleVRPCResponse_NilRPC_DoesNotTearDown — session stays Ready, event ring has LateFrame, ProtocolError absent (kind split regression pin). * TestHandleVRPCResponse_ClosingState_AcceptsMatchingFrame — pins Closing in the whitelist (drain window must accept matching frames). * TestHandleVRPCResponse_WaitServerCloseState_AcceptsMatchingFrame — pins the whitelist widening the parent commit introduced.
…Loop, heartBeatLoop) Third and final PR in the Session core stack. Adds the lifecycle orchestration that drives a Session from Start through teardown, plus the readLoop that dispatches server frames to the vRPC handlers landed in bigtable-session-vrpc. Stacks on bigtable-session-vrpc; deletes the minimal ForceClose stub that PR shipped, replacing it with the full lifecycle-shaped body. New file: session_lifecycle.go (~640 LOC) - Session.Start(ctx, OpenSessionRequest) — transitions New→Starting, Sends the OpenSession frame, fires onStart, spawns readLoop + heartBeatLoop. Wraps a failed Send as codes.Unavailable so retry plumbing treats pre-wire OpenSession loss the same as any other transport-side loss. - ForceClose — full body: transitionTo(Closed), setCloseReason, notifyClosing (once), cancelActiveRPCs, signalQuiescent, notifyClosed. - Close(ctx, CloseSessionRequest) — graceful drain: Ready→Closing, waits on quiescent, sends CloseSession, transitions to WaitServerClose, arms the pool's stuck-session sweep to eventually ForceClose if the server never confirms. - notifyClosing / notifyClosed — once-guarded hook dispatchers with the strict onClosing-precedes-onClose ordering enforced. - readLoop(ctx) — Recv-loop that dispatches every SessionResponse variant via handleSessionResponse, drives handleClose on stream termination, records msgsRecv counters + resets heartbeat deadline on every recognized frame. - handleSessionResponse — dispatch switch to handleOpenSession / handleVRPCResponse / handleErrorResponse / handleSessionParameters / handleGoAway / handleSessionRefreshConfig. Heartbeat frames update the deadline via resetHeartbeatDeadline; the unknown-payload branch records a debug tag but doesn't reset the deadline (so a misbehaving server can't keep the watchdog satisfied with junk). - handleOpenSession — Starting→Ready transition, PeerInfo extract from the bidi header (synchronous, matches Java's onHeaders synchrony), fires onActive. - handleGoAway / handleClose / handleErrorResponse / handleSessionParameters / handleSessionRefreshConfig — protocol-level handlers. - heartBeatLoop — Timer + heartbeatWake reactive-wake pair. Only enforces the deadline while a vRPC is in flight (idle sessions legitimately receive no server heartbeats). See SESSION_SPEC.md #7. - peerInfoExtracter — parses the bigtable-peer-info header, stamps s.peerInfo atomically before onActive fires. - closeReasonLabel / closeReasonToCause — CloseSessionRequest.Reason → string / sentinel error mapping for close-reason attribution. New file: session_lifecycle_test.go (~660 LOC) — 30+ tests covering Start, ForceClose, Close, readLoop, handleSessionResponse dispatch, handleOpenSession + peerInfo parsing, handleGoAway, handleClose, handleErrorResponse (rpc_id=0 harmlessly drops via routeVRPCFrame guards), heartBeatLoop (reactive wake, idle-gate, missed-heartbeat ForceClose). Edits to session_vrpc.go - Delete the minimal ForceClose stub introduced in the prior PR — the full body lives in session_lifecycle.go now. Edits to session_test.go - Add hookCounts helper (start/active/close callback counters). - Add setSlotForTest helper (seed the in-flight slot from tests). Stack: - googleapis#20211 — Session debug surface - googleapis#20213 — Session vRPC dispatch + slot lifecycle - This PR — Session lifecycle
Four doc-comment references to java-bigtable (handleOpenSession "onHeaders synchrony", handleGoAway "onSessionGoAway" + "Java parity: SessionImpl's handleGoAwayResponse", peerInfoExtracter "Base64.getUrlDecoder") rewritten to describe the Go behavior directly. No code change.
Address igor-reviewer: two ForceClose(nil) calls in Session.Close
(the ctx.Done drain-timeout branch and the Send-failure branch) were
triggering spurious tagSessionCloseNoReason via setCloseReason("")
even though Close had already stamped the graceful reason via the
earlier setCloseReason(closeReasonLabel(req)). The debug tag fires
before setCloseReason's CAS, so the tag hit even when the CAS was a
no-op — noise on every drain-timeout or close-send-failure.
Synthesize a specific CloseSessionRequest at each site:
- ctx.Done during drain: REASON_USER "close context deadline
exceeded during drain"
- Send(closeReq) failure: REASON_ERROR "failed to send close
session request"
Five nits from the igor-reviewer-agent review of PR googleapis#20215: - handleGoAway's per-GOAWAY drain goroutine now derives its 30s timeout from s.stream.Context() instead of context.Background(). Pool teardown cancels the stream ctx, which now unblocks the drain wait cleanly instead of hanging up to 30s on a detached background ctx. - handleOpenSession: doc comment now spells out that the response body is intentionally not consumed (PeerInfo comes from the stream Header, not the reply payload). Documents the '_ *spb.OpenSessionResponse' signature so a future reader doesn't wonder if fields were missed. - peerInfoExtracter → extractPeerInfo. Mutator, not a noun; and "extracter" is a non-standard spelling. Rename in impl + all test callers. - notifyClosing docstring adds the OnActive-skip caveat: if a session tears down while still Starting (Close/ForceClose before OpenSession completes), OnActive never runs but OnClosing/OnClose still do. SESSION_SPEC #4 orders hooks that fire, not hooks that don't. - readLoop gains a deferred panic-recover that records a debug tag and ForceCloses with REASON_ERROR + the panic value in the description. Avoids losing the readLoop goroutine silently if handleSessionResponse panics on a malformed frame. New tag constant tagSessionReadLoopPanic in debug_tracer.go.
Two follow-ups after the java-bigtable heartbeat cross-check:
1) handleSessionParameters: drop 3x multiplier on nextHeartbeatDeadline
(already done in resetHeartbeatDeadline in the prior PR-2 commit).
Now 1 * interval, matching Java's SessionImpl.java:440 which sets
nextHeartbeat = now + heartbeatInterval on every reset. Wake-comment
updated to reflect the new arithmetic.
2) heartBeatLoop:
- Delete the `active := 1 // multiplex=1; kept as %d so a future >1
stays greppable` code smell. Inline `1` was documenting itself;
the variable existed only for its comment. Drop it — when
multiplex > 1 lands, the caller will read `s.activeVRPCCount()`
(or similar), not this stale scaffolding.
- `lastFrameAge = interval - remaining` (was `3*interval - remaining`)
to reflect the new 1x deadline.
- Debugf/recordEvent format strings drop the `in_flight=%d` field
that was carrying the constant `1`.
Test adjustments:
- TestHandleSessionParameters_UpdatesIntervalAndDeadline: expect
deadline ~= before+interval (1x) instead of before+3*interval.
- TestHeartBeatLoop_HeartbeatsKeepInflightVRPCAlive: seed deadline
at now+30ms (1x) instead of now+90ms (was 3x).
Four fixes from mutianf's review on googleapis#20215: 1. Start's Unavailable wrap on Send failure — added TODO(sushanb) for distinguishing codes.Unimplemented so the client can fall back to the unary (non-session) path when the server rejects OpenSession with Unimplemented. Today every Send failure is folded into Unavailable, which retries indefinitely against a server that will never accept the RPC. Kept the current wrap for this PR's scope. 2. Close's drain-timeout ForceClose description — was "close context deadline exceeded during drain"; ctx.Done can also fire from user cancel, not only deadline. Now formats the actual ctx.Err() into the description so operators see the concrete cause. 3. handleGoAway's Close description — "client teardown after server GOAWAY" implied full client shutdown; the session is what's closing, not the client. Now "session closing after server GOAWAY". 4. heartBeatLoop → heartbeatLoop — Go convention is single-lowercase for the interior word ("heartbeat" is one word). Cross-file rename including test call sites and one doc-comment reference. Comments 3 (unknown-payload heartbeat reset), 5 (state-based abnormal- close classification), and 7 (peerInfo.Load nil-safety) addressed via inline reply — existing behavior is intentional / already nil-safe. See individual reply threads for justification.
sushanb
force-pushed
the
bigtable-session-lifecycle
branch
from
July 27, 2026 15:42
922cf2a to
c23501d
Compare
sushanb
added a commit
to sushanb/google-cloud-go
that referenced
this pull request
Jul 27, 2026
Same four fixes I landed on PR googleapis#20215 (bigtable-session-lifecycle c23501d), backported to sessionz so the two branches don't drift: 1. Start's Unavailable wrap on Send failure — added TODO(sushanb) for distinguishing codes.Unimplemented so the client can eventually fall back to the unary (non-session) path when the server rejects OpenSession with Unimplemented. Today every Send failure is folded into Unavailable, which retries indefinitely against a server that will never accept the RPC. Kept the wrap for now. 2. Close's drain-timeout ForceClose description — ctx.Done can fire from user cancel, not only deadline. Now formats the actual ctx.Err() into the Description so operators see the concrete cause. 3. handleGoAway's Close description — "client teardown" was too broad (we're only closing the session, not the client). Now "session closing after server GOAWAY". 4. heartBeatLoop → heartbeatLoop across seven files (session.go, session_debug.go, session_lifecycle.go, session_lifecycle_test.go, session_pool_lifecycle.go, session_pool_scaling.go, session_tracer_test.go). Go convention: "heartbeat" is one word. Comments 3 (unknown-payload heartbeat reset), 5 (state-based abnormal- close classification), 7 (peerInfo.Load nil-safety) addressed via inline reply on the PR — existing behavior is intentional and/or already nil-safe at every read site.
sushanb
added a commit
to sushanb/google-cloud-go
that referenced
this pull request
Jul 27, 2026
Three small follow-ups from the sessionz reviewer pass after mutianf review on googleapis#20215: 1. handleSessionResponse docstring — expanded the "unknown frames do NOT reset the heartbeat" note with the Java-parity cite (SessionImpl.handleUnknownResponseMessage also does not reset) and the honest rationale that the watchdog is only armed while a vRPC is in-flight, so a new-variant frame arriving instead of a heartbeat within the interval is itself a server-side keepalive bug worth surfacing. Prevents a future reader from making the forward-compat argument I initially tried in reply to mutianf. 2. TestHeartBeatLoop_* → TestHeartbeatLoop_* — four test function names in session_lifecycle_test.go, one comment cite in session_tracer_test.go, and two comment cites in retry_semantics_test.go (leftovers from the earlier heartBeatLoop → heartbeatLoop identifier rename that missed the CamelCase test names). ErrUnavailableHeartBeatMissed left as-is — exported symbol, treat separately in a later API-shape review. 3. TODO on isAbnormalCloseReason — noted mutianf's suggestion to move to a state-based classifier (New→Ready→Closing→WaitServerClose→ Closed is the "clean" path, anything else is abnormal). Current reason-string scheme encodes state indirectly via CAS-once CloseReason at each transition site and gives finer per-reason attribution for sessionz's close-reasons breakdown, but the whitelist has to be kept in lockstep with every new closeReasonLabel case. Follow up when we add another reason or a downstream consumer needs the state history directly.
Three small follow-ups after the sessionz reviewer pass on the earlier mutianf-review commit: 1. handleSessionResponse docstring — expanded the "unknown frames do NOT reset the heartbeat" note with the Java-parity cite (SessionImpl.handleUnknownResponseMessage also does not reset) and the honest rationale that the watchdog is only armed while a vRPC is in-flight, so a new-variant frame arriving instead of a heartbeat within the interval is itself a server-side keepalive bug worth surfacing. Prevents a future reader from making the forward-compat argument I initially tried in reply to mutianf's comment 3. 2. TestHeartBeatLoop_* → TestHeartbeatLoop_* — four test function names in session_lifecycle_test.go, leftover from the earlier heartBeatLoop → heartbeatLoop identifier rename that missed the CamelCase test names. ErrUnavailableHeartBeatMissed left as-is — exported symbol, treat separately in a later API-shape review. 3. TODO on isAbnormalCloseReason — noted mutianf's suggestion (comment 5 on googleapis#20215) to move to a state-based classifier (New→Ready→Closing→WaitServerClose→Closed is the "clean" path). Current reason-string scheme encodes state indirectly via CAS-once CloseReason at each transition site and gives finer per-reason attribution for sessionz's close-reasons breakdown, but the whitelist has to be kept in lockstep with every new closeReasonLabel case. Follow up when we add another reason or a downstream consumer needs the state history directly.
mutianf
approved these changes
Jul 27, 2026
3 tasks
sushanb
added a commit
that referenced
this pull request
Jul 27, 2026
…#20224) ## Summary First of five PRs porting the session-pool infrastructure from `feat/bigtable-sessionz-debug` (which powers the recycle-repro fleet) upstream. Stacks on #20215 (Session lifecycle, now merged). ### What lands **New file: `session_list.go`** (~600 LOC) — the per-pool sessionList data structure that groups sessions by the AFE (Application Front End) their handshake landed on. Consumed by the two-tier picker (K-choice-over-AFEs → dequeue-idle-session) in a follow-up PR. Key types: - `AfeID` (int64) — AFE identifier from the server's PeerInfo header at session-open. 0 is the sentinel "unknown" bucket for handshakes that did not carry a peer-info header. - `AfeSnapshot` — value-typed view of an afeHandle for pickers to score without holding sl.mu (Checkout re-resolves by ID; no *afeHandle escapes). - `AfeSnapshotRow` — debug-UI row emitted by `sessionList.Snapshot()` (consumed by afez/sessionz in a later debug PR). - `afeHandle` (unexported) — per-AFE bucket: FIFO idle queue, refCount (idle + inFlight + closing), two PeakEwma trackers. - `SessionHandle` — pool bookkeeping wrapper around Session. Carries `inExpectedCount` (I5 guard against WaitServerClose retry storm) and `activated / closingRecorded / closeRecorded` dedup flags for the pool's per-session hook chain. - `sessionList` (unexported) — the state machine, guarded by `sl.mu`. The state model documents **six invariants (I1-I6)** that every method preserves: ``` I1 inExpectedCount ⇒ handleToAfe[sh] != nil I2 readyCount == count of inExpectedCount handles I3 afesWithReady == {afe : len(afe.sessions) > 0} I4 afe.refCount == count of handleToAfe entries pointing at afe I5 sh in afe.sessions ⇒ handleToAfe[sh]==afe AND inExpectedCount I6 refCount-- only on OnSessionClosed (Closing keeps slot warm) ``` Lock order: `sl.mu` ONLY. `RecordVRpcOutcome` deliberately drops `sl.mu` between the map lookup and the `PeakEwma.Update` so the hot vRPC-outcome path doesn't serialize on it. Consolidates AFE types (`AfeID` + `AfeSnapshot`) that previously lived in `afe_snapshot.go` — sessionList now owns all AFE-bucket concepts. Deletes `afe_snapshot.go`. **New file: `session_list_test.go`** (~770 LOC) — I1-I6 coverage plus per-method tests (OnSessionStarted / Checkout / ReleaseToPool / OnSessionClosing / OnSessionClosed / RecordVRpcOutcome / ReadyAfes / Snapshot / AllHandles / Prune) and a concurrency stress test covering the documented lock-drop path in RecordVRpcOutcome. **Edits to `debug_tracer.go`** — three new tag constants for sessionList bookkeeping violations (all unreachable-under-invariants, kept as belt-and-suspenders): - `tagSessionListStartedNilSession` - `tagSessionListRefcountUnderflow` - `tagSessionListReadyCountUnderflow` **Edit to `session.go`** — one-line comment retarget on `AfeID()` (type now lives in `session_list.go`, not the deleted `afe_snapshot.go`). ### Stack 1. #20211 — Session debug surface (merged) 2. #20213 — Session vRPC dispatch (merged) 3. #20215 — Session lifecycle (merged) 4. **This PR** — sessionList (PR-1 of 5) 5. Next — SessionPoolImpl core, pool_lifecycle, pool_scaling, pool_debug ## Test plan - [x] `go build ./internal/transport/` passes - [x] `go vet ./internal/transport/` clean - [x] `go test ./internal/transport/ -race -count=1 -short -timeout=120s` passes — 20+ new sessionList tests plus all pre-existing.
8 tasks
sushanb
added a commit
that referenced
this pull request
Jul 28, 2026
…#20225) ## Summary Second of five PRs porting the session pool infrastructure from `feat/bigtable-sessionz-debug` (which powers the recycle-repro fleet). Adds `SessionPoolImpl`: the concrete two-tier read/write session pool for one resource. ~4000 LOC across five files plus matching tests. ## Stack - [ ] PR-1: sessionList (#20224) — per-AFE bucketing data structure. **Not yet merged.** - [x] **PR-2 (this)** — SessionPoolImpl (pool + scaling + debug + snapshot). - [ ] PR-3 — `SessionPool` / `Invoker` interfaces + `sessionClient` / `sessionTable` factory wiring. - [ ] PR-4 — Debug pages (`sessionz` / `afez` / `flightz` / `loadz` under `bigtable/debugview/`). - [ ] PR-5 — `bigtable.Client` integration + release notes. **Because PR-1 has not landed, this PR is opened against `main` and the diff includes PR-1's commits.** Once #20224 merges the base can be re-targeted (or this branch rebased) so only PR-2's own delta shows. ## What lands **SessionPoolImpl** (5 files, ~2400 LOC prod + ~2200 LOC tests): - `session_pool.go` — struct + constructor + `Invoke` + `CheckoutSession` (waiter queue, deadline propagation) + pluggable picker via the AFE picker from #20204. - `session_pool_lifecycle.go` — `SessionHooks` wiring, consecutive-failure breaker, `Close` (5-phase teardown), `WaitGoroutines` / `spawns.Wait` choreography so no session-owned goroutine outlives the pool. - `session_pool_scaling.go` — `Tick` loop, `createSession` (dial + `OpenSession` + hook registration), `pendingStarts` / `startingSessions` accounting so scale-up decisions never double-count in-flight opens. Uses the channel-pool pick hint (`ChannelPickHintInto`, added to `connpool.go`) to attribute each session to its underlying channel. - `session_pool_debug.go` — `PoolSnapshot` / slow-vRPC ring / per-close-reason counters / scaling-history buffer / `pickHistory` ring — the input to the sessionz / afez / loadz debug pages (landing in a later PR). - `session_snapshot.go` — the value-typed snapshot record the debug surface consumes; no live locks escape. **Session helpers added** (pool-facing additions to files already touched by prior PRs, kept minimal): - `Session.loops sync.WaitGroup` + `WaitGoroutines()` — pool teardown blocks on this so `readLoop` / `heartbeatLoop` and their `notifyClosed → recordClose` callback chains fully unwind before `Close` returns. Prevents session goroutines from racing metric-var writes across test boundaries. - `Session.closeErr atomic.Pointer[error]` + `setCloseErr` / `closeError` — preserves the raw `Recv` error handed to `handleClose`. Pool surfaces this on consecutive-failure breaker trips so operators see the underlying server rejection (e.g. `FailedPrecondition` when the resource is still being created) instead of only the sentinel. **Supporting additions to existing files:** - `afe_picker.go` — const `defaultAfeRandomSubsetSize = 2` (power-of-two-choices K-choice default; matches Java). - `debug_tracer.go` — three new tag constants: `tagSessionPoolCreatePanic`, `tagSessionPoolConsecutiveFailuresTripped`, `tagSessionPoolCheckoutFailedCINil`. - `connpool.go` — `ChannelPickHintInto(ctx, *atomic.Int32)` context helper. No-op when the channel pool doesn't consume the hint. ## What does NOT land yet - `SessionPool` / `Invoker` interfaces (follow-up PR alongside sessionClient / sessionTable). - `bigtable.Client` integration (PR-3+). - Debug pages under `bigtable/debugview/` (later PR). ## Test plan - [x] `go build ./...` passes. - [x] `go vet ./internal/transport/` clean. - [x] `go test ./internal/transport/ -race -count=1 -short -skip 'AfeLbSim' -timeout=180s` — passes (32s wall). ~2200 LOC of new tests across pool lifecycle, scaling, consecutive-failure breaker, AFE integration, debug surface, snapshot rendering, plus a K-choice bench. --- # Reviewer guide ## Guide 1 — mutianf (human) ### What this PR does Adds `SessionPoolImpl`, the layer that sits above the per-AFE `sessionList` shipped in #20224 and consumes it via a two-tier picker (AFE first, then a ready session in that AFE). It owns the session lifecycle (open / active / closing / close hooks), server-driven scaling via `PoolSizer`, a consecutive-failure circuit breaker, and the debug/observability surface (histograms + ring buffers) that feeds sessionz/loadz. New files: 5 source, 6 test, ~4.9k LOC. Nothing outside `session_pool*.go` / `session_snapshot*.go` is new logic — the small edits elsewhere are hook-plumbing scaffolding already vetted by the session/AFE subagent reviewers. ### Recommended read order 1. **`session_pool.go`** — start here. Struct field layout with per-field ownership comments (`:104-179`), the `waiter` FIFO shape (`:94-101`), `CheckoutSession` two-tier pick + parking (`:235-310`), `Invoke` (`:465-559`), `Stats` (`:361-397`), `UpdateConfig` (`:402-431`), `pickerFromLoadBalancing` (`:439-461`). Skim `session_pool_test.go` (28 tests) — the FIFO waiter, Stats, and UpdateConfig behaviors are all covered there. 2. **`session_pool_lifecycle.go`** — hooks (`onActive:255`, `onClosing:308`, `onClose:336`), `recordSessionClose` once-CAS on `Session.poolCloseRecorded` (`:117-130`), `Close`'s 6-phase teardown (`:154-247`), `noteAbnormalCloseIfAny` breaker (`:363-392`), the three ticker loops (`:426-538`). Skim `session_pool_lifecycle_test.go` — every hook + `Close`. 3. **`session_pool_scaling.go`** — `Tick` (`:81-162`), `createSession` worker (`:164-274`), `scalingReason` (`:278-299`), `noDeadlineButCancellableContext` (`:301-311`). Skim `session_pool_scaling_test.go` — the `scalingInProgress` gate and panic-safety are the only non-obvious contracts. 4. **`session_pool_debug.go`** — `poolMetrics` (`:36-72`), `latencyHist` log2 histogram (`:160-228`), the four ring buffers (slow-vRPC, time-series, lifetimes, pick-history), `recordPickDecision` (`:366-387`). Skim `session_pool_debug_test.go` — mostly ring-cap and rate-computation coverage. 5. **`session_snapshot.go`** — mostly type defs. Focus on `PoolSnapshot` (`:452-594`) and `LoadBalancingSnapshot` (`:414-436`) as the debug-view contract. 6. **`session_pool_consecutive_failures_test.go`** and **`session_pool_afe_test.go`** — end-to-end behavior verification; useful for confirming intent. ### Flow of events - **CheckoutSession → Invoke → release.** `CheckoutSession` (`session_pool.go:235`) opportunistically kicks Tick if `sl.ReadyCount()==0`, snapshots the picker under `p.mu`, then two-tier picks outside the lock: `ReadyAfes()` → `PickAfe` → `Checkout(afeID)` (`:259-268`). Miss → park in the FIFO waiter queue (`:286-289`), bracket `waitersCount` for the sizer (`:291,300`). `Invoke` (`:465`) checks out, runs `sh.session.Invoke`, records latencies (`:508-523`), logs a slow-vRPC row if over threshold (`:524-557`); the deferred `sh.DecOutstanding()` + `noteVRpcOutcome` (`:493-496`) hands the OK-gated latency to the per-AFE PeakEwma tracker. Session release itself is driven by `OnSlotDrained` (installed at `session_pool_scaling.go:228-231`), which returns the handle to `sessionList` and calls `signalFree` — separate from the `defer` in `Invoke`. - **Background Tick.** `startTickLoop` (`session_pool_lifecycle.go:426`) fires every 1 s → `tickOnce` debounces via `tickPending` CAS (`:447-458`) → `Tick` (`session_pool_scaling.go:81`) samples uptimes, gates on `scalingInProgress`, calls `sizer.Decide()`, and on a positive delta reserves `pendingStarts += delta` + `spawns.Add(delta)` under `p.mu` (`:131-138`) then fans out one goroutine per session. Each `createSession` acquires the budget outside `p.mu`, dials via `streamFactory`, transfers `pendingStarts → startingSessions` in one lock (`:246-249`), starts the session, and blocks on `WaitGoroutines` so it stays on `p.spawns` until the session dies. - **Abnormal close → breaker trip.** `onClose` (`session_pool_lifecycle.go:336`) CAS's `closeRecorded`, calls `noteAbnormalCloseIfAny` (`:363`), which bumps `consecutiveFailures` and stores the raw error into `lastAbnormalCloseErr`. Crossing the threshold snapshots the poison, CAS-resets the counter, and calls `drainWaitersWithErr` — waiters get `*consecutiveFailureError` wrapping the last cause (so `errors.Is(err, ErrConsecutiveFailures)` and `status.Code(err)` both still work, `:60-82`). Counter only resets in `onActive` (`:292-293`) — a successful open, not a healthy vRPC. ### Key invariants 1. **Two-tier pick, no re-entrant `p.mu`.** `CheckoutSession` reads `p.picker` under `p.mu` (`session_pool.go:249-255`) then unlocks before calling picker/sessionList. `recordPickDecision` takes `pickerName` as a **parameter** (`session_pool_debug.go:366`, `session_pool.go:260-262`) precisely because the caller already holds no lock — but any new pool method that reads `p.picker.Name()` from a hot path must not re-take `p.mu`. 2. **Waiter FIFO with `waitersCount` bracketed.** Every `PushBack` bumps `waitersCount` (`session_pool.go:291`); every wake path (`ctx.Done`, `w.ready`) decrements it (`:294,300`). `removeWaiter` (`:316`) is idempotent via `w.elem != nil`; `signalFree` and `drainWaitersWithErr` nil out `elem` under `waitersMu` (`:329-358`). `Stats().PendingCount` reads `waitersCount.Load()` — this is the sizer's queue-depth input. 3. **Close-exactly-once accounting.** `sessionsClosed` and `closesByReason` bumps are gated by `Session.poolCloseRecorded.CompareAndSwap(false, true)` inside `recordSessionClose` (`session_pool_lifecycle.go:117-130`). `sh.closingRecorded` and `sh.closeRecorded` are per-handle CAS's protecting the lifetime histogram + the `OnClose` branch. `Close`'s Phase 1 pre-flips both CAS's on every handle (`:187-193`) so a concurrent mid-flight onClosing can't double-count. 4. **Breaker resets only on `onActive`.** `consecutiveFailures.Store(0)` and `lastAbnormalCloseErr.Store(nil)` live at `session_pool_lifecycle.go:292-293`. Not on per-vRPC OK — otherwise one long-lived healthy session would mask a run of failed opens. 5. **Hot path is atomics/RLocks; debug views take snapshots.** `Stats` is the only per-request path that briefly takes `p.mu` (`session_pool.go:362`); everything else on the vRPC path is atomic. Debug snapshotters copy under lock and format after release (`session_snapshot.go:452-594`). ### What NOT to worry about - **Session / vRPC layer itself** — shipped in #20213 / #20215 (state machine, one-in-flight, PeerInfo timing, retry oracle, heartbeat). - **Per-AFE `sessionList` I1-I6** — shipped in #20224, has its own tests. - **`PoolSizer` scaling formula** — already upstream (`pool_sizer.go`); this PR only wires it and consumes `ScaleDecision`. - **AFE pickers (`SimpleAfePicker` / `LeastInFlight` / `LeastLatency`)** — already upstream (`afe_picker.go`); this PR only builds them via `pickerFromLoadBalancing`. - **`SessionThrottler` / `AdaptiveSessionThrottler`** — already upstream; this PR consumes `Acquire` / `Release` / `UpdateConfig`. - **`ClientConfigurationManager` polling** — this pool receives `UpdateConfig` calls; the polling itself is elsewhere. ### Danger zones - **Re-entrant `p.mu` on picker access.** `recordPickDecision` intentionally takes `pickerName` as a param (`session_pool_debug.go:366`). Adding a new pool method that reads `p.picker.Name()` from within a `CheckoutSession` code path is a re-entrant deadlock; pass the name in or snapshot up-front. - **`startingSessions` / `pendingStarts` accounting.** Tick reserves `pendingStarts` under `p.mu` (`session_pool_scaling.go:131-138`), `createSession`'s `reserved` defer releases it on any early return (`:172-179`), and the transfer at `:246-249` is atomic under `p.mu`. `onActive` deletes from `startingSessions` (`session_pool_lifecycle.go:265`). Any new failure branch in `createSession` must preserve the invariant `pendingStarts + len(startingSessions) + Ready = "in-flight scale-up capacity"`. - **`budget.Acquire` blocks; must run OUTSIDE `p.mu`.** Currently at `session_pool_scaling.go:181`, deliberately after the `defer reserved` block and before any `p.mu.Lock()`. Moving it under the lock deadlocks scale-up under budget exhaustion. - **Slow-vRPC is fire-and-forget.** `recordSlowVRpc` (`session_pool_debug.go:301`) appends into a mutex-guarded ring; it's on the vRPC return path but bounded by threshold gating and one small mutex. Do not add I/O, logging fanout, or channel sends here. - **`newTestPool` / bootstrap defaults.** Test helpers construct the pool without a `ClientConfigurationManager`, so the bootstrap defaults from `defaultPoolConfig()` (`session_pool.go:218-226`) are the only config those tests see. Real callers always get `UpdateConfig` synchronously on registration — verify by reading `NewSessionPoolImpl` end-to-end, not by trusting the bootstrap values. - **`Close` phase ordering.** Phase 4 (`poolCancel`) runs AFTER Phase 3 (`wg.Wait` on graceful closes) because Phase 2's `closeCtx` derives from `p.poolCtx`; swapping order strands in-flight graceful closes on a cancelled ctx. Phase 5 (`p.spawns.Wait`) blocks on every createSession goroutine's `WaitGoroutines` — that's why createSession blocks on `s.WaitGoroutines()` at `session_pool_scaling.go:272`. --- ## Guide 2 — mutianf-bot (automated reviewer) ### REAL HAZARDS to flag - **Re-entrant `p.mu` in pool methods called from `CheckoutSession`.** Anchor: `session_pool.go:235-310`. `p.mu` is dropped at `:255` before `PickAfe` / `Checkout` / `recordPickDecision` fire. Flag any newly-added helper called from that block that re-acquires `p.mu`, or any new method that reads `p.picker.Name()` without taking the name as a parameter (see the intentional parameter pattern at `session_pool_debug.go:366`). - **`budget.Acquire` under `p.mu`.** Currently correctly outside the lock at `session_pool_scaling.go:181`. `SessionThrottler.Acquire` blocks on the budget semaphore; calling it while holding `p.mu` would deadlock scale-up. Flag any code path that acquires `p.mu` before line `:181` or moves `Acquire` inside a `Lock`/`Unlock` bracket. - **`sync.Map` allocations on hit paths.** `bumpCloseReason` uses `Load` first, `LoadOrStore(k, new(atomic.Int64))` only on miss (`session_pool_lifecycle.go:102-111`) — this is the correct pattern. Flag any new `sync.Map.LoadOrStore(key, new(...))` call on a hot path that isn't gated by a preceding `Load` — that allocates on every hit. - **Waiter counter drift.** `waitersCount.Add(+1)` at `session_pool.go:291`, `Add(-1)` on both the `ctx.Done` branch (`:294`) and the `w.ready` branch (`:300`). Flag any new wake path, timeout branch, or early-return between `:291` and `:308` that doesn't decrement, and any new enqueue site that doesn't increment. Drift here corrupts the sizer's `PendingCount` input. - **Unbalanced `pendingStarts` / `startingSessions`.** Tick increments `pendingStarts` under `p.mu` at `session_pool_scaling.go:131-138`; `createSession`'s `reserved` defer at `:172-179` releases on early return; the transfer to `startingSessions` at `:246-249` is atomic; `onActive` deletes at `session_pool_lifecycle.go:265`; failed-start deletes at `session_pool_scaling.go:253-255`. Flag any new failure branch in `createSession` that returns without either the `reserved` defer or an explicit transfer/cleanup. - **Missing CAS on close-once flags.** `sessionsClosed.Add(1)` sits downstream of `s.poolCloseRecorded.CompareAndSwap(false, true)` at `session_pool_lifecycle.go:121`. Lifetime histogram sits downstream of `sh.closingRecorded.CompareAndSwap(false, true)` at `:187` and `:317`. `recordSessionClose` re-entry is dedup'd at `:190` and `:347`. Flag any new counter bump / histogram record that isn't gated by the corresponding CAS. - **Uncontrolled goroutine spawns.** Every pool-spawned goroutine must be tracked so `Close` phases 3/5/6 can drain: createSession workers use `p.spawns.Add`/`Done` under `p.mu` (`session_pool_scaling.go:137-138,149`); `spawnTickOnce` uses the same pattern (`session_pool_lifecycle.go:464-476`); the mid-Close onActive dispatch adds at `:274`. Session-owned goroutines (`readLoop`/`heartbeatLoop`) are tracked on `Session.loops` and drained via `WaitGoroutines`. Flag any bare `go func(){…}()` in these files with no WaitGroup entry. - **Hot-path lock churn.** `p.mu` is briefly held on the fast path only at `session_pool.go:239-241` and `:249-255`. `Stats` takes `p.mu` (`:362`). Flag any new `p.mu.Lock()` inserted into `Invoke` (`:465-559`) or the pre-park stretch of `CheckoutSession` outside the two existing points. ### FALSE-POSITIVE SUPPRESSION - **"Package name should be `transport`, not `internal`."** The directory is `bigtable/internal/transport/` but the `package` declaration is `package internal` (confirmed at `session_pool.go:20`). Package name is correct as-is; do NOT flag. - **Local variable shadowing an exported type** where the local name is idiomatic (e.g. `afeID` local vs `AfeID` type in `CheckoutSession`). Already resolved intentionally; do NOT re-raise variants. - **`goimports` / `gofmt` / column-alignment / trailing-newline nits.** CI (`goimports -l`, `gofmt -l`, `go vet`) already gates these. Bot echo is noise. - **Comments referencing PR #20213 / #20215 / #20224.** Stacked-PR context, not stale references; do NOT suggest removal. - **Test coverage complaints for `pool_sizer.go`, `afe_picker.go`, `session_list.go`, `session.go`, `session_vrpc.go`, `session_throttler.go`, `client_configuration_manager.go`, `default_client_config.go`.** All shipped in earlier PRs (#20213, #20215, #20224) with their own tests; out of scope here. - **"Missing error wrapping"** on internal-only calls where the caller already annotates via `fmt.Errorf("POOL %s ...: %w", ...)` or via `btopt.Debugf`. Do NOT suggest adding a second wrap. - **Retry loop / context propagation questions on `Session.Invoke`.** That's the Session layer (`session_vrpc.go`), out of scope for this PR. - **"Consider using `sync.RWMutex` instead of `sync.Mutex` on `p.mu`."** The pool holds `p.mu` for tens of nanoseconds at a time and never for read-heavy loops; the added atomic on `RLock`/`RUnlock` would cost more than it saves. Do NOT suggest. - **"Consider extracting anonymous goroutine into named function."** Style-only; do NOT suggest for the three ticker loops or the createSession worker. ### SCOPE BOUNDARY Comment ONLY on: - `bigtable/internal/transport/session_pool.go` - `bigtable/internal/transport/session_pool_lifecycle.go` - `bigtable/internal/transport/session_pool_scaling.go` - `bigtable/internal/transport/session_pool_debug.go` - `bigtable/internal/transport/session_snapshot.go` - `bigtable/internal/transport/session_pool_*_test.go` - `bigtable/internal/transport/session_snapshot_test.go` Do NOT comment on additions to: - `session.go` / `session_vrpc.go` (WaitGoroutines / closeError additions — vetted) - `connpool.go` (`ChannelPickHintInto` helper — vetted) - `afe_picker.go` (`defaultAfeRandomSubsetSize` constant — vetted) - `debug_tracer.go` (3 new tags — vetted) These are supporting scaffolding, already reviewed by the 3 subagent reviewers in this stack. Only re-raise if something looks actively unsafe. ### EFFORT SCALING - ~4.9k LOC across 12 files. Do NOT paginate uniformly. - **First pass — the 4 hot source files, in this order:** 1. `session_pool.go` (559 LOC) 2. `session_pool_lifecycle.go` (538 LOC) 3. `session_pool_scaling.go` (311 LOC) 4. `session_pool_debug.go` (416 LOC) - **Second pass ONLY if a first-pass finding needs corroboration:** `session_snapshot.go` (594 LOC, mostly type defs), and the tests. Tests use `newTestPool`, which skips config wiring — do NOT flag bootstrap defaults on tests as if they were production paths. - If a first-pass finding is a real hazard from the list above, cite the file:line and the exact anchor pattern it violates. Do not file speculative "consider" comments.
sushanb
pushed a commit
that referenced
this pull request
Aug 3, 2026
🤖 I have created a release *beep* *boop* --- ## [1.52.0](bigtable/v1.51.0...bigtable/v1.52.0) (2026-08-03) ### Features * **bigtable:** Add AFE picker (Simple / LeastInFlight / LeastLatency) ([#20204](#20204)) ([bcbf714](bcbf714)) * **bigtable:** Add ClientConfig.DisableSession to opt out of session backend ([#20297](#20297)) ([7ee5e44](7ee5e44)) * **bigtable:** Add getClientConfigDirectAccessChecker for session pools ([#20209](#20209)) ([3b8d30a](3b8d30a)) * **bigtable:** Add NoOpChannelPrimer for session channel pools ([#20208](#20208)) ([d055a8a](d055a8a)) * **bigtable:** Add per-AFE sessionList for the two-tier session pool ([#20224](#20224)) ([dbf0c3f](dbf0c3f)) * **bigtable:** Add protoRowToRow conversion helper for TableShim ([#20257](#20257)) ([1297143](1297143)) * **bigtable:** Add Session debug surface (observability fields + methods) ([#20211](#20211)) ([d8d3e16](d8d3e16)) * **bigtable:** Add Session lifecycle (Start, Close, ForceClose, readLoop, heartBeatLoop) ([#20215](#20215)) ([b9e53c6](b9e53c6)) * **bigtable:** Add Session struct + state machine ([#20117](#20117)) ([09acbb3](09acbb3)) * **bigtable:** Add session.Config.EnableDebug to gate sessionz debug state ([#20247](#20247)) ([ce74c31](ce74c31)) * **bigtable:** Add SessionClient + SessionTable + lazyPool ([#20228](#20228)) ([ab2c96c](ab2c96c)) * **bigtable:** Add SessionPoolImpl (two-tier pool + scaling + debug) ([#20225](#20225)) ([683eda8](683eda8)) * **bigtable:** Rename session pool display to <resource-id>-<PERM> ([#20248](#20248)) ([35e146e](35e146e)) * **bigtable:** Route Client.Open()-returned *Table through the Diverter ([#20273](#20273)) ([2b81c7d](2b81c7d)) * **bigtable:** State-based classification for abnormal session close ([#20243](#20243)) ([f2905b7](f2905b7)) * **bigtable:** TableShim fallback to classic on session UNIMPLEMENTED ([#20269](#20269)) ([36540af](36540af)) * **bigtable:** TTL-on-idle cache for per-resource session.TableAPI ([#20263](#20263)) ([00b2a49](00b2a49)) * **bigtable:** Wire Diverter on Client and route Open* via TableShim ([#20256](#20256)) ([b32fbd7](b32fbd7)) ### Bug Fixes * **bigtable:** AFE picker latency signal — subtract poolWait and compute TransportLatency = wire − backend at source ([#20281](#20281)) ([bb8c4d5](bb8c4d5)) * **bigtable:** Guard NewStream OnFinish against grpc-go double-fire ([#20295](#20295)) ([b51da29](b51da29)) * **bigtable:** Real per-resource pool teardown on sessionTable.Close + cache close-race gate ([#20264](#20264)) ([599aea9](599aea9)) * **bigtable:** Session.durations / session.uptime — set explicit histogram bucket boundaries ([#20276](#20276)) ([97eee22](97eee22)) * **bigtable:** SessionTableHandle self-heals across cache eviction ([#20296](#20296)) ([0dd98cd](0dd98cd)) * **bigtable:** Translate ctx errors to gRPC status on session vRPC ([#20299](#20299)) ([0f3b2a5](0f3b2a5)) * **bigtable:** Treat PingAndWarm NotFound as a successful prime ([#20219](#20219)) ([a1557ad](a1557ad)) ### Performance Improvements * **bigtable:** Delete periodic Tick loop; sizing is event-driven ([#20285](#20285)) ([2c096bd](2c096bd)) * **bigtable:** Drop pick_lost_race debug tag from CheckoutSession hot path ([#20280](#20280)) ([bd0e400](bd0e400)) --- This PR was generated with [Release Please](https://github.com/googleapis/release-please). See [documentation](https://github.com/googleapis/release-please#release-please). Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com>
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
3 participants
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
Third and final PR in the Session core stack. Adds the lifecycle orchestration that drives a Session from `Start` through teardown, plus the `readLoop` that dispatches server frames to the vRPC handlers landed in #20213.
Stacks on #20213 (Session vRPC dispatch + slot lifecycle). Deletes the minimal `ForceClose` stub that PR shipped, replacing it with the full lifecycle-shaped body.
What lands
New file: `session_lifecycle.go` (~640 LOC)
New file: `session_lifecycle_test.go` (~660 LOC) — 30+ tests covering Start, ForceClose, Close, readLoop, handleSessionResponse dispatch, handleOpenSession + peerInfo parsing, handleGoAway, handleClose, handleErrorResponse (rpc_id=0 harmlessly drops via routeVRPCFrame guards), heartBeatLoop (reactive wake, idle-gate, missed-heartbeat ForceClose).
Edits to `session_vrpc.go`
Edits to `session_test.go`
Stack
Test plan