fix(bigtable): fix abnormal session closures and scale-up in session pool

[nimf]

This PR fixes session replacement and pool sizing issues in the Jetstream session pool:

• Preserve pre-failure state on forceClose() (SessionImpl): forceClose() previously set state directly to WAIT_SERVER_CLOSE, tricking SessionPoolImpl into treating abnormal closures (handshake timeouts, missed heartbeats) as graceful closes and skipping replacement. It now captures the originating state (STARTING or READY) and reports it to the pool listener.
• Idempotent cleanup in SessionList.onSessionClosed(): Switched from checking prevState to inspecting handle state (!afe.isPresent() for starting count; afeHandle.sessions.remove() for ready count). This prevents counter leaks on unexpected closure sequences and avoids double-decrements.
• Multi-session scale-up on non-hot paths (SessionPoolImpl): Updated start(), onSessionClose(), onSessionGoAway(), and retry scheduling to create sessions while poolSizer.getScaleDelta() > 0 and budget allows, while intentionally keeping the request hot path (PendingVRpc.start()) to a single session creation to avoid caller latency overhead.

[gemini-code-assist]

Code Review

This pull request refactors session management and pool scaling logic in the Bigtable client, ensuring that abnormally closed or failed sessions are properly replaced to maintain the minimum session count, and simplifying state transition handling. The review feedback suggests improving consistency in SessionPoolImpl.java by using a local openParams variable during session replacement, and refactoring manual polling in tests to use Awaitility for better robustness.

[mutianf]

Code Review Findings

CONFIRMED bugs

1. SessionPoolImpl.java:694 — Force-close incorrectly increments consecutiveFailures (most critical)

The new forceClose path now produces prevState=READY instead of prevState=WAIT_SERVER_CLOSE, so heartbeat-miss closures now increment consecutiveFailures. In a bulk-failure scenario (e.g. all 5 sessions simultaneously miss heartbeat), the counter can hit the threshold and invoke popClosableRpcs(), rejecting all pending RPCs even though the server may be healthy. Previously, the WAIT_SERVER_CLOSE state caused this path to be bypassed.

2. SessionImplTest.java:871 — Test polls for WAIT_SERVER_CLOSE that forceClose never produces

sessionOpenTimeoutForcesClose has a while loop waiting for WAIT_SERVER_CLOSE || CLOSED, but forceClose now goes directly STARTING → CLOSED, skipping WAIT_SERVER_CLOSE entirely. The loop is a dead no-op that passes only because CLOSED is also accepted. If a future change breaks the CLOSED path, this test will spin with no meaningful timeout protection.

3. SessionPoolImplTest.java:686 — Flaky exact-count assertion after non-atomic wait

abnormalCloseScalesUpMultipleSessions calls waitForOpenRequestCount(openCountBefore + 5, ...) then immediately asserts isEqualTo(openCountBefore + 5). The count can increment between the two calls (e.g. if onSessionReady fires and triggers another scale-up), causing a spurious failure.

4. SessionImpl.java:86 and :500 — Stale comments reference the removed WAIT_SERVER_CLOSE path

Both comments describe the STARTING → WAIT_SERVER_CLOSE → onSessionClose flow, which no longer exists after the forceClose refactor. These are actively misleading for anyone debugging the force-close path or tracing session-budget accounting.

---

PLAUSIBLE bugs

5. SessionPoolImplTest.java:626 — Same racy exact-count assertion in pendingCallOnHotPathCreatesOnlyOneSession

Same pattern as finding #3: asserts isEqualTo(openCountBefore + 1) immediately after waitForOpenRequestCount(openCountBefore + 1).

6. SessionList.java:279 — Comment explaining READY-state startingCount-- edge case was removed

The refactored onSessionClosed now silently runs startingCount-- for prevState=READY when afe is absent (pool closed before onSessionStarted ran). This is correct behavior, but the comment that explained why was removed. A future maintainer will likely see startingCount-- for prevState=READY and revert it as wrong, reintroducing the counter leak.

