googleapis
diff --git a/‎bigtable/internal/transport/afe_picker.go‎
Lines changed: 6 additions & 0 deletions b/‎bigtable/internal/transport/afe_picker.go‎
Lines changed: 6 additions & 0 deletions
diff --git a/‎bigtable/internal/transport/connpool.go‎
Lines changed: 19 additions & 0 deletions b/‎bigtable/internal/transport/connpool.go‎
Lines changed: 19 additions & 0 deletions
diff --git a/‎bigtable/internal/transport/debug_tracer.go‎
Lines changed: 26 additions & 1 deletion b/‎bigtable/internal/transport/debug_tracer.go‎
Lines changed: 26 additions & 1 deletion
diff --git a/‎bigtable/internal/transport/pool_sizer.go‎
Lines changed: 32 additions & 12 deletions b/‎bigtable/internal/transport/pool_sizer.go‎
Lines changed: 32 additions & 12 deletions
diff --git a/‎bigtable/internal/transport/session.go‎
Lines changed: 12 additions & 0 deletions b/‎bigtable/internal/transport/session.go‎
Lines changed: 12 additions & 0 deletions
diff --git a/‎bigtable/internal/transport/session_debug.go‎
Lines changed: 20 additions & 0 deletions b/‎bigtable/internal/transport/session_debug.go‎
Lines changed: 20 additions & 0 deletions
diff --git a/‎bigtable/internal/transport/session_lifecycle.go‎
Lines changed: 15 additions & 2 deletions b/‎bigtable/internal/transport/session_lifecycle.go‎
Lines changed: 15 additions & 2 deletions
@@ -18,6 +18,12 @@ import (
 	"math/rand/v2"
 )
 
+// defaultAfeRandomSubsetSize is the default K for K-choice random draws
+// in LeastInFlightAfePicker / LeastLatencyAfePicker when the caller
+// doesn't specify one. Two candidates ("power of two choices") is the
+// standard K-choice draw size.
+const defaultAfeRandomSubsetSize = 2
+
 // PickCandidate is one AFE the picker considered during a K-choice draw,
 // with the cost value the picker's decision rule used to score it.
 // Cost's interpretation depends on the picker in play: NumOutstanding
 
@@ -1066,3 +1066,22 @@ func (m multiError) Error() string {
 	}
 	return fmt.Sprintf("%s (and %d other errors)", s, n-1)
 }
