Skip to content

Commit f2905b7

Browse files
sushanbsushanb
authored
feat(bigtable): state-based classification for abnormal session close (#20243)
## Summary Bring the Go client's abnormal session-close classification in line with Java by keying on the pre-close `State`, and stamp a `StreamEnd:<code>` close-reason so sessionz surfaces exit codes that would otherwise land in the \"Unspecified\" bucket. - **`refactor(bigtable): state-based abnormal-close classification (Java parity)`** — `noteAbnormalCloseIfAny` now consults `Session.prevStateAtClose`. Any session whose prior state was NOT `StateWaitServerClose` counts as abnormal; sessions that gracefully entered `WAIT_SERVER_CLOSE` (client-initiated `CloseSession` → server ACK) do not. Matches Java's `SessionImpl.handleAbnormalClose` gating on `prevState != WAIT_SERVER_CLOSE`. Heartbeat-miss sessions (ForceClose from `Ready` → `prev=Ready`) correctly still count as abnormal. - **`observe(bigtable): stamp StreamEnd:<code> close-reason in handleClose`** — the terminal `handleClose` path CAS-stamps a `StreamEnd:<grpc-code>` reason so a bare stream error no longer renders as \"Unspecified\" on sessionz. Upstream stampers (`GoAway`, `MissedHeartbeat`, `Error`, `User`) still win the CAS because `setCloseReason` is one-shot — they run first and this stamp is a strict fallback. - **`fix(bigtable): special-case io.EOF in StreamEnd close-reason`** — Igor caught during review: `status.Code(io.EOF) == codes.Unknown`, so the default arm rendered a graceful server-side shutdown as `\"StreamEnd:Unknown\"` — hiding a distinct-and-common signal. Special-cases `errors.Is(err, io.EOF)` → `\"StreamEnd:EOF\"` before falling through. ## Test plan - [ ] `go test ./bigtable/internal/transport/ -count=1 -race -short -timeout=120s` - [ ] Manual sessionz check: run a workload where the server GOAWAYs a session; confirm `close_reason` renders as `GoAway` (not `StreamEnd:EOF`), i.e. the upstream GoAway stamper still wins the CAS. - [ ] Manual sessionz check: run a workload where an in-flight session sees a plain graceful stream close; confirm `close_reason = StreamEnd:EOF` (not `Unspecified`, not `StreamEnd:Unknown`). - [ ] Heartbeat-miss path: kill a session's inbound frames; confirm `MissedHeartbeat` still wins the CAS (this PR doesn't change that path, but verify no regression). ## Reviewer notes Three reviewers ran clean: `session-reviewer` (behavioral — 4 specs), `session-component-review` (boundaries — Part B + Part C), `igor-reviewer` (persona). No spec drift; `setCloseReason` remains CAS-once and no upstream stampers are demoted.
1 parent 3b8d30a commit f2905b7

4 files changed

Lines changed: 86 additions & 112 deletions

File tree

‎bigtable/internal/transport/session.go‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -156,8 +156,20 @@ type Session struct {
156156

157157
state atomic.Int32
158158

159-
// closingOnce/closeOnce fire hooks.OnClosing/OnClose exactly once each
160-
// even when multiple teardown paths race.
159+
// prevStateAtClose is the state the session was in immediately
160+
// before its final transition to StateClosed — captured as the
161+
// prev return of transitionTo(StateClosed, ...) at the two
162+
// transition sites (ForceClose, handleClose). Set-once by
163+
// construction (transitionTo(StateClosed) applies at most once),
164+
// then read from hooks.OnClose consumers. Lets the pool
165+
// distinguish a client-initiated clean-close (prev == WSC) from a
166+
// server-initiated / transport-error close without carrying a
167+
// side-channel bool.
168+
prevStateAtClose atomic.Int32
169+
170+
// closingOnce serializes hooks.OnClosing so it fires exactly once
171+
// across the four transition sites that can drive a session out of
172+
// Ready (Close, ForceClose, handleGoAway, handleClose).
161173
closingOnce sync.Once
162174
closeOnce sync.Once
163175

‎bigtable/internal/transport/session_lifecycle.go‎

Lines changed: 34 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ func (s *Session) ForceClose(req *spb.CloseSessionRequest) {
100100
if !ok {
101101
return
102102
}
103+
s.prevStateAtClose.Store(int32(prev))
103104
if prev == StateNew {
104105
// Force-closing a NEW session means the pool decided to tear us
105106
// down before Start ran — a bookkeeping oddity worth flagging.
@@ -226,7 +227,9 @@ func (s *Session) Close(ctx context.Context, req *spb.CloseSessionRequest) error
226227
return fmt.Errorf("send close session request: %w", err)
227228
}
228229
// Advance to WaitServerClose so the pool monitor can see we're waiting
229-
// on the server. handleClose accepts StateWaitServerClose → Closed.
230+
// on the server. handleClose accepts StateWaitServerClose → Closed
231+
// and will capture WSC as prevStateAtClose — that's the signal
232+
// noteAbnormalCloseIfAny reads to skip the abnormal-close counter.
230233
s.transitionTo(StateWaitServerClose, isState(StateClosing))
231234
return nil
232235
}
@@ -279,11 +282,10 @@ func (s *Session) readLoop(ctx context.Context) {
279282
// Receiving any recognized frame resets the heartbeat watchdog; unknown
280283
// frames do NOT, so a misbehaving server cannot keep the watchdog
281284
// satisfied with junk payloads and a rogue future oneof variant can't
282-
// mask a broken stream. Java parity: `SessionImpl.handleUnknownResponseMessage`
283-
// also does not reset the heartbeat. The watchdog is only armed while
284-
// a vRPC is in-flight anyway — during that window the server MUST be
285-
// sending heartbeats, so a new-variant frame arriving instead of a
286-
// heartbeat within the interval is itself a signal worth surfacing.
285+
// mask a broken stream. The watchdog is only armed while a vRPC is
286+
// in-flight anyway — during that window the server MUST be sending
287+
// heartbeats, so a new-variant frame arriving instead of a heartbeat
288+
// within the interval is itself a signal worth surfacing.
287289
func (s *Session) handleSessionResponse(resp *spb.SessionResponse) {
288290
switch p := resp.GetPayload().(type) {
289291
case *spb.SessionResponse_OpenSession:
@@ -446,111 +448,53 @@ func (s *Session) handleGoAway(goAway *spb.GoAwayResponse) {
446448
// when the server's EOF arrives after a CloseSession we sent) and cancels
447449
// every remaining in-flight RPC.
448450
//
449-
// The close reason is derived from the Recv error if no more-specific
450-
// reason was recorded earlier — see streamEndReason. setCloseReason is
451-
// CompareAndSwap-once, so a GoAway / MissedHeartbeat / Error stamp from
452-
// upstream always wins; the categorized StreamEnd label only sticks when
453-
// the stream ended without any other path classifying it first.
451+
// The close reason is stamped by upstream paths (Close / ForceClose /
452+
// handleGoAway / heartbeat trip / handleErrorResponse) before handleClose
453+
// runs, and setCloseReason is CompareAndSwap-once — so no reason
454+
// classification is needed here.
454455
func (s *Session) handleClose(err error) {
455-
if _, ok := s.transitionTo(StateClosed, notState(StateClosed)); !ok {
456+
prev, ok := s.transitionTo(StateClosed, notState(StateClosed))
457+
if !ok {
456458
return
457459
}
460+
s.prevStateAtClose.Store(int32(prev))
458461
// Ready → Closed can happen directly here (server EOFed without a
459462
// prior GoAway or CloseSession). Guarantee onClosing fires before the
460463
// notifyClosed below drives onClose. closingOnce makes this a no-op
461464
// when handleGoAway or Close already fired earlier.
462465
s.notifyClosing()
463-
reason := streamEndReason(err)
464-
s.setCloseReason(reason)
465-
s.setCloseErr(err)
466-
// After setCloseReason (CompareAndSwap-once), the *final* reason may
467-
// be an earlier stamp (GoAway / MissedHeartbeat / Error) or the
468-
// streamEndReason we just computed. Only flag as abnormal when the
469-
// final reason is a StreamEnd category that isn't a clean shutdown.
470-
if isAbnormalCloseReason(s.CloseReason()) {
471-
recordDebugTag(tagSessionAbnormalClose)
466+
// Fallback close-reason stamp for observability: paths that reach
467+
// handleClose with no prior stamp (transport-level EOF/Unavailable
468+
// without a GoAway or client Close) would otherwise fall into
469+
// sessionz's "Unspecified" bucket. setCloseReason is CAS-once, so
470+
// upstream stampers (GoAway, MissedHeartbeat, Error, User) still win.
471+
if err != nil {
472+
// Special-case io.EOF (graceful server-side shutdown): status.Code(io.EOF)
473+
// returns codes.Unknown, which would render as "StreamEnd:Unknown" and
474+
// hide a distinct-and-common signal. Ctx errors (Canceled,
475+
// DeadlineExceeded) are already mapped by grpc-go's status helpers so
476+
// the default branch covers them correctly.
477+
if errors.Is(err, io.EOF) {
478+
s.setCloseReason("StreamEnd:EOF")
479+
} else {
480+
s.setCloseReason("StreamEnd:" + status.Code(err).String())
481+
}
472482
}
483+
s.setCloseErr(err)
473484
inFlight := 0
474485
if s.activeVRPC() != nil {
475486
inFlight = 1
476487
}
477488
age := time.Since(s.StartedAt())
478489
lastRPC := s.nextRPCID.Load()
479490
peer := s.peerInfoSummary()
480-
s.recordEvent("close", "reason=%s age=%v in_flight=%d last_rpc_id=%d %s raw_err=%v",
481-
reason, age, inFlight, lastRPC, peer, err)
491+
s.recordEvent(SessionEventClose, "age=%v in_flight=%d last_rpc_id=%d %s raw_err=%v",
492+
age, inFlight, lastRPC, peer, err)
482493
s.cancelActiveRPCs(unavailable(err, "session closed: %v", err))
483494
s.signalQuiescent()
484495
s.notifyClosed(err)
485496
}
486497

487-
// streamEndReason classifies the Recv error that ended the stream. The
488-
// returned label is what shows up in sessionz's Close-reasons breakdown
489-
// when no upstream path stamped a more specific reason (GoAway,
490-
// MissedHeartbeat, Error, etc.).
491-
//
492-
// Categories the operator typically cares about:
493-
//
494-
// StreamEnd:EOF — server closed the stream cleanly with
495-
// io.EOF (graceful shutdown from server's
496-
// side that didn't go through GoAway)
497-
// StreamEnd:Canceled — local ctx cancel (pool teardown,
498-
// client app exit) or grpc CANCELED
499-
// StreamEnd:DeadlineExceeded — ctx deadline or grpc DEADLINE_EXCEEDED
500-
// StreamEnd:Unavailable — transport-level break (TCP drop,
501-
// connection recycler killed the channel,
502-
// load balancer evicted the backend)
503-
// StreamEnd:Internal — server INTERNAL error
504-
// StreamEnd:{Code} — any other gRPC status code (verbatim)
505-
// StreamEnd:Other — no recognizable category (extremely rare)
506-
// StreamEnd — err was nil (shouldn't happen since Recv
507-
// only returns on error)
508-
func streamEndReason(err error) string {
509-
if err == nil {
510-
return "StreamEnd"
511-
}
512-
if errors.Is(err, io.EOF) {
513-
return "StreamEnd:EOF"
514-
}
515-
if errors.Is(err, context.Canceled) {
516-
return "StreamEnd:Canceled"
517-
}
518-
if errors.Is(err, context.DeadlineExceeded) {
519-
return "StreamEnd:DeadlineExceeded"
520-
}
521-
if st, ok := status.FromError(err); ok {
522-
return "StreamEnd:" + st.Code().String()
523-
}
524-
return "StreamEnd:Other"
525-
}
526-
527-
// isAbnormalCloseReason returns true when the recorded close reason
528-
// looks like something we did NOT initiate cleanly. Clean paths:
529-
// EOF (server graceful), Canceled (client teardown / ctx cancel), and
530-
// the explicit client-initiated reasons stamped by handleGoAway /
531-
// heartbeatLoop / handleErrorResponse. Anything else — a StreamEnd
532-
// tagged with a transport-failure code, or the bare "StreamEnd" that
533-
// indicates Recv returned nil (which shouldn't happen) — is abnormal
534-
// and worth flagging.
535-
//
536-
// TODO(sushanb): move to a state-based classifier per mutianf's review
537-
// on #20215. Current reason-string scheme encodes state indirectly (via
538-
// CAS-once CloseReason stamped at each transition site) and gives finer
539-
// per-reason attribution for sessionz's close-reasons breakdown, but a
540-
// state-transition source of truth ("did we go New→Ready→Closing→
541-
// WaitServerClose→Closed cleanly?") is more robust — the whitelist
542-
// here has to be kept in lockstep with every new closeReasonLabel case.
543-
// Refactor when we add a new close-reason (or when a downstream
544-
// consumer wants the state-transition history directly).
545-
func isAbnormalCloseReason(reason string) bool {
546-
switch reason {
547-
case "StreamEnd:EOF", "StreamEnd:Canceled",
548-
"GoAway", "MissedHeartbeat", "Error", "":
549-
return false
550-
}
551-
return strings.HasPrefix(reason, "StreamEnd")
552-
}
553-
554498
// heartbeatLoop watches the session's heartbeat deadline using a single Timer
555499
// that re-arms itself when a frame extends the deadline. The watchdog is
556500
// only enforced while at least one VRPC is in flight: the server emits

‎bigtable/internal/transport/session_pool_consecutive_failures_test.go‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,12 @@ func abnormalOnCloseFor(t testing.TB, p *SessionPoolImpl, abnormal bool) {
5151
return
5252
}
5353
sh := injectActiveSession(t, p, "active", time.Now())
54+
// Simulate the state history that real session.Close() → handleClose
55+
// produces: Ready → Closing → WSC → Closed, so prevStateAtClose = WSC.
56+
// The fixture bypasses that path (calls onClose directly), so stamp
57+
// prevStateAtClose here to match what noteAbnormalCloseIfAny expects
58+
// on a client-initiated clean close.
59+
sh.session.prevStateAtClose.Store(int32(StateWaitServerClose))
5460
p.onClose(sh, nil)
5561
}
5662

@@ -94,11 +100,15 @@ func TestConsecutiveFailures_UserReasonNotAbnormal(t *testing.T) {
94100
// the second close instead of needing 10.
95101
p.consecutiveFailureThreshold.Store(2)
96102

97-
// Fire "User" close-reason twice on ACTIVATED sessions — the
98-
// closest fixture to Pool.Close's Phase-2 (activated sessions
99-
// closed with REASON_USER). State-based gate must NOT count them.
103+
// Fire "User" close-reason twice on ACTIVATED sessions that also
104+
// went through WSC — the closest fixture to Pool.Close's Phase-2
105+
// (session.Close() transitions through WSC, then onClose fires).
106+
// State-history gate must NOT count them.
100107
for i := 0; i < 2; i++ {
101108
sh := injectActiveSession(t, p, "user-close", time.Now())
109+
// Matches the state history real session.Close() → handleClose
110+
// produces: prevStateAtClose = WSC when the WSC → Closed step ran.
111+
sh.session.prevStateAtClose.Store(int32(StateWaitServerClose))
102112
stampCloseReason(sh.session, "User")
103113
p.onClose(sh, nil)
104114
}

‎bigtable/internal/transport/session_pool_lifecycle.go‎

Lines changed: 25 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -356,18 +356,19 @@ func (p *SessionPoolImpl) onClose(sh *SessionHandle, err error) {
356356
p.noteAbnormalCloseIfAny(sh)
357357
}
358358

359-
// noteAbnormalCloseIfAny bumps the consecutive-failure counter when the
360-
// session died before it ever became usable. Classification is
361-
// state-based via sh.activated: onActive is the sole writer of that
362-
// flag, so `!sh.activated.Load()` is the exact "never reached
363-
// StateReady" signal — no reason-string whitelist to keep in lockstep
364-
// with new close reasons, no accidental mis-classification when the
365-
// server invents a new REASON. A session that activated (even briefly)
366-
// is treated as a healthy open regardless of how it later died: the
367-
// counter is reset on every onActive so a churn pattern of activate →
368-
// server-side GoAway → replace converges to zero. Crossing the
369-
// threshold drains every parked waiter with ErrConsecutiveFailures and
370-
// resets. CAS on reset guards against two goroutines double-draining.
359+
// noteAbnormalCloseIfAny bumps the consecutive-failure counter when a
360+
// session's terminal transition did NOT come through StateWaitServerClose.
361+
// Classification is state-based via Session.prevStateAtClose, captured
362+
// at the two transitionTo(StateClosed, …) call sites — no reason-string
363+
// whitelist to keep in lockstep with new close reasons; the
364+
// state-transition history is the source of truth. A clean shutdown
365+
// (Close() → WSC → server ack → Closed) skips the counter; a
366+
// server-initiated GoAway / heartbeat trip / stream error on a Ready
367+
// session counts. Also emits `tagSessionAbnormalClose` on the counted
368+
// path so operators can see per-abnormal-close volume in debug-tag
369+
// counters. Crossing the threshold drains every parked waiter with
370+
// ErrConsecutiveFailures and resets. CAS on reset guards against two
371+
// goroutines double-draining.
371372
func (p *SessionPoolImpl) noteAbnormalCloseIfAny(sh *SessionHandle) {
372373
// Defensive nil-guard: production callers always pass a live sh
373374
// with sh.session backfilled (createSession sets it before wiring
@@ -377,13 +378,20 @@ func (p *SessionPoolImpl) noteAbnormalCloseIfAny(sh *SessionHandle) {
377378
if sh == nil || sh.session == nil {
378379
return
379380
}
380-
// State-based gate: activated=true means the session reached
381-
// StateReady at least once (onActive fired). Skip the trip counter —
382-
// server-initiated GoAway / heartbeat missed / stream errors on an
383-
// already-Ready session are transport hiccups, not open failures.
384-
if sh.activated.Load() {
381+
// State-history gate: skip the trip counter only when the session's
382+
// state immediately before Closed was StateWaitServerClose — i.e.,
383+
// the client-initiated clean-close path (Close() sent CloseSession,
384+
// server acked, handleClose completed WSC → Closed). Every other
385+
// terminal transition counts:
386+
// - never activated → open failure (prev = Starting)
387+
// - activated then server GoAway / heartbeat miss / stream error
388+
// → prev = Closing (transport failure worth surfacing)
389+
// - sweep of stuck WSC session via ForceClose → prev = WSC
390+
// already, so those stay exempt.
391+
if State(sh.session.prevStateAtClose.Load()) == StateWaitServerClose {
385392
return
386393
}
394+
recordDebugTag(tagSessionAbnormalClose)
387395
s := sh.session
388396
if e := s.closeError(); e != nil {
389397
p.lastAbnormalCloseErr.Store(&e)

0 commit comments

Comments
 (0)