googleapis · sushanb · Jul 28, 2026 · Jul 28, 2026 · Jul 28, 2026 · Jul 28, 2026
@@ -156,8 +156,20 @@ type Session struct {
 
 	state atomic.Int32
 
-	// closingOnce/closeOnce fire hooks.OnClosing/OnClose exactly once each
-	// even when multiple teardown paths race.
+	// prevStateAtClose is the state the session was in immediately
+	// before its final transition to StateClosed — captured as the
+	// prev return of transitionTo(StateClosed, ...) at the two
+	// transition sites (ForceClose, handleClose). Set-once by
+	// construction (transitionTo(StateClosed) applies at most once),
+	// then read from hooks.OnClose consumers. Lets the pool
+	// distinguish a client-initiated clean-close (prev == WSC) from a
+	// server-initiated / transport-error close without carrying a
+	// side-channel bool.
+	prevStateAtClose atomic.Int32
+
+	// closingOnce serializes hooks.OnClosing so it fires exactly once
+	// across the four transition sites that can drive a session out of
+	// Ready (Close, ForceClose, handleGoAway, handleClose).
 	closingOnce sync.Once
 	closeOnce   sync.Once
 

@@ -100,6 +100,7 @@ func (s *Session) ForceClose(req *spb.CloseSessionRequest) {
 	if !ok {
 		return
 	}
+	s.prevStateAtClose.Store(int32(prev))
 	if prev == StateNew {
 		// Force-closing a NEW session means the pool decided to tear us
 		// down before Start ran — a bookkeeping oddity worth flagging.
@@ -226,7 +227,9 @@ func (s *Session) Close(ctx context.Context, req *spb.CloseSessionRequest) error
 		return fmt.Errorf("send close session request: %w", err)
 	}
 	// Advance to WaitServerClose so the pool monitor can see we're waiting
-	// on the server. handleClose accepts StateWaitServerClose → Closed.
+	// on the server. handleClose accepts StateWaitServerClose → Closed
+	// and will capture WSC as prevStateAtClose — that's the signal
+	// noteAbnormalCloseIfAny reads to skip the abnormal-close counter.
 	s.transitionTo(StateWaitServerClose, isState(StateClosing))
 	return nil
 }
@@ -279,11 +282,10 @@ func (s *Session) readLoop(ctx context.Context) {
 // Receiving any recognized frame resets the heartbeat watchdog; unknown
 // frames do NOT, so a misbehaving server cannot keep the watchdog
 // satisfied with junk payloads and a rogue future oneof variant can't
-// mask a broken stream. Java parity: `SessionImpl.handleUnknownResponseMessage`
-// also does not reset the heartbeat. The watchdog is only armed while
-// a vRPC is in-flight anyway — during that window the server MUST be
-// sending heartbeats, so a new-variant frame arriving instead of a
-// heartbeat within the interval is itself a signal worth surfacing.
+// mask a broken stream. The watchdog is only armed while a vRPC is
+// in-flight anyway — during that window the server MUST be sending
+// heartbeats, so a new-variant frame arriving instead of a heartbeat
+// within the interval is itself a signal worth surfacing.
 func (s *Session) handleSessionResponse(resp *spb.SessionResponse) {
 	switch p := resp.GetPayload().(type) {
 	case *spb.SessionResponse_OpenSession:
@@ -446,111 +448,53 @@ func (s *Session) handleGoAway(goAway *spb.GoAwayResponse) {
 // when the server's EOF arrives after a CloseSession we sent) and cancels
 // every remaining in-flight RPC.
 //
-// The close reason is derived from the Recv error if no more-specific
-// reason was recorded earlier — see streamEndReason. setCloseReason is
-// CompareAndSwap-once, so a GoAway / MissedHeartbeat / Error stamp from
-// upstream always wins; the categorized StreamEnd label only sticks when
-// the stream ended without any other path classifying it first.
+// The close reason is stamped by upstream paths (Close / ForceClose /
+// handleGoAway / heartbeat trip / handleErrorResponse) before handleClose
+// runs, and setCloseReason is CompareAndSwap-once — so no reason
+// classification is needed here.
 func (s *Session) handleClose(err error) {
-	if _, ok := s.transitionTo(StateClosed, notState(StateClosed)); !ok {
+	prev, ok := s.transitionTo(StateClosed, notState(StateClosed))
+	if !ok {
 		return
 	}
+	s.prevStateAtClose.Store(int32(prev))
 	// Ready → Closed can happen directly here (server EOFed without a
 	// prior GoAway or CloseSession). Guarantee onClosing fires before the
 	// notifyClosed below drives onClose. closingOnce makes this a no-op
 	// when handleGoAway or Close already fired earlier.
 	s.notifyClosing()
-	reason := streamEndReason(err)
-	s.setCloseReason(reason)
-	s.setCloseErr(err)
-	// After setCloseReason (CompareAndSwap-once), the *final* reason may
-	// be an earlier stamp (GoAway / MissedHeartbeat / Error) or the
-	// streamEndReason we just computed. Only flag as abnormal when the
-	// final reason is a StreamEnd category that isn't a clean shutdown.
-	if isAbnormalCloseReason(s.CloseReason()) {
-		recordDebugTag(tagSessionAbnormalClose)
+	// Fallback close-reason stamp for observability: paths that reach
+	// handleClose with no prior stamp (transport-level EOF/Unavailable
+	// without a GoAway or client Close) would otherwise fall into
+	// sessionz's "Unspecified" bucket. setCloseReason is CAS-once, so
+	// upstream stampers (GoAway, MissedHeartbeat, Error, User) still win.
+	if err != nil {
+		// Special-case io.EOF (graceful server-side shutdown): status.Code(io.EOF)
+		// returns codes.Unknown, which would render as "StreamEnd:Unknown" and
+		// hide a distinct-and-common signal. Ctx errors (Canceled,
+		// DeadlineExceeded) are already mapped by grpc-go's status helpers so
+		// the default branch covers them correctly.
+		if errors.Is(err, io.EOF) {
+			s.setCloseReason("StreamEnd:EOF")
+		} else {
+			s.setCloseReason("StreamEnd:" + status.Code(err).String())
+		}
 	}
+	s.setCloseErr(err)
 	inFlight := 0
 	if s.activeVRPC() != nil {
 		inFlight = 1
 	}
 	age := time.Since(s.StartedAt())
 	lastRPC := s.nextRPCID.Load()
 	peer := s.peerInfoSummary()
-	s.recordEvent("close", "reason=%s age=%v in_flight=%d last_rpc_id=%d %s raw_err=%v",
-		reason, age, inFlight, lastRPC, peer, err)
+	s.recordEvent(SessionEventClose, "age=%v in_flight=%d last_rpc_id=%d %s raw_err=%v",
+		age, inFlight, lastRPC, peer, err)
 	s.cancelActiveRPCs(unavailable(err, "session closed: %v", err))
 	s.signalQuiescent()
 	s.notifyClosed(err)
 }
 
-// streamEndReason classifies the Recv error that ended the stream. The
-// returned label is what shows up in sessionz's Close-reasons breakdown
-// when no upstream path stamped a more specific reason (GoAway,
-// MissedHeartbeat, Error, etc.).
-//
-// Categories the operator typically cares about:
-//
-//	StreamEnd:EOF              — server closed the stream cleanly with
-//	                              io.EOF (graceful shutdown from server's
-//	                              side that didn't go through GoAway)
-//	StreamEnd:Canceled         — local ctx cancel (pool teardown,
-//	                              client app exit) or grpc CANCELED
-//	StreamEnd:DeadlineExceeded — ctx deadline or grpc DEADLINE_EXCEEDED
-//	StreamEnd:Unavailable      — transport-level break (TCP drop,
-//	                              connection recycler killed the channel,
-//	                              load balancer evicted the backend)
-//	StreamEnd:Internal         — server INTERNAL error
-//	StreamEnd:{Code}           — any other gRPC status code (verbatim)
-//	StreamEnd:Other            — no recognizable category (extremely rare)
-//	StreamEnd                  — err was nil (shouldn't happen since Recv
-//	                              only returns on error)
-func streamEndReason(err error) string {
-	if err == nil {
-		return "StreamEnd"
-	}
-	if errors.Is(err, io.EOF) {
-		return "StreamEnd:EOF"
-	}
-	if errors.Is(err, context.Canceled) {
-		return "StreamEnd:Canceled"
-	}
-	if errors.Is(err, context.DeadlineExceeded) {
-		return "StreamEnd:DeadlineExceeded"
-	}
-	if st, ok := status.FromError(err); ok {
-		return "StreamEnd:" + st.Code().String()
-	}
-	return "StreamEnd:Other"
-}
-
-// isAbnormalCloseReason returns true when the recorded close reason
-// looks like something we did NOT initiate cleanly. Clean paths:
-// EOF (server graceful), Canceled (client teardown / ctx cancel), and
-// the explicit client-initiated reasons stamped by handleGoAway /
-// heartbeatLoop / handleErrorResponse. Anything else — a StreamEnd
-// tagged with a transport-failure code, or the bare "StreamEnd" that
-// indicates Recv returned nil (which shouldn't happen) — is abnormal
-// and worth flagging.
-//
-// TODO(sushanb): move to a state-based classifier per mutianf's review
-// on #20215. Current reason-string scheme encodes state indirectly (via
-// CAS-once CloseReason stamped at each transition site) and gives finer
-// per-reason attribution for sessionz's close-reasons breakdown, but a
-// state-transition source of truth ("did we go New→Ready→Closing→
-// WaitServerClose→Closed cleanly?") is more robust — the whitelist
-// here has to be kept in lockstep with every new closeReasonLabel case.
-// Refactor when we add a new close-reason (or when a downstream
-// consumer wants the state-transition history directly).
-func isAbnormalCloseReason(reason string) bool {
-	switch reason {
-	case "StreamEnd:EOF", "StreamEnd:Canceled",
-		"GoAway", "MissedHeartbeat", "Error", "":
-		return false
-	}
-	return strings.HasPrefix(reason, "StreamEnd")
-}
-
 // heartbeatLoop watches the session's heartbeat deadline using a single Timer
 // that re-arms itself when a frame extends the deadline. The watchdog is
 // only enforced while at least one VRPC is in flight: the server emits

@@ -51,6 +51,12 @@ func abnormalOnCloseFor(t testing.TB, p *SessionPoolImpl, abnormal bool) {
 		return
 	}
 	sh := injectActiveSession(t, p, "active", time.Now())
+	// Simulate the state history that real session.Close() → handleClose
+	// produces: Ready → Closing → WSC → Closed, so prevStateAtClose = WSC.
+	// The fixture bypasses that path (calls onClose directly), so stamp
+	// prevStateAtClose here to match what noteAbnormalCloseIfAny expects
+	// on a client-initiated clean close.
+	sh.session.prevStateAtClose.Store(int32(StateWaitServerClose))
 	p.onClose(sh, nil)
 }
 
@@ -94,11 +100,15 @@ func TestConsecutiveFailures_UserReasonNotAbnormal(t *testing.T) {
 	// the second close instead of needing 10.
 	p.consecutiveFailureThreshold.Store(2)
 
-	// Fire "User" close-reason twice on ACTIVATED sessions — the
-	// closest fixture to Pool.Close's Phase-2 (activated sessions
-	// closed with REASON_USER). State-based gate must NOT count them.
+	// Fire "User" close-reason twice on ACTIVATED sessions that also
+	// went through WSC — the closest fixture to Pool.Close's Phase-2
+	// (session.Close() transitions through WSC, then onClose fires).
+	// State-history gate must NOT count them.
 	for i := 0; i < 2; i++ {
 		sh := injectActiveSession(t, p, "user-close", time.Now())
+		// Matches the state history real session.Close() → handleClose
+		// produces: prevStateAtClose = WSC when the WSC → Closed step ran.
+		sh.session.prevStateAtClose.Store(int32(StateWaitServerClose))
 		stampCloseReason(sh.session, "User")
 		p.onClose(sh, nil)
 	}

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