+
+// channelPickHintKey identifies the *atomic.Int32 destination the caller
+// wants BigtableChannelPool to publish the picked connEntry index into.
+type channelPickHintKey struct{}
+
+// ChannelPickHintInto returns a context that BigtableChannelPool will use to
+// publish the picked connEntry index into the supplied *atomic.Int32. The
+// caller can then read the value once the stream/invoke has returned.
+//
+// Used by Session creation to link sessions back to the channel they ride
+// on — surfaced in the sessionz / channelz debug UIs. Untouched by callers
+// that don't care: stampChannelPickHint short-circuits when the context
+// lacks the key. Passing dst == nil returns ctx unchanged.
+func ChannelPickHintInto(ctx context.Context, dst *atomic.Int32) context.Context {
+	if dst == nil {
+		return ctx
+	}
+	return context.WithValue(ctx, channelPickHintKey{}, dst)
+}
@@ -116,7 +116,32 @@ const (
 	tagSessionPoolStuckSessionSwept = "session_pool_stuck_session_swept"
 	tagSessionPoolDrainTimeout      = "session_pool_drain_timeout"
 	tagSessionPoolCreateFailed      = "session_pool_create_failed"
-	tagSessionPoolPickLostRace      = "session_pool_pick_lost_race"
+	// tagSessionPoolCreatePanic distinguishes a recovered panic inside
+	// Tick's createSession fanout (streamFactory / NewSession / hook
+	// wiring) from a plain error return (tagSessionPoolCreateFailed).
+	// The two paths have very different root causes — a panic indicates
+	// a client-side bug, an error is typically transient — so ops
+	// should be able to grep them apart in the debug-tag counters.
+	tagSessionPoolCreatePanic                = "session_pool_create_panic"
+	tagSessionPoolPickLostRace               = "session_pool_pick_lost_race"
+	tagSessionPoolConsecutiveFailuresTripped = "session_pool_consecutive_failures_tripped"
+	// tagSessionPoolNoBudget fires when createSession's budget.Acquire
+	// returns an error — either poolCtx cancel (teardown) or the
+	// throttler's NewSessionCreationPenalty window expired without an
+	// existing reservation being released. The count is the pool's
+	// "opens throttled" signal; sustained emission means the budget
+	// ceiling is too low for the offered load OR opens are hanging past
+	// the penalty window.
+	tagSessionPoolNoBudget = "session_pool_no_budget"
+
+	// tagSessionPoolCheckoutFailedCINil fires on SessionPoolImpl.Invoke's
+	// early return when CheckoutSession failed — pool returns
+	// InvokeResult{} with nil ClusterInfo, so stampAttempt downstream
+	// records TagSessionAttemptNilClusterInfo without any session ever
+	// being picked. Empirically dominates the nil-ClusterInfo population
+	// during pool cold-start (waiters ctx.Done before first session
+	// reaches Ready) and pool-close bursts (drainWaitersWithErr).
+	tagSessionPoolCheckoutFailedCINil = "session_pool_checkout_failed_ci_nil"
 
 	// sessionList bookkeeping violations.
 	//
 
@@ -17,6 +17,7 @@ package internal
 import (
 	"math"
 	"sync"
+	"sync/atomic"
 
 	spb "cloud.google.com/go/bigtable/apiv2/bigtablepb"
 )
@@ -53,10 +54,15 @@ type StatsFetcher func() *PoolStats
 // outside any peak-window cannot trigger a mass prune followed by a
 // cold-start scale-up on the next crest.
 type PoolSizer struct {
-	mu              sync.Mutex
-	fetcher         StatsFetcher
-	minSessions     int
-	maxSessions     int
+	mu      sync.Mutex
+	fetcher StatsFetcher
+	// minSessions / maxSessions are the authoritative pool-size bounds.
+	// Stored as atomics so hot-path readers (CheckoutSession gate,
+	// onClosing replace-check, createSession cap) can Load without
+	// taking s.mu — pool.go used to duplicate these into its own
+	// atomics; that duplication is now gone.
+	minSessions     atomic.Int32
+	maxSessions     atomic.Int32
 	headroomPct     float64 // Idle headroom as a fraction of sessions in use (e.g., 0.10 = 10%)
 	newSessionQLen  int     // server-driven per-session pending queue length; divides PendingCount
 	minIdleSessions int     // floor on the idle cushion so headroom never collapses to 0
@@ -87,16 +93,26 @@ func NewPoolSizer(fetcher StatsFetcher, minSessions, maxSessions int, headroomPc
 	if headroomPct <= 0 {
 		headroomPct = float64(defaultPoolConfig().GetHeadroom())
 	}
-	return &PoolSizer{
+	s := &PoolSizer{
 		fetcher:         fetcher,
-		minSessions:     minSessions,
-		maxSessions:     maxSessions,
 		headroomPct:     headroomPct,
 		newSessionQLen:  int(defaultPoolConfig().GetNewSessionQueueLength()),
 		minIdleSessions: defaultMinIdleSessions,
 	}
+	s.minSessions.Store(int32(minSessions))
+	s.maxSessions.Store(int32(maxSessions))
+	return s
 }
 
+// MinSessions returns the current pool-floor session count. Atomic load;
+// safe from any goroutine without taking the sizer lock. Callers that
+// need min/max/headroom as a consistent triple should use Decide instead.
+func (s *PoolSizer) MinSessions() int { return int(s.minSessions.Load()) }
+
+// MaxSessions returns the current pool-ceiling session count. Atomic
+// load; same non-locking contract as MinSessions.
+func (s *PoolSizer) MaxSessions() int { return int(s.maxSessions.Load()) }
+
 // UpdateConfig dynamically adjusts the sizer's capacity bounds,
 // headroom cushion, and per-session queue length at runtime. Called
 // from the server-config listener path; safe against concurrent
@@ -110,8 +126,12 @@ func (s *PoolSizer) UpdateConfig(config *spb.SessionClientConfiguration_SessionP
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	s.minSessions = int(config.MinSessionCount)
-	s.maxSessions = int(config.MaxSessionCount)
+	// min/max are atomic so hot-path readers (pool gate + snapshot) see
+	// the new bound the moment UpdateConfig returns; s.mu still brackets
+	// the write with headroom/qlen so a Decide walking the sizer sees a
+	// coherent config triple.
+	s.minSessions.Store(config.MinSessionCount)
+	s.maxSessions.Store(config.MaxSessionCount)
 	// Mirror the constructor guard: a zero or negative headroom from
 	// the server would render as HeadroomPct=0 on the loadz trace and
 	// collapse IdleHeadroom to the MinIdleSessions floor — the pool
@@ -196,8 +216,8 @@ func (s *PoolSizer) Decide() ScaleDecision {
 	defer s.mu.Unlock()
 
 	d := ScaleDecision{
-		MinSessions:     s.minSessions,
-		MaxSessions:     s.maxSessions,
+		MinSessions:     int(s.minSessions.Load()),
+		MaxSessions:     int(s.maxSessions.Load()),
 		HeadroomPct:     s.headroomPct,
 		NewSessionQLen:  s.newSessionQLen,
 		MinIdleSessions: s.minIdleSessions,
@@ -233,7 +253,7 @@ func (s *PoolSizer) Decide() ScaleDecision {
 		d.IdleHeadroom = s.minIdleSessions
 	}
 	d.DesiredRaw = d.SessionsInUse + d.IdleHeadroom
-	d.DesiredCapacity = clamp(d.DesiredRaw, s.minSessions, s.maxSessions)
+	d.DesiredCapacity = clamp(d.DesiredRaw, d.MinSessions, d.MaxSessions)
 
 	d.ImmediateCapacity = stats.ReadyCount
 	d.EventualCapacity = stats.ReadyCount + stats.StartingCount
 
@@ -195,6 +195,18 @@ type Session struct {
 	// Embedded so bare field access (s.tracer, s.okRpcs, s.recordEvent,
 	// ...) continues to compile once vRPC / lifecycle land.
 	sessionDebug
+
+	// loops tracks readLoop + heartbeatLoop so a supervising owner
+	// (SessionPoolImpl.Close) can wait for them to fully unwind — through
+	// their notifyClosing / notifyClosed callback chains — before it
+	// returns. Prevents readLoop's recordClose from racing package-level
+	// metric var writes across test boundaries.
+	loops sync.WaitGroup
+
+	// closeErr preserves the raw Recv error handed to handleClose. The
+	// pool surfaces this on consecutive-failure breaker trips so operators
+	// see the underlying server rejection instead of only the sentinel.
+	closeErr atomic.Pointer[error]
 }
 
 // SessionOption configures a Session at construction time.
 
@@ -405,3 +405,23 @@ func WithSessionLogger(logger *log.Logger) SessionOption {
 func WithSessionPoolName(name string) SessionOption {
 	return func(s *Session) { s.tracer.setPoolName(name) }
 }
+
+// setCloseErr records the raw error that ended the stream. First writer
+// wins so a follow-up close path (e.g. cancelActiveRPCs) can't overwrite
+// the original cause. Nil is treated as "no error" and ignored.
+func (s *Session) setCloseErr(err error) {
+	if err == nil {
+		return
+	}
+	s.closeErr.CompareAndSwap(nil, &err)
+}
+
+// closeError returns the raw error that ended the stream, or nil if
+// none was recorded. Consulted by the pool to surface the underlying
+// server rejection on consecutive-failure trips.
+func (s *Session) closeError() error {
+	if p := s.closeErr.Load(); p != nil {
+		return *p
+	}
+	return nil
+}
@@ -75,11 +75,23 @@ func (s *Session) Start(ctx context.Context, req *spb.OpenSessionRequest) error
 	// is safe.
 	s.hooks.onStart(ctx)
 
-	go s.readLoop(ctx)
-	go s.heartbeatLoop(ctx)
+	// Track readLoop + heartbeatLoop so WaitGoroutines can block until
+	// their callback chains (notifyClosed → recordClose, etc.) have
+	// unwound. Owners (SessionPoolImpl.Close) call WaitGoroutines during
+	// teardown so no session-owned goroutine outlives the pool.
+	s.loops.Add(2)
+	go func() { defer s.loops.Done(); s.readLoop(ctx) }()
+	go func() { defer s.loops.Done(); s.heartbeatLoop(ctx) }()
 	return nil
 }
 
+// WaitGoroutines blocks until readLoop and heartbeatLoop have fully
+// returned (including their notifyClosed / recordClose callback
+// chains). No-op if Start was never called.
+func (s *Session) WaitGoroutines() {
+	s.loops.Wait()
+}
+
 // ForceClose immediately transitions the session to StateClosed and cancels
 // every in-flight RPC. It is safe to call multiple times; only the first call
 // fires the tracer.recordClose and hooks.onClose callbacks.
@@ -450,6 +462,7 @@ func (s *Session) handleClose(err error) {
 	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