fix(bigtable): release session-creation budget on GO_AWAY during STARTING

A session that received a server GO_AWAY (or was force-closed) while still
STARTING transitions STARTING -> CLOSING -> WAIT_SERVER_CLOSE and closes with
prevState == WAIT_SERVER_CLOSE. The only failure-side budget release in
onSessionClose lived inside the abnormal-close branch
(if prevState != WAIT_SERVER_CLOSE), guarded by prevState == STARTING, so that
path was skipped and the reserved concurrency slot leaked permanently. Over time
this exhausts the session-creation budget even at very low QPS.

Track each reservation on the SessionHandle and release it exactly once: as a
success when the session reaches READY (onSessionReady), or as a failure when it
terminates without ever becoming READY (onSessionClose), before the pool-state
and abnormal-close gates. Keying off the reservation instead of the close-time
prevState makes the release exactly-once and also covers the latent leak where
onSessionReady early-returns (poolState != STARTED) without releasing.

Adds a regression test that drives a session receiving a GO_AWAY before the open
handshake, backed by a new go_away_before_open fake-proto field and matching
SessionHandler branch. The test fails without the fix and passes with it.

Author: mutianf

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java

@@ -154,6 +154,18 @@ private enum PoolState {
   @GuardedBy("this")
   private final SessionCreationBudget budget;
 
+  // Handles that reserved a session-creation budget slot in createSession() and have not yet
+  // released it. A slot is released exactly once: as a success when the session reaches READY
+  // (onSessionReady), or as a failure when the session terminates without ever becoming READY
+  // (onSessionClose). Tracking the reservation on the handle -- instead of inferring it from the
+  // close-time prevState -- is what makes the release exactly-once: a session that goes
+  // STARTING -> WAIT_SERVER_CLOSE (e.g. a server GO_AWAY before the open handshake completes)
+  // closes with prevState == WAIT_SERVER_CLOSE, which the abnormal-close branch skips, so the
+  // prevState == STARTING check alone would leak the slot forever. SessionHandle uses identity
+  // equality, so a plain HashSet keys on the handle instance.
+  @GuardedBy("this")
+  private final Set<SessionHandle> sessionsHoldingBudget = new HashSet<>();
+
   private final ClientConfigurationManager configManager;
   private final ClientConfigurationManager.ListenerHandle configListenerHandle;
 
@@ -424,6 +436,9 @@ private synchronized void createSession(OpenParams openParams) {
         SessionStream stream = factory.createNew();
         Session session = new SessionImpl(metrics, info, sessionNum++, stream, timer);
         SessionHandle handle = sessions.newHandle(session);
+        // Bind the budget reservation made by tryReserveSession() above to this handle so it is
+        // released exactly once when the session becomes READY or terminates.
+        sessionsHoldingBudget.add(handle);
 
         Metadata localMd = new Metadata();
         localMd.merge(openParams.metadata());
@@ -498,7 +513,11 @@ private synchronized void onSessionReady(SessionHandle handle, OpenSessionRespon
     }
     handle.onSessionStarted();
 
-    budget.onSessionCreationSuccess();
+    // Release the reservation as a success. Guarded by the set so we release exactly once even if
+    // onSessionReady were ever delivered more than once.
+    if (sessionsHoldingBudget.remove(handle)) {
+      budget.onSessionCreationSuccess();
+    }
 
     // handle pending rpcs
     tryDrainPendingRpcs();
@@ -558,6 +577,17 @@ private void onSessionClose(
 
       handle.onSessionClosed(prevState);
 
+      // Release the budget reservation as a failure if this session still holds one, i.e. it
+      // terminated without ever reaching READY. Done here -- before the pool-state and
+      // abnormal-close gates below -- because a session that received a GO_AWAY / was force-closed
+      // while STARTING closes with prevState == WAIT_SERVER_CLOSE, which the abnormal-close branch
+      // skips; keying off the reservation instead of prevState is what prevents the slot from
+      // leaking. Sessions that reached READY were already removed from the set in onSessionReady,
+      // so this never double-releases.
+      if (sessionsHoldingBudget.remove(handle)) {
+        budget.onSessionCreationFailure();
+      }
+
       // If the pool is closed then there is nothing else to do
       // dont need to create a replacement session and pending vRpcs get cleaned up in close()
       if (poolState == PoolState.CLOSED) {
@@ -582,9 +612,8 @@ private void onSessionClose(
           toBeClosed = popClosableRpcs();
         }
 
-        if (prevState == SessionState.STARTING) {
-          budget.onSessionCreationFailure();
-        }
+        // Budget release for STARTING-phase closes is handled above via sessionsHoldingBudget,
+        // which also covers the STARTING -> WAIT_SERVER_CLOSE (GO_AWAY) path this branch skips.
 
         // TODO: backoff creating a new session when consecutive failures > max?
         if (poolSizer.handleSessionClose(StatusProto.fromStatusAndTrailers(status, trailers))) {

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java

@@ -612,6 +612,108 @@ public void test() throws Exception {
     }
   }
 
+  @Nested
+  class GoAwayBeforeReadyBudgetRelease {
+
+    private FakeClock fakeClock;
+    private BigtableTimer mockTimer;
+    private FakeSessionService fakeService;
+    private Server server;
+    private ChannelPool channelPool;
+    private SessionPoolImpl<OpenFakeSessionRequest> sessionPool;
+    private SessionCreationBudget budget;
+
+    private final Duration penalty = Duration.ofMinutes(1);
+
+    @BeforeEach
+    void setUp() throws Exception {
+      fakeClock = new FakeClock(Instant.now());
+      // Mock timer so the pool's watchdog / AFE-prune / create-session-retry ticks never fire on
+      // their own -- the budget stays quiescent after the session lifecycle completes, leaving the
+      // test thread as the only actor that touches it.
+      mockTimer = mock(BigtableTimer.class, Mockito.RETURNS_DEEP_STUBS);
+
+      // Budget of 1 so exactly one session goes STARTING; any extra create attempts fail to reserve
+      // and just schedule a (no-op) retry on the mock timer.
+      budget = new SessionCreationBudget(1, penalty, fakeClock);
+
+      fakeService = new FakeSessionService(executor);
+      server = FakeServiceBuilder.create(fakeService).intercept(new PeerInfoInterceptor()).start();
+
+      channelPool =
+          new SingleChannelPool(
+              Suppliers.ofInstance(
+                  Grpc.newChannelBuilderForAddress(
+                          "localhost", server.getPort(), InsecureChannelCredentials.create())
+                      .build()));
+      channelPool.start();
+
+      sessionPool =
+          new SessionPoolImpl<>(
+              metrics,
+              FeatureFlags.getDefaultInstance(),
+              CLIENT_INFO,
+              configManager,
+              channelPool,
+              CallOptions.DEFAULT,
+              FakeDescriptor.FAKE_SESSION,
+              "fake-pool",
+              mockTimer,
+              MoreExecutors.directExecutor(),
+              budget);
+    }
+
+    @AfterEach
+    void tearDown() {
+      sessionPool.close(
+          CloseSessionRequest.newBuilder()
+              .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER)
+              .setDescription("close session")
+              .build());
+      channelPool.close();
+      server.shutdownNow();
+    }
+
+    @Test
+    void goAwayWhileStartingReleasesBudget() throws Exception {
+      // Drive a single session that receives a GO_AWAY before the open handshake, so it transitions
+      // STARTING -> WAIT_SERVER_CLOSE and closes with prevState == WAIT_SERVER_CLOSE. The
+      // abnormal-close branch (prevState != WAIT_SERVER_CLOSE) is skipped, so before the fix the
+      // STARTING budget slot -- which was released only on the prevState == STARTING path -- leaked
+      // permanently. The fix keys the release off the reservation set, releasing it here as a
+      // creation failure.
+      sessionPool.start(
+          OpenFakeSessionRequest.newBuilder().setGoAwayBeforeOpen(true).build(), new Metadata());
+
+      // Poll for the leaked slot to become reclaimable. onSessionClose(WAIT_SERVER_CLOSE) records a
+      // creation failure (fix only); advancing past the penalty lets tryReserveSession drain it and
+      // succeed. Access is guarded by the pool monitor (the pool's synchronized methods lock on the
+      // pool instance) so it doesn't race the session-close callback that mutates the budget. With
+      // the bug no slot is ever released and this stays false until the deadline.
+      boolean recovered = false;
+      long deadlineMs = System.currentTimeMillis() + 5_000;
+      while (System.currentTimeMillis() < deadlineMs) {
+        synchronized (sessionPool) {
+          fakeClock.increment(penalty.plusMillis(1));
+          if (budget.tryReserveSession()) {
+            recovered = true;
+            break;
+          }
+        }
+        Thread.sleep(20);
+      }
+
+      assertThat(recovered).isTrue();
+
+      // Sanity check that the client actually walked the GO_AWAY-before-open path: on receiving the
+      // GO_AWAY while STARTING it moves to WAIT_SERVER_CLOSE and sends a CloseSession.
+      boolean sentClose =
+          fakeService.getSessionRequests().stream()
+              .anyMatch(r -> r.getPayloadCase() == SessionRequest.PayloadCase.CLOSE_SESSION);
+      assertThat(sentClose).isTrue();
+    }
+  }
+
   @Test
   public void refreshConfigTest() throws Exception {
     OpenSessionRequest refreshRequest =

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/SessionHandler.java

@@ -141,6 +141,23 @@ public State onMessage(SessionRequest sReq) {
         throw new RuntimeException(e);
       }
 
+      if (oReq.getGoAwayBeforeOpen()) {
+        // Send a GO_AWAY before the SessionParameters / OpenSession handshake, so the client
+        // receives it while the session is still STARTING. The client will transition to
+        // WAIT_SERVER_CLOSE, send a CloseSession, and half-close; RunningState terminates the
+        // stream successfully on either of those. goAwayDelay is left large so RunningState's own
+        // lifecycle GO_AWAY timer never fires during the test.
+        handler.writeResponse(
+            SessionResponse.newBuilder()
+                .setGoAway(
+                    GoAwayResponse.newBuilder()
+                        .setReason("session_expiry")
+                        .setDescription("go away before open handshake")
+                        .build())
+                .build());
+        return new RunningState(handler, Duration.ofMinutes(1), oReq);
+      }
+
       if (oReq.hasStreamError()) {
         Status status =
             Status.fromCodeValue(oReq.getStreamError().getStatus().getCode())

java-bigtable/google-cloud-bigtable/src/test/proto/google/bigtable/v2/fake.proto

@@ -37,6 +37,12 @@ message OpenFakeSessionRequest {
   StreamError stream_error = 3;
   SessionParametersResponse session_params = 4;
 
+  // If set, the server sends a GO_AWAY immediately without ever sending the
+  // SessionParameters / OpenSession handshake responses, i.e. the session receives a
+  // GO_AWAY while it is still STARTING. Used to exercise the STARTING -> WAIT_SERVER_CLOSE
+  // session-creation-budget release path.
+  bool go_away_before_open = 9;
+
   // session steady phase
   message Action {
     google.protobuf.Duration delay = 1;