From fec3892f519c499ee658907e0b905e158cc022a3 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Thu, 12 Nov 2026 03:22:48 +0000 Subject: [PATCH 1/4] fix(bigtable): stop SessionPoolImpl watchdog when sessions finish draining SessionPoolImpl.close() previously left the watchdog running until awaitTerminated() was called. Callers that close a pool without calling awaitTerminated() (such as TableBase.close() when an idle table is evicted from compat.Util.createSessionMap) leaked the self-rescheduling watchdog on the client's shared HashedWheelTimer, retaining the closed SessionPoolImpl in memory. Close the watchdog as soon as all sessions drain after close() (immediately if the pool is empty, or in onSessionClose when the last session terminates), and guard onSessionGoAway against creating replacement sessions after the pool is closed. --- .../v2/internal/session/SessionPoolImpl.java | 22 ++++++++++--- .../internal/session/SessionPoolImplTest.java | 33 +++++++++++++++++++ 2 files changed, 51 insertions(+), 4 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java index 502cdcacb637..3314d2b1fde5 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java @@ -134,7 +134,7 @@ private enum PoolState { @GuardedBy("poolLock") private final Deque> pendingRpcs = new ArrayDeque<>(); - private final Watchdog watchdog; + @VisibleForTesting final Watchdog watchdog; // Shared by every SessionImpl, PendingVRpc, Watchdog, AFE pruner, and retry-create-session in // this pool. One tick thread per Client (owned by Client); O(1) insert / O(1) cancel. @@ -354,12 +354,15 @@ public void close(CloseSessionRequest req) { retryCreateSessionFuture.cancel(); retryCreateSessionFuture = null; } - // Watchdog stays alive past close() so it can escalate any session that lingers in - // WAIT_SERVER_CLOSE during shutdown. awaitTerminated() takes ownership of closing it. + // Watchdog stays alive past close() while sessions drain so it can escalate any session + // that lingers in WAIT_SERVER_CLOSE during shutdown. Once all sessions have closed (or + // immediately if none exist), the watchdog is closed so callers that only invoke close() + // (such as TableBase.close() on cache eviction) don't leak the watchdog timer. sessions.close(req); // If the pool had no sessions, drainedFuture would never be completed by onSessionClose. if (sessions.getAllSessions().isEmpty()) { drainedFuture.complete(null); + watchdog.close(); } } finally { poolLock.unlock(); @@ -640,6 +643,9 @@ private void onSessionGoAway(SessionHandle handle, GoAwayResponse msg) { poolLock.lock(); try { handle.onSessionClosing(); + if (poolState == PoolState.CLOSED) { + return; + } // If the session received refresh config or pool requires a replacement, keep the current // session and request a replacement. @@ -688,9 +694,10 @@ private void onSessionClose( // 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) { - // Signal awaitTerminated() once the last session has drained. + // Signal awaitTerminated() and stop the watchdog once the last session has drained. if (sessions.getAllSessions().isEmpty()) { drainedFuture.complete(null); + watchdog.close(); } return; } @@ -1209,5 +1216,12 @@ public void close() { } } } + + @VisibleForTesting + boolean isClosed() { + synchronized (scheduleLock) { + return watchdogClosed; + } + } } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java index 207265693c93..566841728ad9 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java @@ -304,6 +304,7 @@ void awaitTerminatedReturnsTrueWhenPoolIsEmpty() throws InterruptedException { .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) .setDescription("empty pool") .build()); + assertThat(sessionPool.watchdog.isClosed()).isTrue(); assertThat(sessionPool.awaitTerminated(Duration.ofMillis(100))).isTrue(); } @@ -332,6 +333,38 @@ void awaitTerminatedReturnsTrueAfterSessionsDrain() assertThat(sessionPool.awaitTerminated(Duration.ofSeconds(10))).isTrue(); } + @Test + void closeStopsWatchdogAfterSessionsDrainWithoutAwaitTerminated() + throws InterruptedException, ExecutionException, TimeoutException { + // Callers such as TableBase.close() (invoked when an idle table is evicted from the + // per-table session cache in compat.Util.createSessionMap) only call sessionPool.close() + // and never call awaitTerminated(). Verify that the watchdog is still closed once all + // sessions finish draining so the closed pool is not retained by the shared timer. + sessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); + + VRpc vrpc = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + UnaryResponseFuture resultFuture = new UnaryResponseFuture<>(); + vrpc.start( + SessionFakeScriptedRequest.getDefaultInstance(), + VRpcCallContext.create(Deadline.after(10, TimeUnit.SECONDS), true, vrpcTracer), + resultFuture); + resultFuture.get(10, TimeUnit.SECONDS); + assertThat(sessionPool.watchdog.isClosed()).isFalse(); + + sessionPool.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("evicted table pool") + .build()); + + long deadlineMs = System.currentTimeMillis() + 10_000; + while (!sessionPool.watchdog.isClosed() && System.currentTimeMillis() < deadlineMs) { + Thread.sleep(10); + } + assertThat(sessionPool.watchdog.isClosed()).isTrue(); + } + @Test void closeCancelsInFlightVRpcAndDrains() throws InterruptedException { // Park the session-open response in the client interceptor so the vRPC cannot drain to a From c08f8bf13d246ec35dcbe74acfeae0cf905d00bc Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Thu, 12 Nov 2026 18:07:42 +0000 Subject: [PATCH 2/4] refactor(bigtable): replace pool Watchdog with per-session closeTimeout Remove SessionPoolImpl.Watchdog and instead arm a one-shot closeTimeout (CLOSE_SESSION_TIMEOUT = 5 minutes) on SessionImpl when entering WAIT_SERVER_CLOSE in startGracefulClose(), cancelled in updateState() when leaving WAIT_SERVER_CLOSE. This eliminates the recurring per-pool watchdog timer task on Client.sessionTimer and matches how openTimeout and heartbeatTimeout are managed per session. --- .../data/v2/internal/session/SessionImpl.java | 54 +++++ .../v2/internal/session/SessionPoolImpl.java | 182 +---------------- .../v2/internal/session/SessionImplTest.java | 99 +++++++++ .../internal/session/SessionPoolImplTest.java | 47 +---- .../v2/internal/session/WatchdogTest.java | 191 ------------------ 5 files changed, 165 insertions(+), 408 deletions(-) delete mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java index 39366e0e09f6..cabb86e1c76e 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java @@ -93,6 +93,11 @@ public class SessionImpl implements Session, VRpcSessionApi { .setDescription("session open timeout") .build(); + @VisibleForTesting + // Upper bound on how long a session may remain in WAIT_SERVER_CLOSE waiting for the server to + // close the stream after CloseSessionRequest is sent. + static final Duration CLOSE_SESSION_TIMEOUT = Duration.ofMinutes(5); + private final Clock clock; private final BigtableTimer timer; // Serializes all session state mutations. Stream callbacks and the heartbeat tick dispatch @@ -149,6 +154,10 @@ public class SessionImpl implements Session, VRpcSessionApi { // session leaves STARTING so a session that opens normally never trips it. @Nullable private BigtableTimer.Timeout openTimeout; + // Handle for the graceful-close deadline (armed while WAIT_SERVER_CLOSE). Cancelled the moment + // the session leaves WAIT_SERVER_CLOSE. + @Nullable private BigtableTimer.Timeout closeTimeout; + // 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; @@ -419,6 +428,7 @@ private void startGracefulClose() { // TODO: send metrics updateState(SessionState.WAIT_SERVER_CLOSE); + scheduleCloseTimeoutCheck(); // Should never happen if (closeReason == null) { @@ -532,6 +542,47 @@ private void checkOpenTimeout() { forceClose(OPEN_TIMEOUT_CLOSE_REQUEST); } + private void scheduleCloseTimeoutCheck() { + closeTimeout = + timer.newTimeout( + this::checkCloseTimeout, + sessionSyncContext, + CLOSE_SESSION_TIMEOUT.toMillis(), + TimeUnit.MILLISECONDS); + } + + private void cancelCloseTimeout() { + if (closeTimeout != null) { + closeTimeout.cancel(); + closeTimeout = null; + } + } + + // Runs on sessionSyncContext (dispatched from the wheel-timer tick body). If the session is + // still in WAIT_SERVER_CLOSE, the server never closed the stream after CloseSessionRequest; + // force-close so the session reaches CLOSED. + private void checkCloseTimeout() { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); + if (state != SessionState.WAIT_SERVER_CLOSE) { + return; + } + debugTagTracer.record(TelemetryConfiguration.Level.WARN, "watchdog_close_session"); + logger.fine( + String.format( + "Found session %s that lingered too long in WAIT_SERVER_CLOSE state, canceling" + + " session", + info.getLogName())); + forceClose( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR) + .setDescription( + String.format( + "Watchdog found session %s in awaiting close for too long. Last" + + " activity was %s ago.", + info.getLogName(), Duration.between(lastStateChangedAt, clock.instant()))) + .build()); + } + private void scheduleHeartbeatCheck() { heartbeatTimeout = timer.newTimeout( @@ -935,6 +986,9 @@ private void updateState(SessionState newState) { if (newState.phase >= SessionState.WAIT_SERVER_CLOSE.phase) { cancelHeartbeatTimeout(); } + if (newState != SessionState.WAIT_SERVER_CLOSE) { + cancelCloseTimeout(); + } } private static String formatPeerInfo(PeerInfo peerInfo) { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java index 3314d2b1fde5..bcab2476a3a5 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java @@ -18,7 +18,6 @@ import com.google.bigtable.v2.ClientConfiguration; import com.google.bigtable.v2.CloseSessionRequest; -import com.google.bigtable.v2.CloseSessionRequest.CloseSessionReason; import com.google.bigtable.v2.FeatureFlags; import com.google.bigtable.v2.GoAwayResponse; import com.google.bigtable.v2.LoadBalancingOptions; @@ -53,7 +52,6 @@ import io.grpc.Status; import io.grpc.protobuf.StatusProto; import io.opentelemetry.context.Scope; -import java.time.Clock; import java.time.Duration; import java.time.Instant; import java.util.ArrayDeque; @@ -69,7 +67,6 @@ import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; -import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; import java.util.logging.Level; import java.util.logging.Logger; @@ -134,14 +131,12 @@ private enum PoolState { @GuardedBy("poolLock") private final Deque> pendingRpcs = new ArrayDeque<>(); - @VisibleForTesting final Watchdog watchdog; - - // Shared by every SessionImpl, PendingVRpc, Watchdog, AFE pruner, and retry-create-session in + // Shared by every SessionImpl, PendingVRpc, AFE pruner, and retry-create-session in // this pool. One tick thread per Client (owned by Client); O(1) insert / O(1) cancel. private final BigtableTimer timer; // Executor for timer-scheduled pool maintenance bodies (AFE prune, retry-create-session, - // deadline monitor, watchdog tick). The timer's tick thread hands work here so it doesn't + // deadline monitor). The timer's tick thread hands work here so it doesn't // run pool-lock-holding work inline. private final Executor backgroundExecutor; @@ -267,11 +262,8 @@ public SessionPoolImpl( debugTagTracer = metrics.getDebugTagTracer(); - // Watchdog checks for sessions in WAIT_SERVER_CLOSE state and runs every 5 minutes - watchdog = - new Watchdog( - poolLock, timer, backgroundExecutor, Duration.ofMinutes(5), sessions, debugTagTracer); - // Heartbeat monitoring is now done per-session via SessionImpl.scheduleHeartbeatCheck. + // Heartbeat, open-handshake, and graceful-close (WAIT_SERVER_CLOSE) timeouts are managed + // per-session in SessionImpl. scheduleNextAfePrune(); this.budget = budget; @@ -354,15 +346,10 @@ public void close(CloseSessionRequest req) { retryCreateSessionFuture.cancel(); retryCreateSessionFuture = null; } - // Watchdog stays alive past close() while sessions drain so it can escalate any session - // that lingers in WAIT_SERVER_CLOSE during shutdown. Once all sessions have closed (or - // immediately if none exist), the watchdog is closed so callers that only invoke close() - // (such as TableBase.close() on cache eviction) don't leak the watchdog timer. sessions.close(req); // If the pool had no sessions, drainedFuture would never be completed by onSessionClose. if (sessions.getAllSessions().isEmpty()) { drainedFuture.complete(null); - watchdog.close(); } } finally { poolLock.unlock(); @@ -392,9 +379,6 @@ public boolean awaitTerminated(Duration timeout) throws InterruptedException { debugTagTracer.record( TelemetryConfiguration.Level.ERROR, "session_pool_drained_future_failed"); throw new IllegalStateException("drainedFuture failed unexpectedly", e); - } finally { - // Close the watchdog on the way out — drained or timed out, its job is done. - watchdog.close(); } } @@ -427,8 +411,6 @@ public void start(OpenReqT openReq, Metadata md) { break; } } - - watchdog.start(); } finally { poolLock.unlock(); } @@ -694,10 +676,9 @@ private void onSessionClose( // 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) { - // Signal awaitTerminated() and stop the watchdog once the last session has drained. + // Signal awaitTerminated() once the last session has drained. if (sessions.getAllSessions().isEmpty()) { drainedFuture.complete(null); - watchdog.close(); } return; } @@ -1071,157 +1052,4 @@ private BigtableTimer.Timeout monitorDeadline(Deadline deadline) { TimeUnit.NANOSECONDS); } } - - static class Watchdog implements Runnable { - private static final Logger LOG = Logger.getLogger(Watchdog.class.getName()); - - private final Lock lock; - private final BigtableTimer timer; - private final Executor backgroundExecutor; - private final Duration interval; - private final SessionList sessions; - private final Clock clock; - private final DebugTagTracer debugTagTracer; - - // Guards currentTimeout and watchdogClosed. Self-contained — never composed with any external - // lock. - private final Object scheduleLock = new Object(); - - @GuardedBy("scheduleLock") - private BigtableTimer.Timeout currentTimeout; - - @GuardedBy("scheduleLock") - private boolean watchdogClosed = false; - - // The `lock` parameter is the pool-wide lock (SessionPoolImpl.poolLock). - public Watchdog( - Lock lock, - BigtableTimer timer, - Executor backgroundExecutor, - Duration interval, - SessionList sessionList, - DebugTagTracer debugTagTracer) { - this( - lock, - timer, - backgroundExecutor, - interval, - sessionList, - debugTagTracer, - Clock.systemUTC()); - } - - @VisibleForTesting - Watchdog( - Lock lock, - BigtableTimer timer, - Executor backgroundExecutor, - Duration interval, - SessionList sessionList, - DebugTagTracer debugTagTracer, - Clock clock) { - this.lock = lock; - this.timer = timer; - this.backgroundExecutor = backgroundExecutor; - this.interval = interval; - this.sessions = sessionList; - this.debugTagTracer = debugTagTracer; - this.clock = clock; - } - - public void start() { - scheduleNext(); - } - - // Self-reschedule. Called once from start() and again at the end of each tick. - private void scheduleNext() { - synchronized (scheduleLock) { - if (watchdogClosed) return; - currentTimeout = - timer.newTimeout( - this::runAndReschedule, - backgroundExecutor, - interval.toMillis(), - TimeUnit.MILLISECONDS); - } - } - - private void runAndReschedule() { - try { - run(); - } catch (Throwable t) { - // Preserve the watchdog across body exceptions — unlike scheduleAtFixedRate, which silently - // stops the schedule on the first exception, we keep going so a transient fault doesn't - // permanently disable session leak detection. - LOG.log(Level.WARNING, "Watchdog tick threw; continuing", t); - } finally { - scheduleNext(); - } - } - - @Override - public void run() { - // Snapshot under the pool lock: getAllSessions() returns the live HashSet, and pool state - // mutation (session create/close on another thread) during iteration would throw - // ConcurrentModificationException. Most common trigger is a heartbeat-miss cascade that - // churns sessions while the watchdog is walking the list. - Set allSessions; - lock.lock(); - try { - allSessions = new HashSet<>(sessions.getAllSessions()); - } finally { - lock.unlock(); - } - - for (SessionHandle handle : allSessions) { - Session s = handle.getSession(); - - if (s.getState() != SessionState.WAIT_SERVER_CLOSE) { - continue; - } - - Duration stateDuration = Duration.between(s.getLastStateChange(), Instant.now(clock)); - if (stateDuration.compareTo(interval) < 0) { - continue; - } - - debugTagTracer.record(TelemetryConfiguration.Level.WARN, "watchdog_close_session"); - LOG.fine( - String.format( - "Found session %s that lingered too long in WAIT_SERVER_CLOSE state, canceling" - + " session", - handle.getSession().getLogName())); - handle - .getSession() - .forceClose( - CloseSessionRequest.newBuilder() - .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR) - .setDescription( - String.format( - "Watchdog found session %s in awaiting close for too long. Last" - + " activity was %s ago.", - handle.getSession().getLogName(), - Duration.between( - handle.getSession().getLastStateChange(), Instant.now(clock)))) - .build()); - } - } - - public void close() { - synchronized (scheduleLock) { - watchdogClosed = true; - if (currentTimeout != null) { - currentTimeout.cancel(); - currentTimeout = null; - } - } - } - - @VisibleForTesting - boolean isClosed() { - synchronized (scheduleLock) { - return watchdogClosed; - } - } - } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java index d995a32f6d10..2f2128d2f38b 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java @@ -36,11 +36,14 @@ import com.google.bigtable.v2.SessionFakeScriptedRequest; import com.google.bigtable.v2.SessionFakeScriptedResponse; import com.google.bigtable.v2.SessionParametersResponse; +import com.google.bigtable.v2.SessionRequest; import com.google.bigtable.v2.VirtualRpcResponse; import com.google.cloud.bigtable.data.v2.internal.api.InstanceName; import com.google.cloud.bigtable.data.v2.internal.api.UnaryResponseFuture; import com.google.cloud.bigtable.data.v2.internal.api.VRpcException; import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool; +import com.google.cloud.bigtable.data.v2.internal.channels.ForwardingSessionStream; +import com.google.cloud.bigtable.data.v2.internal.channels.SessionStream; import com.google.cloud.bigtable.data.v2.internal.channels.SingleChannelPool; import com.google.cloud.bigtable.data.v2.internal.csm.Metrics; import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; @@ -893,6 +896,102 @@ void sessionOpenTimeoutForcesClose() throws Exception { assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED); } + @Test + void sessionCloseTimeoutForcesClose() throws Exception { + CapturingBigtableTimer capturing = new CapturingBigtableTimer(timer); + SessionStream underlying = sessionFactory.createNew(); + // Drop CloseSession and halfClose so the server never closes the stream and the session stays + // in WAIT_SERVER_CLOSE until the per-session closeTimeout fires. + SessionStream hangingCloseStream = + new ForwardingSessionStream(underlying) { + @Override + public void sendMessage(SessionRequest request) { + if (request.hasCloseSession()) { + return; + } + super.sendMessage(request); + } + + @Override + public void halfClose() {} + }; + SessionImpl session = new SessionImpl(metrics, poolInfo, 0, hangingCloseStream, capturing); + + FakeSessionListener sessionListener = new FakeSessionListener(); + session.start( + OpenSessionRequest.newBuilder() + .setPayload(OpenFakeSessionRequest.getDefaultInstance().toByteString()) + .build(), + new Metadata(), + sessionListener); + assertThat(sessionListener.popUntil(OpenSessionResponse.class)) + .isInstanceOf(OpenSessionResponse.class); + + capturing.tasks.clear(); + capturing.executors.clear(); + + session.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("test close timeout") + .build()); + + Stopwatch sw = Stopwatch.createStarted(); + while ((capturing.tasks.isEmpty() + || session.getState() != Session.SessionState.WAIT_SERVER_CLOSE) + && sw.elapsed(TimeUnit.SECONDS) < 5) { + Thread.sleep(10); + } + assertThat(session.getState()).isEqualTo(Session.SessionState.WAIT_SERVER_CLOSE); + assertWithMessage("close-timeout tick should have been armed while WAIT_SERVER_CLOSE") + .that(capturing.tasks) + .hasSize(1); + + // Fire the close-timeout tick on its executor (sessionSyncContext). + capturing.executors.get(0).execute(capturing.tasks.get(0)); + + Status status = sessionListener.popUntil(Status.class); + assertThat(status.getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(status.getDescription()).contains("awaiting close for too long"); + assertThat(sessionListener.getLastPrevState()) + .isEqualTo(Session.SessionState.WAIT_SERVER_CLOSE); + assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED); + } + + @Test + void sessionNormalCloseCancelsCloseTimeout() throws Exception { + CountingBigtableTimer counting = new CountingBigtableTimer(timer); + SessionImpl session = + new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), counting); + + FakeSessionListener sessionListener = new FakeSessionListener(); + session.start( + OpenSessionRequest.newBuilder() + .setPayload(OpenFakeSessionRequest.getDefaultInstance().toByteString()) + .build(), + new Metadata(), + sessionListener); + assertThat(sessionListener.popUntil(OpenSessionResponse.class)) + .isInstanceOf(OpenSessionResponse.class); + + counting.scheduleCount.set(0); + counting.cancelCount.set(0); + + session.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("normal close") + .build()); + assertThat(sessionListener.popUntil(Status.class)).isOk(); + + assertWithMessage("startGracefulClose must arm the close-timeout tick") + .that(counting.scheduleCount.get()) + .isEqualTo(1); + assertWithMessage("reaching CLOSED must cancel the close-timeout tick") + .that(counting.cancelCount.get()) + .isEqualTo(1); + } + // endregion // Captures newTimeout tasks/executors so a scheduled tick can be fired deterministically instead diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java index 566841728ad9..5f9612842d62 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java @@ -304,7 +304,6 @@ void awaitTerminatedReturnsTrueWhenPoolIsEmpty() throws InterruptedException { .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) .setDescription("empty pool") .build()); - assertThat(sessionPool.watchdog.isClosed()).isTrue(); assertThat(sessionPool.awaitTerminated(Duration.ofMillis(100))).isTrue(); } @@ -333,38 +332,6 @@ void awaitTerminatedReturnsTrueAfterSessionsDrain() assertThat(sessionPool.awaitTerminated(Duration.ofSeconds(10))).isTrue(); } - @Test - void closeStopsWatchdogAfterSessionsDrainWithoutAwaitTerminated() - throws InterruptedException, ExecutionException, TimeoutException { - // Callers such as TableBase.close() (invoked when an idle table is evicted from the - // per-table session cache in compat.Util.createSessionMap) only call sessionPool.close() - // and never call awaitTerminated(). Verify that the watchdog is still closed once all - // sessions finish draining so the closed pool is not retained by the shared timer. - sessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); - - VRpc vrpc = - sessionPool.newCall(FakeDescriptor.SCRIPTED); - UnaryResponseFuture resultFuture = new UnaryResponseFuture<>(); - vrpc.start( - SessionFakeScriptedRequest.getDefaultInstance(), - VRpcCallContext.create(Deadline.after(10, TimeUnit.SECONDS), true, vrpcTracer), - resultFuture); - resultFuture.get(10, TimeUnit.SECONDS); - assertThat(sessionPool.watchdog.isClosed()).isFalse(); - - sessionPool.close( - CloseSessionRequest.newBuilder() - .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) - .setDescription("evicted table pool") - .build()); - - long deadlineMs = System.currentTimeMillis() + 10_000; - while (!sessionPool.watchdog.isClosed() && System.currentTimeMillis() < deadlineMs) { - Thread.sleep(10); - } - assertThat(sessionPool.watchdog.isClosed()).isTrue(); - } - @Test void closeCancelsInFlightVRpcAndDrains() throws InterruptedException { // Park the session-open response in the client interceptor so the vRPC cannot drain to a @@ -865,16 +832,16 @@ void tearDown() { @Test public void test() throws Exception { ArgumentCaptor runnableCaptor = ArgumentCaptor.forClass(Runnable.class); - // Filter out watchdog (5 min, exact), AFE prune (10 min, exact), and per-session open-timeout - // (OPEN_SESSION_TIMEOUT, exact) ticks. The retry-create-session site computes its delay - // against the real wall clock and the fake budget clock, so it can land anywhere from - // sub-second to a couple of penalty intervals. Match anything that isn't one of the fixed - // cadences. - long watchdogMs = Duration.ofMinutes(5).toMillis(); + // Filter out close-timeout (5 min, exact), AFE prune (10 min, exact), and per-session + // open-timeout (OPEN_SESSION_TIMEOUT, exact) ticks. The retry-create-session site computes + // its delay against the real wall clock and the fake budget clock, so it can land anywhere + // from sub-second to a couple of penalty intervals. Match anything that isn't one of the + // fixed cadences. + long closeTimeoutMs = SessionImpl.CLOSE_SESSION_TIMEOUT.toMillis(); long afePruneMs = SessionList.SESSION_LIST_PRUNE_INTERVAL.toMillis(); long openTimeoutMs = SessionImpl.OPEN_SESSION_TIMEOUT.toMillis(); LongPredicate isRetrySchedule = - d -> d > 0 && d != watchdogMs && d != afePruneMs && d != openTimeoutMs; + d -> d > 0 && d != closeTimeoutMs && d != afePruneMs && d != openTimeoutMs; // start the pool sessionPool.start( diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java deleted file mode 100644 index 121fae54df22..000000000000 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java +++ /dev/null @@ -1,191 +0,0 @@ -/* - * Copyright 2026 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.google.cloud.bigtable.data.v2.internal.session; - -import static com.google.common.truth.Truth.assertThat; - -import com.google.bigtable.v2.CloseSessionRequest; -import com.google.bigtable.v2.OpenSessionRequest; -import com.google.bigtable.v2.PeerInfo; -import com.google.bigtable.v2.PeerInfo.TransportType; -import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; -import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc; -import com.google.cloud.bigtable.data.v2.internal.session.Session.SessionState; -import com.google.cloud.bigtable.data.v2.internal.session.SessionPoolImpl.Watchdog; -import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeClock; -import com.google.common.util.concurrent.MoreExecutors; -import com.google.protobuf.Message; -import io.grpc.Metadata; -import java.time.Duration; -import java.time.Instant; -import java.util.concurrent.Executor; -import java.util.concurrent.TimeUnit; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.junit.jupiter.MockitoExtension; - -@ExtendWith(MockitoExtension.class) -public class WatchdogTest { - private final Duration interval = Duration.ofMinutes(5); - private BigtableTimer timer; - private SessionPoolImpl.Watchdog watchdog; - private SessionList sessions; - private FakeSession fakeSession = new FakeSession(); - - private Instant now; - private FakeClock fakeClock; - - @BeforeEach - void setUp() { - // run() is invoked synchronously in tests; the timer is wired in only so the constructor - // signature is satisfied. start() / close() are not exercised here. - timer = new NoOpBigtableTimer(); - - now = Instant.now(); - fakeClock = new FakeClock(now); - sessions = new SessionList(); - watchdog = - new Watchdog( - new java.util.concurrent.locks.ReentrantLock(), - timer, - MoreExecutors.directExecutor(), - interval, - sessions, - NoopMetrics.NoopDebugTracer.INSTANCE, - fakeClock); - } - - // A BigtableTimer that drops every newTimeout(). Used because awaitCloseTest drives the watchdog - // by calling run() directly; the scheduling layer is not under test. - private static final class NoOpBigtableTimer implements BigtableTimer { - @Override - public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) { - return new Timeout() { - @Override - public boolean cancel() { - return false; - } - - @Override - public boolean isCancelled() { - return true; - } - }; - } - - @Override - public Registration onStop(Runnable hook) { - return () -> {}; - } - - @Override - public void stop() {} - } - - @Test - public void awaitCloseTest() { - fakeSession.lastStateChange = fakeClock.instant(); - - fakeSession.state = SessionState.STARTING; - SessionList.SessionHandle handle = sessions.newHandle(fakeSession); - - fakeSession.state = SessionState.READY; - handle.onSessionStarted(); - - fakeSession.state = SessionState.WAIT_SERVER_CLOSE; - handle.onSessionClosing(); - - watchdog.run(); - assertThat(fakeSession.forceCloseReason).isNull(); - - fakeClock.increment(interval.plus(Duration.ofNanos(1))); - watchdog.run(); - - assertThat(fakeSession.forceCloseReason).isNotNull(); - assertThat(fakeSession.forceCloseReason.getDescription()) - .contains("awaiting close for too long"); - } - - private static class FakeSession implements Session { - private static final String sessionName = "fake-session"; - private SessionState state = SessionState.NEW; - private Instant lastStateChange = Instant.now(); - private CloseSessionRequest forceCloseReason = null; - - @Override - public PeerInfo getPeerInfo() { - return PeerInfo.newBuilder() - .setTransportType(TransportType.TRANSPORT_TYPE_SESSION_DIRECT_ACCESS) - .setApplicationFrontendId(1234) - .setApplicationFrontendRegion("us-east1") - .setApplicationFrontendSubzone("xx") - .build(); - } - - @Override - public SessionState getState() { - return state; - } - - @Override - public Instant getLastStateChange() { - return lastStateChange; - } - - @Override - public String getLogName() { - return sessionName; - } - - @Override - public Instant getNextHeartbeat() { - return Instant.now().plus(Duration.ofMinutes(5)); - } - - @Override - public OpenParams getOpenParams() { - throw new UnsupportedOperationException(); - } - - @Override - public boolean isOpenParamsUpdated() { - return false; - } - - @Override - public void start(OpenSessionRequest req, Metadata headers, Listener sessionListener) { - throw new UnsupportedOperationException(); - } - - @Override - public void close(CloseSessionRequest req) { - throw new UnsupportedOperationException(); - } - - @Override - public void forceClose(CloseSessionRequest reason) { - forceCloseReason = reason; - } - - @Override - public - VRpc newCall(VRpcDescriptor descriptor) - throws IllegalStateException { - throw new UnsupportedOperationException(); - } - } -} From bdd7996bb5b2fcaa0e0e29e425c1a6895fc15288 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Thu, 12 Nov 2026 19:35:44 +0000 Subject: [PATCH 3/4] chore(bigtable): address PR review comments on session closeTimeout Remove unused Session.getLastStateChange(), add a unit test for onSessionGoAway on a closed pool using monotonic System.nanoTime(), pass FakeClock in sessionCloseTimeoutForcesClose, and update stale watchdog comments in Client and SessionPoolImplTest. --- .../bigtable/data/v2/internal/api/Client.java | 23 ++++---- .../data/v2/internal/session/Session.java | 2 - .../data/v2/internal/session/SessionImpl.java | 9 +--- .../v2/internal/session/SessionImplTest.java | 7 ++- .../v2/internal/session/SessionListTest.java | 5 -- .../internal/session/SessionPoolImplTest.java | 53 +++++++++++++++++-- 6 files changed, 68 insertions(+), 31 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java index 4893d17eddf1..5743a4cc9e88 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java @@ -61,9 +61,9 @@ public class Client implements AutoCloseable { private static final Logger logger = Logger.getLogger(Client.class.getName()); - // Per-pool drain budget during close. One full watchdog tick (5 min) plus 1 min buffer; if a - // pool can't drain in that window, something is genuinely wrong on the server side and we give - // up on it so close() returns. The watchdog interval is what makes the worst case finite. + // Per-pool drain budget during close. One full session close timeout (5 min) plus 1 min buffer; + // if a pool can't drain in that window, something is genuinely wrong on the server side and we + // give up on it so close() returns. private static final Duration POOL_DRAIN_TIMEOUT = Duration.ofMinutes(6); public static final String USER_AGENT = "bigtable-java/" + Version.VERSION; @@ -100,7 +100,7 @@ public class Client implements AutoCloseable { // blocked user callback from starving heartbeats, retry delays, or other vRPCs (which all run on // backgroundExecutor). private final Resource userCallbackExecutor; - // Hashed-wheel timer for heartbeat / deadline / watchdog / retry scheduling. Built over + // Hashed-wheel timer for heartbeat / deadline / close-timeout / retry scheduling. Built over // backgroundExecutor (the timer's tick thread dispatches bodies onto it). Single tick thread per // Client, shared across every SessionPoolImpl. private final BigtableTimer sessionTimer; @@ -290,14 +290,15 @@ public void close() { // CLOSING → graceful CloseSessionRequest → WAIT_SERVER_CLOSE → CLOSED asynchronously. toClose.forEach(p -> p.close(closeReq)); - // Phase 2: wait for sessions to drain. The pool's watchdog stays alive during this wait and - // escalates anything stuck in WAIT_SERVER_CLOSE longer than its tick interval (5 min). Once - // a pool's last session reaches CLOSED, drainedFuture completes and awaitTerminated returns. + // Phase 2: wait for sessions to drain. Each closing session arms a per-session closeTimeout + // (5 min) when entering WAIT_SERVER_CLOSE to force-close if the server never closes the stream. + // Once a pool's last session reaches CLOSED, drainedFuture completes and awaitTerminated + // returns. // - // Shared deadline across all pools: every pool's watchdog started at Phase 1 and ticks in - // parallel on the shared sessionTimer, so one watchdog interval is enough for the whole - // fleet to escalate. Bounding on a single deadline keeps close() finite even if multiple - // pools' force-close paths hang. + // Shared deadline across all pools: every closing session's closeTimeout started at Phase 1 + // on the shared sessionTimer, so one closeTimeout interval is enough for the whole fleet to + // escalate. Bounding on a single deadline keeps close() finite even if multiple pools' + // force-close paths hang. long deadlineNanos = System.nanoTime() + POOL_DRAIN_TIMEOUT.toNanos(); for (SessionPool pool : toClose) { long remainingNanos = deadlineNanos - System.nanoTime(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Session.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Session.java index 35f8f98bc571..e025ce71518b 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Session.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Session.java @@ -76,8 +76,6 @@ OpenParams withIncrementedAttempts() { SessionState getState(); - Instant getLastStateChange(); - OpenParams getOpenParams(); boolean isOpenParamsUpdated(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java index cabb86e1c76e..e792919cbe86 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java @@ -113,8 +113,6 @@ public class SessionImpl implements Session, VRpcSessionApi { private volatile SessionState state = SessionState.NEW; - private volatile Instant lastStateChangedAt; - // All fields below are owned by sessionSyncContext: every writer runs inside a // sessionSyncContext task, and the in-class readers do too (handlers, scheduled heartbeat // tick). They lost their volatile / lock guards when synchronized(lock) was removed; @@ -124,6 +122,8 @@ public class SessionImpl implements Session, VRpcSessionApi { // either run on sessionSyncContext themselves (e.g. via a Session.Listener callback, which is // dispatched from inside the context) or accept a possibly-stale snapshot. There are no // off-context production readers today. + private Instant lastStateChangedAt; + private Listener sessionListener; private OpenParams openParams; @@ -264,11 +264,6 @@ public SessionState getState() { return state; } - @Override - public Instant getLastStateChange() { - return lastStateChangedAt; - } - @Override public OpenParams getOpenParams() { return openParams; diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java index 2f2128d2f38b..56297c7e49ee 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java @@ -898,6 +898,7 @@ void sessionOpenTimeoutForcesClose() throws Exception { @Test void sessionCloseTimeoutForcesClose() throws Exception { + FakeClock clock = new FakeClock(Instant.now()); CapturingBigtableTimer capturing = new CapturingBigtableTimer(timer); SessionStream underlying = sessionFactory.createNew(); // Drop CloseSession and halfClose so the server never closes the stream and the session stays @@ -915,7 +916,8 @@ public void sendMessage(SessionRequest request) { @Override public void halfClose() {} }; - SessionImpl session = new SessionImpl(metrics, poolInfo, 0, hangingCloseStream, capturing); + SessionImpl session = + new SessionImpl(metrics, clock, poolInfo, 0, hangingCloseStream, capturing); FakeSessionListener sessionListener = new FakeSessionListener(); session.start( @@ -947,12 +949,15 @@ public void halfClose() {} .that(capturing.tasks) .hasSize(1); + clock.increment(SessionImpl.CLOSE_SESSION_TIMEOUT); + // Fire the close-timeout tick on its executor (sessionSyncContext). capturing.executors.get(0).execute(capturing.tasks.get(0)); Status status = sessionListener.popUntil(Status.class); assertThat(status.getCode()).isEqualTo(Status.Code.UNAVAILABLE); assertThat(status.getDescription()).contains("awaiting close for too long"); + assertThat(status.getDescription()).contains("PT5M ago"); assertThat(sessionListener.getLastPrevState()) .isEqualTo(Session.SessionState.WAIT_SERVER_CLOSE); assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionListTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionListTest.java index e44d795f13e5..469401766b12 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionListTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionListTest.java @@ -508,11 +508,6 @@ public PeerInfo getPeerInfo() { .build(); } - @Override - public Instant getLastStateChange() { - throw new UnsupportedOperationException(); - } - @Override public OpenParams getOpenParams() { throw new UnsupportedOperationException(); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java index 5f9612842d62..abdf8a4473d3 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java @@ -187,8 +187,7 @@ void tearDown() throws InterruptedException { .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) .setDescription("close session") .build()); - // Wait for sessions to drain so the watchdog can be closed before testTimer.stop() races - // with its self-reschedule loop. + // Wait for sessions to drain before stopping the shared timer. sessionPool.awaitTerminated(Duration.ofSeconds(10)); channelPool.close(); // channel gets shutdown in channelPool.close() @@ -332,6 +331,50 @@ void awaitTerminatedReturnsTrueAfterSessionsDrain() assertThat(sessionPool.awaitTerminated(Duration.ofSeconds(10))).isTrue(); } + @Test + @SuppressWarnings("GuardedBy") + void goAwayOnClosedPoolDoesNotCreateReplacementSession() throws Exception { + sessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); + + long deadlineNano = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + ReentrantLock poolLock = extractPoolLock(sessionPool); + while (System.nanoTime() < deadlineNano) { + poolLock.lock(); + try { + if (sessionPool.sessions.getStats().getReadyCount() == 5) { + break; + } + } finally { + poolLock.unlock(); + } + Thread.sleep(20); + } + + SessionHandle handle; + poolLock.lock(); + try { + handle = sessionPool.sessions.getAllSessions().iterator().next(); + } finally { + poolLock.unlock(); + } + + int openCountBeforeClose = fakeService.getOpenRequestCount().get(); + sessionPool.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("closed before goaway") + .build()); + + java.lang.reflect.Method onSessionGoAway = + SessionPoolImpl.class.getDeclaredMethod( + "onSessionGoAway", SessionHandle.class, com.google.bigtable.v2.GoAwayResponse.class); + onSessionGoAway.setAccessible(true); + onSessionGoAway.invoke( + sessionPool, handle, com.google.bigtable.v2.GoAwayResponse.getDefaultInstance()); + + assertThat(fakeService.getOpenRequestCount().get()).isEqualTo(openCountBeforeClose); + } + @Test void closeCancelsInFlightVRpcAndDrains() throws InterruptedException { // Park the session-open response in the client interceptor so the vRPC cannot drain to a @@ -339,7 +382,7 @@ void closeCancelsInFlightVRpcAndDrains() throws InterruptedException { // vRPC's listener gets a CANCELLED onClose with the "SessionPool closed:" description from // SessionPoolImpl.close()'s cancellation loop, (b) awaitTerminated returns true (the // OPENING session is driven to CLOSED by sessions.close(req), not parked until the - // watchdog tick), and (c) no second onClose arrives once the delayed responses finally + // closeTimeout tick), and (c) no second onClose arrives once the delayed responses finally // flush past the interceptor. SessionPool testSessionPool = null; try (ChannelPool delayedPool = @@ -912,7 +955,7 @@ class GoAwayBeforeReadyBudgetRelease { @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 + // Mock timer so the pool's 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); @@ -1018,7 +1061,7 @@ class CreateSessionFailureBudgetRelease { @BeforeEach void setUp() { fakeClock = new FakeClock(Instant.now()); - // Mock timer so the pool's watchdog / AFE-prune / create-session-retry ticks never fire on + // Mock timer so the pool's AFE-prune / create-session-retry ticks never fire on // their own, leaving the test thread as the only actor touching the budget. mockTimer = mock(BigtableTimer.class, Mockito.RETURNS_DEEP_STUBS); From 5bba152a2816d0627139bb15f5b997761afa3f31 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Thu, 12 Nov 2026 19:49:43 +0000 Subject: [PATCH 4/4] chore(bigtable): check newState.phase > WAIT_SERVER_CLOSE.phase in updateState --- .../cloud/bigtable/data/v2/internal/session/SessionImpl.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java index e792919cbe86..b155ffd4d7ad 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java @@ -981,7 +981,7 @@ private void updateState(SessionState newState) { if (newState.phase >= SessionState.WAIT_SERVER_CLOSE.phase) { cancelHeartbeatTimeout(); } - if (newState != SessionState.WAIT_SERVER_CLOSE) { + if (newState.phase > SessionState.WAIT_SERVER_CLOSE.phase) { cancelCloseTimeout(); } }