7. SessionPoolImpl.java:719 — Scale-up loops in onSessionClose/onSessionGoAway missing CLOSED guard

The inline scale-up while loops lack the poolState != PoolState.CLOSED guard that maybeScheduleCreateSessionRetry has. Safe today because poolLock is held throughout, but brittle if lock scope is ever relaxed — sessions could be created after the pool starts draining with no guard to catch it.

8. SessionPoolImpl.java:648 — GoAway replacement vs. scale-up use different openParams with no explanation

The replacement session is created with handle.getSession().getOpenParams() while the scale-up loop uses pool-level openParams. The net result appears correct but the inconsistency is undocumented and could confuse future debugging or get broken in a refactor.

9. .cloud/helpers/common.sh:57 — contains() leaves $2 unquoted in a bash [[ glob pattern

contains() uses [[ " $1 " == *" $2 "* ]]. If $2 contains glob metacharacters (*, ?, [), the pattern is treated as a glob rather than a literal string match, producing false positives. Should be [[ " $1 " == *" ${2} "* ]] with quoting, or use [[ " $1 " == *" $2 "* ]] only with set -f or fnmatch alternatives.

[mutianf]

Review on my agent's confirmed bug 1 :)

I think changing forceClose from WAIT_SERVER_CLOSE to CLOSE actually makes it semantically correct. It matches the original design:

Client hard close
Client cancels the underlying stream. Client hard close means that client stops sending vRPCs on the sessions immediately and cancels all active vRPCs on the session. 

Scenarios:
Client stops receiving heartbeat or responses on a session after the keepalive time. 
Client logic threw an exception due to a bug

Actions:
Client: stops receiving heartbeat or responses on a session after the keepalive time.
Session pool: remove the session from session pool and create a new session.
Active vRPCs: client cancels all active vRPCs on the session. If the vRPCs is idempotent and there’s still a deadline, queue it for retry on other available sessions. 
Pending vRPCs: continue to queue on other available sessions until the continuous failure budget runs out. After that the client will cancel all pending vRPCs.

Continuous failure budget: this will be a number from ClientConfiguration. We want to limit the number of retries a vRPCs can have when the retry is caused by underlying session hanging. If we have a query of death, we want to avoid the problematic vRPC crashing all the AFEs in the pools by limiting the number of retries on other available sessions. 

However, with this change, there is the risk:

Heartbeat misses ARE session failures from the pool's perspective — the session is gone, a replacement must be created, and pending vRPCs need to tolerate this. The budget is supposed to cap how many times a vRPC retries before giving up, precisely to prevent a "query of death" from thrashing the pool.

But there's a shape problem: consecutiveFailures is a pool-level counter, not a per-vRPC counter. The scenario that breaks it:

t=0: Sessions S1–S5 all miss heartbeat simultaneously
→ forceClose fires on all 5
→ consecutiveFailures increments 5 times before any replacement is READY
→ if threshold=5, popClosableRpcs() fires immediately
→ all pending vRPCs fail with REJECTED (no retry)

But: S1–S5 failing simultaneously is one "event" (e.g., brief network blip),
not five independent consecutive failures.
Pending vRPCs never even got a chance to run on a replacement.

The counter is meant to detect persistent failure (every new session I create immediately dies), but it's sensitive to burst failure (many sessions die at once in one wave). Two possible fixes:

Option A — Per-session failure counts toward budget only when it dies before serving any vRPC:
consecutiveFailures++ only if the session never became READY
(i.e., prevState == STARTING when it closed abnormally)
Heartbeat misses happen on READY sessions → don't count. Only sessions that fail to open increment the counter. This directly models "the pool keeps trying to create sessions and they all immediately fail."

Option B — Threshold scales with pool size:
The threshold from ClientConfiguration should be tuned so that threshold > pool_size to avoid a single burst triggering it.

Maybe we should consider option A so this change is less risky.