fix(bigtable): fix session creation leaks

Author: mutianf

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

@@ -78,6 +78,21 @@ public class SessionImpl implements Session, VRpcSessionApi {
           .setDescription("missed heartbeat")
           .build();
 
+  @VisibleForTesting
+  // Upper bound on how long a session may remain STARTING before we give up on the open handshake
+  // and force-close it. Without this, a stream that connects but never delivers the OpenSession
+  // response (or a GoAway) leaves the session wedged in STARTING forever, holding the
+  // session-creation-budget slot it reserved and never firing the terminal callback that would
+  // release it. Force-closing routes STARTING -> WAIT_SERVER_CLOSE -> onSessionClose, which frees
+  // the budget.
+  static final Duration OPEN_SESSION_TIMEOUT = Duration.ofSeconds(90);
+
+  private static final CloseSessionRequest OPEN_TIMEOUT_CLOSE_REQUEST =
+      CloseSessionRequest.newBuilder()
+          .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR)
+          .setDescription("session open timeout")
+          .build();
+
   private final Clock clock;
   private final BigtableTimer timer;
   // Serializes all session state mutations. Stream callbacks and the heartbeat tick dispatch
@@ -130,6 +145,10 @@ public class SessionImpl implements Session, VRpcSessionApi {
   // transitions so the wheel doesn't carry a no-op entry until the next fire.
   @Nullable private BigtableTimer.Timeout heartbeatTimeout;
 
+  // Handle for the open-handshake deadline (armed while STARTING). Cancelled the moment the
+  // session leaves STARTING so a session that opens normally never trips it.
+  @Nullable private BigtableTimer.Timeout openTimeout;
+
   // Set by the global SyncContext handler when an uncaught exception triggers an abort. Read on
   // re-entry to break out instead of looping. Only accessed inside sessionSyncContext.
   private boolean isAborting = false;
@@ -310,6 +329,9 @@ public void start(OpenSessionRequest req, Metadata headers, Listener sessionList
           tracer.onStart();
 
           updateState(SessionState.STARTING);
+          // Bound the open handshake: if we're still STARTING when this fires, force-close so the
+          // reserved session-creation-budget slot is released instead of leaking forever.
+          scheduleOpenTimeoutCheck();
           openParams = OpenParams.create(headers, req);
 
           SessionRequest wrappedReq = SessionRequest.newBuilder().setOpenSession(req).build();
@@ -457,6 +479,41 @@ public void cancelRpc(long rpcId, @Nullable String message, @Nullable Throwable
         });
   }
 
