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

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

@@ -300,12 +300,30 @@ public void forceClose(CloseSessionRequest closeReason) {
             return;
           }
 
-          updateState(SessionState.WAIT_SERVER_CLOSE);
+          SessionState prevState = state;
           this.closeReason = closeReason;
+          VRpcImpl<?, ?, ?> localRpc = currentRpc;
+          currentRpc = null;
+          updateState(SessionState.CLOSED);
+          debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_abnormal_close");
 
           // Not sending the CloseSessionRequest because cancel() will just drop it
-          stream.forceClose(closeReason.getDescription(), null);
-          // Listeners will be notified by dispatchStreamClosed
+          try {
+            stream.forceClose(closeReason.getDescription(), null);
+          } catch (Throwable t) {
+            debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_force_close_failed");
+            logger.log(
+                Level.WARNING,
+                String.format(
+                    "Session error: %s Exception while force-closing stream", info.getLogName()),
+                t);
+          }
+
+          notifyTerminalClose(
+              Status.UNAVAILABLE.withDescription(closeReason.getDescription()),
+              new Metadata(),
+              localRpc,
+              prevState);
         });
   }
 
@@ -778,6 +796,9 @@ private void handleUnknownResponseMessage(SessionResponse message) {
 
   private void dispatchStreamClosed(Status status, Metadata trailers) {
     sessionSyncContext.throwIfNotInThisSynchronizationContext();
+    if (state == SessionState.CLOSED) {
+      return;
+    }
     SessionState prevState = state;
 
     if (!status.isOk()) {
@@ -800,15 +821,13 @@ private void dispatchStreamClosed(Status status, Metadata trailers) {
               info.getLogName(), state, status);
       logger.warning(msg);
 
-      if (state == SessionState.CLOSED) {
-        return;
+      if (closeReason == null) {
+        closeReason =
+            CloseSessionRequest.newBuilder()
+                .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR)
+                .setDescription("Unexpected session close with status: " + status.getCode())
+                .build();
       }
-
-      closeReason =
-          CloseSessionRequest.newBuilder()
-              .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR)
-              .setDescription("Unexpected session close with status: " + status.getCode())
-              .build();
     }
 
     VRpcImpl<?, ?, ?> localVRpc = currentRpc;

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

@@ -261,6 +261,12 @@ void onSessionClosing() {
     }
 
     void onSessionClosed(SessionState prevState) {
+      if (prevState == SessionState.NEW) {
+        throw new IllegalStateException("NEW session was closed");
+      }
+      if (prevState == SessionState.CLOSED) {
+        throw new IllegalStateException("double close");
+      }
       // Always drop from allSessions on the way out, even if the branch below throws — otherwise a
       // stranded handle blocks drainedFuture and pool.awaitTerminated hangs until the shared close
       // deadline elapses.
@@ -273,35 +279,17 @@ void onSessionClosed(SessionState prevState) {
         afe.ifPresent(afeHandle -> afeHandle.refCount--);
 
         // NOTE: don't need to update vRpc counters, onVRpcFinish will have been invoked already
-        switch (prevState) {
-          case NEW:
-            throw new IllegalStateException("NEW session was closed");
-          case STARTING:
-            poolStats.startingCount--;
-            break;
-          case READY:
-            {
-              // afe may be empty if SessionPoolImpl.onSessionReady early-returned on poolState !=
-              // STARTED (pool closed after SessionImpl transitioned to READY but before
-              // handle.onSessionStarted ran). Skip the AFE bookkeeping cleanly rather than NPE.
-              if (afe.isPresent()) {
-                AfeHandle afeHandle = afe.get();
-                // If the session was available & idle, then we need to remove it
-                if (afeHandle.sessions.remove(this)) {
-                  poolStats.readyCount--;
-                  if (afeHandle.sessions.isEmpty()) {
-                    afesWithReadySessions.remove(afeHandle);
-                  }
-                }
-              }
-              break;
+        if (!afe.isPresent()) {
+          poolStats.startingCount--;
+        } else {
+          AfeHandle afeHandle = afe.get();
+          // If the session was available & idle, then we need to remove it
+          if (afeHandle.sessions.remove(this)) {
+            poolStats.readyCount--;
+            if (afeHandle.sessions.isEmpty()) {
+              afesWithReadySessions.remove(afeHandle);
             }
-          case CLOSING:
-          case WAIT_SERVER_CLOSE:
-            // noop
-            break;
-          case CLOSED:
-            throw new IllegalStateException("double close");
+          }
         }
       } finally {
         allSessions.remove(this);

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

@@ -419,8 +419,10 @@ public void start(OpenReqT openReq, Metadata md) {
       poolState = PoolState.STARTED; // TODO: maybe need a READY state as well?
 
       // Pre-start
-      for (int i = poolSizer.getScaleDelta(); i > 0; i--) {
-        createSession(openParams);
+      while (poolSizer.getScaleDelta() > 0) {
+        if (!createSession(openParams)) {
+          break;
+        }
       }
 
       watchdog.start();
@@ -449,7 +451,7 @@ private int getMaxConsecutiveFailures(ClientConfigurationManager configManager)
   }
 
   @GuardedBy("poolLock")
-  private void createSession(OpenParams openParams) {
+  private boolean createSession(OpenParams openParams) {
     if (!budget.tryReserveSession()) {
       debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_pool_no_budget");
       logger.fine(
@@ -461,7 +463,7 @@ private void createSession(OpenParams openParams) {
       // after failing to create any sessions and exhausting all the budget, we'll retry session
       // creation once the budget becomes available so the VRpc still has a chance to succeed.
       maybeScheduleCreateSessionRetry();
-      return;
+      return false;
     }
 
     // Explicit create session streams in a detached context
@@ -507,6 +509,7 @@ public void onClose(SessionState prevState, Status status, Metadata trailers) {
               }
             });
       }
+      return true;
     } catch (RuntimeException | Error e) {
       // A synchronous failure here (e.g. factory.createNew, the SessionImpl constructor, or
       // metadata merge) means no terminal session callback will ever run for this reservation, so
@@ -526,6 +529,7 @@ public void onClose(SessionState prevState, Status status, Metadata trailers) {
       }
       // Let the pool recover instead of running permanently short a session.
       maybeScheduleCreateSessionRetry();
+      return false;
     } finally {
       Context.ROOT.detach(prevContext);
     }
@@ -549,8 +553,10 @@ private void maybeScheduleCreateSessionRetry() {
               poolLock.lock();
               try {
                 retryCreateSessionFuture = null;
-                if (poolState != PoolState.CLOSED && poolSizer.getScaleDelta() > 0) {
-                  createSession(openParams);
+                while (poolState != PoolState.CLOSED && poolSizer.getScaleDelta() > 0) {
+                  if (!createSession(openParams)) {
+                    break;
+                  }
                 }
               } finally {
                 poolLock.unlock();
@@ -692,15 +698,20 @@ private void onSessionClose(
       // Handle abnormal close. This can happen if the Session was aborted due to underlying stream
       // termination
       if (prevState != SessionState.WAIT_SERVER_CLOSE) {
-        consecutiveFailures++;
-        if (status.getCode() == Status.Code.UNIMPLEMENTED) {
-          consecutiveUnimplementedFailures++;
-        } else {
-          consecutiveUnimplementedFailures = 0;
-        }
-        // TODO: decide if max consecutive failures should be capped per client
-        if (consecutiveFailures >= getMaxConsecutiveFailures(configManager)) {
-          toBeClosed = popClosableRpcs();
+        // Only count failures during STARTING against the consecutive failure budget.
+        // Drops of established READY sessions (e.g. heartbeat misses) are transport failures
+        // rather than session establishment failures and should not fail pending vRPCs.
+        if (prevState == SessionState.STARTING) {
+          consecutiveFailures++;
+          if (status.getCode() == Status.Code.UNIMPLEMENTED) {
+            consecutiveUnimplementedFailures++;
+          } else {
+            consecutiveUnimplementedFailures = 0;
+          }
+          // TODO: decide if max consecutive failures should be capped per client
+          if (consecutiveFailures >= getMaxConsecutiveFailures(configManager)) {
+            toBeClosed = popClosableRpcs();
+          }
         }
 
         // Budget release for STARTING-phase closes is handled above via sessionsHoldingBudget,

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

@@ -524,6 +524,29 @@ void testHeartbeat() throws Exception {
     assertThat(sessionListener.popUntil(Status.class)).isOk();
   }
 
+  @Test
+  void missedHeartbeatDeliversUnavailableStatus() throws Exception {
+    SessionImpl session = new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer);
+
+    FakeSessionListener sessionListener = new FakeSessionListener();
+    OpenSessionRequest openSessionRequest =
+        OpenSessionRequest.newBuilder()
+            .setPayload(OpenFakeSessionRequest.getDefaultInstance().toByteString())
+            .build();
+    session.start(openSessionRequest, new Metadata(), sessionListener);
+    assertThat(sessionListener.popUntil(OpenSessionResponse.class))
+        .isInstanceOf(OpenSessionResponse.class);
+
+    session.forceClose(
+        CloseSessionRequest.newBuilder()
+            .setReason(CloseSessionReason.CLOSE_SESSION_REASON_MISSED_HEARTBEAT)
+            .setDescription("missed heartbeat")
+            .build());
+
+    Status status = sessionListener.popUntil(Status.class);
+    assertThat(status.getCode()).isEqualTo(Status.Code.UNAVAILABLE);
+  }
+
   @Test
   void testCancel() throws Exception {
     SessionImpl session = new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer);
@@ -864,17 +887,10 @@ void sessionOpenTimeoutForcesClose() throws Exception {
     capturing.executors.get(0).execute(capturing.tasks.get(0));
 
     // The session must force-close and drive the listener to a terminal Status.
-    assertWithMessage("terminal status should be delivered after open timeout")
-        .that(sessionListener.popUntil(Status.class))
-        .isNotNull();
-    sw.reset().start();
-    while (session.getState() != Session.SessionState.WAIT_SERVER_CLOSE
-        && session.getState() != Session.SessionState.CLOSED
-        && sw.elapsed(TimeUnit.SECONDS) < 5) {
-      Thread.sleep(10);
-    }
-    assertThat(session.getState())
-        .isAnyOf(Session.SessionState.WAIT_SERVER_CLOSE, Session.SessionState.CLOSED);
+    Status status = sessionListener.popUntil(Status.class);
+    assertThat(status.getCode()).isEqualTo(Status.Code.UNAVAILABLE);
+    assertThat(sessionListener.getLastPrevState()).isEqualTo(Session.SessionState.STARTING);
+    assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED);
   }
 
   // endregion

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

@@ -412,6 +412,54 @@ void testClosingInUseToClosed() {
     assertThat(stats.getExpectedCapacity()).isEqualTo(0);
   }
 
+  @Test
+  void testStartingSessionClosedAfterWaitServerClose_decrementsStartingCount() {
+    SessionList list = new SessionList();
+    PoolStats stats = list.getStats();
+
+    fakeSession.state = SessionState.STARTING;
+    SessionHandle handle = list.newHandle(fakeSession);
+
+    assertThat(stats.getStartingCount()).isEqualTo(1);
+    assertThat(stats.getExpectedCapacity()).isEqualTo(1);
+
+    // When a starting session terminates with WAIT_SERVER_CLOSE (e.g. server sent GOAWAY before
+    // open response, or normal close before open), startingCount and expectedCapacity must
+    // decrement cleanly.
+    handle.onSessionClosed(SessionState.WAIT_SERVER_CLOSE);
+
+    assertThat(list.getAllSessions()).isEmpty();
+    assertThat(stats.getStartingCount()).isEqualTo(0);
+    assertThat(stats.getExpectedCapacity()).isEqualTo(0);
+  }
+
+  @Test
+  void testReadyIdleSessionClosedAfterWaitServerClose_cleansUpReadySessionsAndReadyCount() {
+    SessionList list = new SessionList();
+    PoolStats stats = list.getStats();
+
+    fakeSession.state = SessionState.STARTING;
+    SessionHandle handle = list.newHandle(fakeSession);
+
+    fakeSession.state = SessionState.READY;
+    handle.onSessionStarted();
+
+    assertThat(stats.getReadyCount()).isEqualTo(1);
+    assertThat(list.getAfesWithReadySessions()).hasSize(1);
+    assertThat(stats.getStartingCount()).isEqualTo(0);
+
+    // When an idle session terminates with WAIT_SERVER_CLOSE without prior onSessionClosing()
+    // (e.g. session.close() or direct termination), readyCount and AFE handles must be cleaned up
+    // cleanly.
+    handle.onSessionClosed(SessionState.WAIT_SERVER_CLOSE);
+
+    assertThat(list.getAfesWithReadySessions()).isEmpty();
+    assertThat(list.getAllSessions()).isEmpty();
+    assertThat(stats.getStartingCount()).isEqualTo(0);
+    assertThat(stats.getReadyCount()).isEqualTo(0);
+    assertThat(stats.getExpectedCapacity()).isEqualTo(0);
+  }
+
   @Test
   void testAwaitCloseToSoftClosed() {
     SessionList list = new SessionList();