+  private void scheduleOpenTimeoutCheck() {
+    openTimeout =
+        timer.newTimeout(
+            this::checkOpenTimeout,
+            sessionSyncContext,
+            OPEN_SESSION_TIMEOUT.toMillis(),
+            TimeUnit.MILLISECONDS);
+  }
+
+  private void cancelOpenTimeout() {
+    if (openTimeout != null) {
+      openTimeout.cancel();
+      openTimeout = null;
+    }
+  }
+
+  // Runs on sessionSyncContext (dispatched from the wheel-timer tick body). If the session is
+  // still STARTING, the open handshake never completed; force-close so the reserved
+  // session-creation-budget slot is released via the STARTING -> WAIT_SERVER_CLOSE ->
+  // onSessionClose
+  // path. Any transition out of STARTING cancels this timer, so reaching here in another state is a
+  // benign race (tick already dispatched) and is ignored.
+  private void checkOpenTimeout() {
+    sessionSyncContext.throwIfNotInThisSynchronizationContext();
+    if (state != SessionState.STARTING) {
+      return;
+    }
+    logger.warning(
+        String.format(
+            "Session %s did not complete the open handshake within %s, forcing session close",
+            info.getLogName(), OPEN_SESSION_TIMEOUT));
+    debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_open_timeout");
+    forceClose(OPEN_TIMEOUT_CLOSE_REQUEST);
+  }
+
   private void scheduleHeartbeatCheck() {
     heartbeatTimeout =
         timer.newTimeout(
@@ -849,6 +906,12 @@ private PeerInfo safeGetPeerInfo() {
   private void updateState(SessionState newState) {
     this.state = newState;
     this.lastStateChangedAt = clock.instant();
+    // The open deadline only applies while STARTING. Any other transition (READY on success, or a
+    // terminal state) means the handshake resolved, so drop the pending tick. NEW -> STARTING keeps
+    // it, since start() arms the timer right after this call.
+    if (newState != SessionState.STARTING) {
+      cancelOpenTimeout();
+    }
     // Once we're past READY, no further heartbeat checks are useful: checkHeartbeat short-circuits
     // on state.phase >= WAIT_SERVER_CLOSE. Cancel any pending tick to keep the wheel clean during
     // session churn.

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

@@ -77,6 +77,24 @@ SessionHandle newHandle(Session session) {
     return h;
   }
 
+  /**
+   * Unwinds a handle produced by {@link #newHandle} whose session never started, i.e. a synchronous
+   * failure in {@code SessionPoolImpl.createSession} between {@code newHandle} and {@code
+   * session.start}. Such a session stays in {@link SessionState#NEW}, so it never receives a
+   * terminal {@link SessionHandle#onSessionClosed} callback (which additionally rejects NEW
+   * sessions). This reverses exactly the bookkeeping {@code newHandle} performed so the handle
+   * doesn't strand {@link #allSessions} (blocking pool drain) or skew the pool stats.
+   */
+  void removeUnstartedHandle(SessionHandle handle) {
+    if (allSessions.remove(handle)) {
+      poolStats.startingCount--;
+      if (handle.inExpectedCount) {
+        poolStats.expectedCapacity--;
+        handle.inExpectedCount = false;
+      }
+    }
+  }
+
   /** Get {@link PoolStats} */
   PoolStats getStats() {
     return poolStats;

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

@@ -175,6 +175,18 @@ private enum PoolState {
   @GuardedBy("poolLock")
   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("poolLock")
+  private final Set<SessionHandle> sessionsHoldingBudget = new HashSet<>();
+
   private final ClientConfigurationManager configManager;
   private final ClientConfigurationManager.ListenerHandle configListenerHandle;
 
@@ -455,35 +467,65 @@ private void createSession(OpenParams openParams) {
     // Explicit create session streams in a detached context
     // We don't want to propagate the rpc deadline nor the trace context
     Context prevContext = Context.ROOT.attach();
+    // Tracks the handle once registered so the catch below can tell whether the reservation was
+    // ever bound to a handle (and therefore what cleanup is required).
+    SessionHandle handle = null;
     try {
       try (Scope ignored = io.opentelemetry.context.Context.root().makeCurrent()) {
 
+        // Build the metadata before registering the handle so the only step between newHandle and
+        // session.start (whose own failures are handled by the session's abort path once the
+        // listener is published) is the non-throwing sessionsHoldingBudget.add.
+        Metadata localMd = new Metadata();
+        localMd.merge(openParams.metadata());
+
         SessionStream stream = factory.createNew();
         Session session = new SessionImpl(metrics, info, sessionNum++, stream, timer);
-        SessionHandle handle = sessions.newHandle(session);
+        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());
+        final SessionHandle startedHandle = handle;
         session.start(
             openParams.request(),
             localMd,
             new Listener() {
               @Override
               public void onReady(OpenSessionResponse msg) {
-                SessionPoolImpl.this.onSessionReady(handle, msg);
+                SessionPoolImpl.this.onSessionReady(startedHandle, msg);
               }
 
               @Override
               public void onGoAway(GoAwayResponse msg) {
-                SessionPoolImpl.this.onSessionGoAway(handle, msg);
+                SessionPoolImpl.this.onSessionGoAway(startedHandle, msg);
               }
 
               @Override
               public void onClose(SessionState prevState, Status status, Metadata trailers) {
-                SessionPoolImpl.this.onSessionClose(handle, prevState, status, trailers);
+                SessionPoolImpl.this.onSessionClose(startedHandle, prevState, status, trailers);
               }
             });
       }
+    } 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
+      // the budget slot and any partially-registered handle would leak for the life of the pool.
+      // Release the slot as a failure and unwind the handle ourselves. Note session.start's own
+      // failures do NOT land here: it publishes the listener first and routes exceptions through
+      // the session's abort path, which fires onClose -> onSessionClose (the normal release path).
+      debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_create_failed");
+      logger.log(Level.WARNING, "Failed to create session, releasing reserved budget", e);
+      // If a handle was registered it owns the reservation via the set; otherwise the reservation
+      // was never bound to a handle and must be released directly.
+      if (handle == null || sessionsHoldingBudget.remove(handle)) {
+        budget.onSessionCreationFailure();
+      }
+      if (handle != null) {
+        sessions.removeUnstartedHandle(handle);
+      }
+      // Let the pool recover instead of running permanently short a session.
+      maybeScheduleCreateSessionRetry();
     } finally {
       Context.ROOT.detach(prevContext);
     }
@@ -540,7 +582,11 @@ private void onSessionReady(SessionHandle handle, OpenSessionResponse ignored) {
       }
       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();
@@ -613,6 +659,21 @@ private void onSessionClose(
 
     poolLock.lock();
     try {
+      // Release the budget reservation FIRST, before any code below that can throw. This is the
+      // session's terminal callback, so there is no later backstop: if the reservation is not
+      // released here it leaks forever. In particular handle.onSessionClosed(prevState) below can
+      // throw (e.g. IllegalStateException on an unexpected NEW / double-CLOSED prevState), which
+      // would otherwise strand the slot -- the same leak this fix exists to prevent.
+      //
+      // Release as a failure if this session still holds a slot, i.e. it terminated without ever
+      // reaching READY; sessions that reached READY were already removed from the set in
+      // onSessionReady, so this never double-releases. Keying off the reservation instead of the
+      // close-time prevState is also what covers the STARTING -> WAIT_SERVER_CLOSE (GO_AWAY) path
+      // that the abnormal-close branch below skips.
+      if (sessionsHoldingBudget.remove(handle)) {
+        budget.onSessionCreationFailure();
+      }
+
       logger.fine(
           String.format("Removing closed session from pool %s", handle.getSession().getLogName()));
 
@@ -642,9 +703,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/SessionImplTest.java

@@ -70,6 +70,9 @@
 import java.io.IOException;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
@@ -596,6 +599,12 @@ void testHeartbeatNotScheduledWithoutVRpc() throws Exception {
     assertThat(sessionListener.popUntil(OpenSessionResponse.class))
         .isInstanceOf(OpenSessionResponse.class);
 
+    // The open-handshake deadline is armed while STARTING and cancelled once READY (before onReady,
+    // so it's already settled here). Reset the counters to scope the assertion below to the
+    // heartbeat lifecycle only.
+    counting.scheduleCount.set(0);
+    counting.cancelCount.set(0);
+
     // After session is READY with no vRPC, no Timeout should ever have been scheduled. Wait a
     // bit so that any background tick (none expected) would have shown up.
     Thread.sleep(50);
@@ -640,6 +649,11 @@ void testHeartbeatScheduledOnlyDuringVRpc() throws Exception {
     assertThat(sessionListener.popUntil(OpenSessionResponse.class))
         .isInstanceOf(OpenSessionResponse.class);
 
+    // The open-handshake deadline is armed while STARTING and cancelled once READY. Reset the
+    // counters so the assertions below track only the heartbeat lifecycle.
+    counting.scheduleCount.set(0);
+    counting.cancelCount.set(0);
+
     assertThat(counting.scheduleCount.get()).isEqualTo(0);
 
     VRpc<SessionFakeScriptedRequest, SessionFakeScriptedResponse> rpc =
@@ -811,8 +825,99 @@ public void onClose(Session.SessionState prevState, Status status, Metadata trai
     assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED);
   }
 
+  // Regression test: a session that connects but never completes the open handshake must not stay
+  // wedged in STARTING forever. Without the open-handshake deadline the session would never fire a
+  // terminal callback, so the pool's reserved session-creation-budget slot would leak. Arming
+  // OPEN_SESSION_TIMEOUT force-closes such a session, routing STARTING -> WAIT_SERVER_CLOSE and
+  // driving the listener to a terminal Status (which releases the budget in the pool).
+  @Test
+  void sessionOpenTimeoutForcesClose() throws Exception {
+    // Capture the scheduled open-timeout tick so we can fire it deterministically instead of
+    // waiting the real OPEN_SESSION_TIMEOUT.
+    CapturingBigtableTimer capturing = new CapturingBigtableTimer(timer);
+    SessionImpl session =
+        new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), capturing);
+
+    FakeSessionListener sessionListener = new FakeSessionListener();
+    session.start(
+        OpenSessionRequest.newBuilder()
+            .setPayload(
+                OpenFakeSessionRequest.newBuilder().setHangBeforeOpen(true).build().toByteString())
+            .build(),
+        new Metadata(),
+        sessionListener);
+
+    // start() runs on the sessionSyncContext; wait until it has armed the open-timeout tick and the
+    // session is STARTING. The server never responds, so it stays STARTING.
+    Stopwatch sw = Stopwatch.createStarted();
+    while ((capturing.tasks.isEmpty() || session.getState() != Session.SessionState.STARTING)
+        && sw.elapsed(TimeUnit.SECONDS) < 5) {
+      Thread.sleep(10);
+    }
+    assertThat(session.getState()).isEqualTo(Session.SessionState.STARTING);
+    assertWithMessage("open-timeout tick should have been armed while STARTING")
+        .that(capturing.tasks)
+        .isNotEmpty();
+
+    // Fire the open-timeout tick on its executor (the sessionSyncContext), simulating the deadline
+    // elapsing while still STARTING.
+    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);
+  }
+
   // endregion
 
+  // Captures newTimeout tasks/executors so a scheduled tick can be fired deterministically instead
+  // of waiting out the real delay.
+  private static final class CapturingBigtableTimer implements BigtableTimer {
+    private final BigtableTimer delegate;
+    final List<Runnable> tasks = Collections.synchronizedList(new ArrayList<>());
+    final List<Executor> executors = Collections.synchronizedList(new ArrayList<>());
+
+    CapturingBigtableTimer(BigtableTimer delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) {
+      tasks.add(task);
+      executors.add(executor);
+      return new Timeout() {
+        @Override
+        public boolean cancel() {
+          return true;
+        }
+
+        @Override
+        public boolean isCancelled() {
+          return false;
+        }
+      };
+    }
+
+    @Override
+    public Registration onStop(Runnable hook) {
+      return delegate.onStop(hook);
+    }
+
+    @Override
+    public void stop() {
+      delegate.stop();
+    }
+  }
+
   // Wraps a real BigtableTimer and counts newTimeout / cancel calls. Used to assert that the
   // heartbeat tick is only armed while a vRPC is in flight.
   private static final class CountingBigtableTimer implements BigtableTimer {