From baab434c7269d1e0f0e45fdd7a7aea11f93d5231 Mon Sep 17 00:00:00 2001 From: Yuri Golobokov Date: Fri, 18 Sep 2026 03:43:54 +0000 Subject: [PATCH 1/2] fix(bigtable): fix abnormal session closures and scale-up in session pool --- .../data/v2/internal/session/SessionImpl.java | 40 +++- .../data/v2/internal/session/SessionList.java | 44 ++-- .../v2/internal/session/SessionPoolImpl.java | 28 ++- .../v2/internal/session/SessionImplTest.java | 1 + .../v2/internal/session/SessionListTest.java | 48 ++++ .../internal/session/SessionPoolImplTest.java | 205 ++++++++++++++++++ .../session/fake/FakeSessionListener.java | 7 + 7 files changed, 328 insertions(+), 45 deletions(-) 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 fa0c833b6cf7..a5e8557d4269 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 @@ -300,12 +300,29 @@ 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) { + logger.log( + Level.WARNING, + String.format( + "Session error: %s Exception while force-closing stream", info.getLogName()), + t); + } + + notifyTerminalClose( + Status.CANCELLED.withDescription(closeReason.getDescription()), + new Metadata(), + localRpc, + prevState); }); } @@ -778,6 +795,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 +820,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; diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java index 8021e2f78199..543d41b1833c 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java +++ b/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); 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 d2c6892d800a..1c5b48258e3f 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 @@ -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(); @@ -646,6 +652,11 @@ private void onSessionGoAway(SessionHandle handle, GoAwayResponse msg) { "Adding new session to replace a going away session %s", handle.getSession().getLogName())); createSession(handle.getSession().getOpenParams()); + while (poolSizer.getScaleDelta() > 0) { + if (!createSession(openParams)) { + break; + } + } } } finally { poolLock.unlock(); @@ -712,6 +723,11 @@ private void onSessionClose( String.format( "Replacing abnormally closed session %s", handle.getSession().getLogName())); createSession(openParams.withIncrementedAttempts()); + while (poolSizer.getScaleDelta() > 0) { + if (!createSession(openParams)) { + break; + } + } } } } finally { 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 326f971af88a..5e07e86ad9f3 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 @@ -867,6 +867,7 @@ void sessionOpenTimeoutForcesClose() throws Exception { assertWithMessage("terminal status should be delivered after open timeout") .that(sessionListener.popUntil(Status.class)) .isNotNull(); + assertThat(sessionListener.getLastPrevState()).isEqualTo(Session.SessionState.STARTING); sw.reset().start(); while (session.getState() != Session.SessionState.WAIT_SERVER_CLOSE && session.getState() != 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 442df965e809..ae33b60a1879 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 @@ -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. + fakeSession.state = SessionState.CLOSED; + 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); + + // 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. + fakeSession.state = SessionState.CLOSED; + handle.onSessionClosed(SessionState.WAIT_SERVER_CLOSE); + + assertThat(list.getAfesWithReadySessions()).isEmpty(); + assertThat(list.getAllSessions()).isEmpty(); + assertThat(stats.getReadyCount()).isEqualTo(0); + assertThat(stats.getExpectedCapacity()).isEqualTo(0); + } + @Test void testAwaitCloseToSoftClosed() { SessionList list = new SessionList(); 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 04292f8e3133..8d17c987a100 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 @@ -27,14 +27,18 @@ import static org.mockito.Mockito.when; import com.google.bigtable.v2.CloseSessionRequest; +import com.google.bigtable.v2.CloseSessionRequest.CloseSessionReason; import com.google.bigtable.v2.FeatureFlags; import com.google.bigtable.v2.OpenFakeSessionRequest; +import com.google.bigtable.v2.OpenFakeSessionRequest.Action; +import com.google.bigtable.v2.OpenFakeSessionRequest.ActionList; import com.google.bigtable.v2.OpenFakeSessionRequest.StreamError; import com.google.bigtable.v2.OpenSessionRequest; import com.google.bigtable.v2.SessionFakeScriptedRequest; import com.google.bigtable.v2.SessionFakeScriptedResponse; import com.google.bigtable.v2.SessionRefreshConfig; 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; @@ -47,6 +51,7 @@ import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc; import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcCallContext; import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcResult; +import com.google.cloud.bigtable.data.v2.internal.session.SessionList.SessionHandle; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeClock; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeSessionService; @@ -491,6 +496,196 @@ void testCreateSessionDoesntPropagateDeadline() { } } + @Test + @SuppressWarnings("GuardedBy") + void readyIdleSessionAbnormallyClosedIsReplaced() throws Exception { + sessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); + + // Wait until sessions are READY + long deadline = System.currentTimeMillis() + 5000; + ReentrantLock poolLock = extractPoolLock(sessionPool); + while (System.currentTimeMillis() < deadline) { + poolLock.lock(); + try { + if (sessionPool.sessions.getStats().getReadyCount() > 0) { + break; + } + } finally { + poolLock.unlock(); + } + Thread.sleep(20); + } + + poolLock.lock(); + SessionHandle handle; + try { + assertThat(sessionPool.sessions.getStats().getReadyCount()).isGreaterThan(0); + handle = sessionPool.sessions.getAllSessions().iterator().next(); + } finally { + poolLock.unlock(); + } + + int requestCountBefore = fakeService.getOpenRequestCount().get(); + + // Now simulate heartbeat failure / forceClose on the idle session: + // This transitions state -> WAIT_SERVER_CLOSE and cancels stream. + handle + .getSession() + .forceClose( + CloseSessionRequest.newBuilder() + .setReason( + CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_MISSED_HEARTBEAT) + .setDescription("missed heartbeat") + .build()); + + waitForOpenRequestCount(requestCountBefore + 1, Duration.ofSeconds(5)); + + // Expected behavior: The pool should recognize that an idle session was lost below + // min_session_count + // and replace it with a new session. + assertThat(fakeService.getOpenRequestCount().get()).isGreaterThan(requestCountBefore); + } + + @Test + void startingSessionFailedHandshakeIsReplaced() throws Exception { + // When a starting session receives GoAway before open, the pool should replace the session + // to maintain min_session_count. + sessionPool.start( + OpenFakeSessionRequest.newBuilder().setGoAwayBeforeOpen(true).build(), new Metadata()); + + waitForOpenRequestCount(6, Duration.ofSeconds(5)); + + assertThat(fakeService.getOpenRequestCount().get()).isGreaterThan(5); + } + + @Test + @SuppressWarnings("GuardedBy") + void pendingCallOnHotPathCreatesOnlyOneSession() throws Exception { + ActionList delayedAction = + ActionList.newBuilder() + .addActions( + Action.newBuilder() + .setDelay(Durations.fromMillis(2000)) + .setResponse(VirtualRpcResponse.getDefaultInstance())) + .build(); + + sessionPool.start( + OpenFakeSessionRequest.newBuilder().putVrpcActions(0, delayedAction).build(), + new Metadata()); + + // Wait until initial 5 sessions are READY + long deadline = System.currentTimeMillis() + 5000; + ReentrantLock poolLock = extractPoolLock(sessionPool); + while (System.currentTimeMillis() < deadline) { + poolLock.lock(); + try { + if (sessionPool.sessions.getStats().getReadyCount() == 5) { + break; + } + } finally { + poolLock.unlock(); + } + Thread.sleep(20); + } + + int openCountBefore = fakeService.getOpenRequestCount().get(); + assertThat(openCountBefore).isEqualTo(5); + + // Occupy all 5 sessions with active calls + for (int i = 0; i < 5; i++) { + VRpc rpc = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + UnaryResponseFuture f = new UnaryResponseFuture<>(); + rpc.start( + SessionFakeScriptedRequest.newBuilder().setTag(0).build(), + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, vrpcTracer), + f); + } + + // Now all 5 sessions are in-use. Start an additional call. + // On the hot path, only 1 session is created even though scaleDelta > 1. + VRpc pendingCall = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + UnaryResponseFuture pf = new UnaryResponseFuture<>(); + pendingCall.start( + SessionFakeScriptedRequest.newBuilder().setTag(0).build(), + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, vrpcTracer), + pf); + + waitForOpenRequestCount(openCountBefore + 1, Duration.ofSeconds(5)); + + assertThat(fakeService.getOpenRequestCount().get()).isEqualTo(openCountBefore + 1); + } + + @Test + @SuppressWarnings("GuardedBy") + void abnormalCloseScalesUpMultipleSessions() throws Exception { + ActionList delayedAction = + ActionList.newBuilder() + .addActions( + Action.newBuilder() + .setDelay(Durations.fromMillis(2000)) + .setResponse(VirtualRpcResponse.getDefaultInstance())) + .build(); + + sessionPool.start( + OpenFakeSessionRequest.newBuilder().putVrpcActions(0, delayedAction).build(), + new Metadata()); + + // Wait until initial 5 sessions are READY + long deadline = System.currentTimeMillis() + 5000; + ReentrantLock poolLock = extractPoolLock(sessionPool); + while (System.currentTimeMillis() < deadline) { + poolLock.lock(); + try { + if (sessionPool.sessions.getStats().getReadyCount() == 5) { + break; + } + } finally { + poolLock.unlock(); + } + Thread.sleep(20); + } + + int openCountBefore = fakeService.getOpenRequestCount().get(); + assertThat(openCountBefore).isEqualTo(5); + + // Occupy 4 of the 5 sessions with active calls + for (int i = 0; i < 4; i++) { + VRpc rpc = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + UnaryResponseFuture f = new UnaryResponseFuture<>(); + rpc.start( + SessionFakeScriptedRequest.newBuilder().setTag(0).build(), + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, vrpcTracer), + f); + } + + // Pick the 1 remaining idle session + SessionHandle idleHandle; + poolLock.lock(); + try { + idleHandle = + sessionPool.sessions.getAfesWithReadySessions().get(0).sessions.iterator().next(); + } finally { + poolLock.unlock(); + } + + // When the idle session abnormally closes, onSessionClose scales up all needed sessions + idleHandle + .getSession() + .forceClose( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_MISSED_HEARTBEAT) + .setDescription("missed heartbeat") + .build()); + + waitForOpenRequestCount(openCountBefore + 5, Duration.ofSeconds(5)); + + // 1 replacement + 4 scale delta = 5 new sessions created + assertThat(fakeService.getOpenRequestCount().get()).isEqualTo(openCountBefore + 5); + } + @Nested class RetrySessionCreation { @@ -1048,6 +1243,16 @@ void healthyPool_preExistingInterrupt_doesNotSpuriouslyDegrade() throws Exceptio assertPoolServesVRpc(sessionPool, "after an interrupted caller (pool must remain healthy)"); } + private void waitForOpenRequestCount(int minCount, Duration timeout) throws InterruptedException { + long deadline = System.currentTimeMillis() + timeout.toMillis(); + while (System.currentTimeMillis() < deadline) { + if (fakeService.getOpenRequestCount().get() >= minCount) { + return; + } + Thread.sleep(10); + } + } + private static ReentrantLock extractPoolLock(SessionPoolImpl pool) throws Exception { Field field = SessionPoolImpl.class.getDeclaredField("poolLock"); field.setAccessible(true); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/FakeSessionListener.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/FakeSessionListener.java index d7063eddeda7..f12c235c6ac3 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/FakeSessionListener.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/FakeSessionListener.java @@ -33,6 +33,8 @@ public class FakeSessionListener implements Session.Listener { private BlockingDeque msgs = new LinkedBlockingDeque<>(); + private volatile SessionState lastPrevState; + @Override public void onReady(OpenSessionResponse msg) { msgs.add(msg); @@ -45,9 +47,14 @@ public void onGoAway(GoAwayResponse msg) { @Override public void onClose(SessionState prevState, Status status, Metadata trailers) { + this.lastPrevState = prevState; msgs.add(status); } + public SessionState getLastPrevState() { + return lastPrevState; + } + @SuppressWarnings("unchecked") public T popUntil(Class cls) throws InterruptedException, TimeoutException { List seen = new ArrayList<>(); From 1b961603397a0c99474a55c07bc4171df853f17e Mon Sep 17 00:00:00 2001 From: Yuri Golobokov Date: Fri, 18 Sep 2026 21:45:55 +0000 Subject: [PATCH 2/2] address review comments --- .../data/v2/internal/session/SessionImpl.java | 3 +- .../v2/internal/session/SessionPoolImpl.java | 33 +++---- .../v2/internal/session/SessionImplTest.java | 37 +++++--- .../v2/internal/session/SessionListTest.java | 4 +- .../internal/session/SessionPoolImplTest.java | 92 ++++++++++++++++++- 5 files changed, 131 insertions(+), 38 deletions(-) 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 a5e8557d4269..39366e0e09f6 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 @@ -311,6 +311,7 @@ public void forceClose(CloseSessionRequest closeReason) { try { stream.forceClose(closeReason.getDescription(), null); } catch (Throwable t) { + debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_force_close_failed"); logger.log( Level.WARNING, String.format( @@ -319,7 +320,7 @@ public void forceClose(CloseSessionRequest closeReason) { } notifyTerminalClose( - Status.CANCELLED.withDescription(closeReason.getDescription()), + Status.UNAVAILABLE.withDescription(closeReason.getDescription()), new Metadata(), localRpc, prevState); 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 1c5b48258e3f..502cdcacb637 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 @@ -652,11 +652,6 @@ private void onSessionGoAway(SessionHandle handle, GoAwayResponse msg) { "Adding new session to replace a going away session %s", handle.getSession().getLogName())); createSession(handle.getSession().getOpenParams()); - while (poolSizer.getScaleDelta() > 0) { - if (!createSession(openParams)) { - break; - } - } } } finally { poolLock.unlock(); @@ -703,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, @@ -723,11 +723,6 @@ private void onSessionClose( String.format( "Replacing abnormally closed session %s", handle.getSession().getLogName())); createSession(openParams.withIncrementedAttempts()); - while (poolSizer.getScaleDelta() > 0) { - if (!createSession(openParams)) { - break; - } - } } } } finally { 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 5e07e86ad9f3..d995a32f6d10 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 @@ -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,18 +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(); + Status status = sessionListener.popUntil(Status.class); + assertThat(status.getCode()).isEqualTo(Status.Code.UNAVAILABLE); assertThat(sessionListener.getLastPrevState()).isEqualTo(Session.SessionState.STARTING); - 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); + assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED); } // endregion 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 ae33b60a1879..e44d795f13e5 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 @@ -426,7 +426,6 @@ void testStartingSessionClosedAfterWaitServerClose_decrementsStartingCount() { // 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. - fakeSession.state = SessionState.CLOSED; handle.onSessionClosed(SessionState.WAIT_SERVER_CLOSE); assertThat(list.getAllSessions()).isEmpty(); @@ -447,15 +446,16 @@ void testReadyIdleSessionClosedAfterWaitServerClose_cleansUpReadySessionsAndRead 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. - fakeSession.state = SessionState.CLOSED; 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); } 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 8d17c987a100..207265693c93 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 @@ -84,6 +84,7 @@ import java.lang.reflect.Field; import java.time.Duration; import java.time.Instant; +import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CopyOnWriteArrayList; @@ -619,7 +620,7 @@ void pendingCallOnHotPathCreatesOnlyOneSession() throws Exception { @Test @SuppressWarnings("GuardedBy") - void abnormalCloseScalesUpMultipleSessions() throws Exception { + void abnormalCloseReplacesSession() throws Exception { ActionList delayedAction = ActionList.newBuilder() .addActions( @@ -671,7 +672,7 @@ void abnormalCloseScalesUpMultipleSessions() throws Exception { poolLock.unlock(); } - // When the idle session abnormally closes, onSessionClose scales up all needed sessions + // When the idle session abnormally closes, onSessionClose performs a 1:1 replacement idleHandle .getSession() .forceClose( @@ -680,10 +681,91 @@ void abnormalCloseScalesUpMultipleSessions() throws Exception { .setDescription("missed heartbeat") .build()); - waitForOpenRequestCount(openCountBefore + 5, Duration.ofSeconds(5)); + waitForOpenRequestCount(openCountBefore + 1, Duration.ofSeconds(5)); + + assertThat(fakeService.getOpenRequestCount().get()).isEqualTo(openCountBefore + 1); + } + + @Test + @SuppressWarnings("GuardedBy") + void readySessionsHeartbeatMissDoesNotIncrementConsecutiveFailures() throws Exception { + ActionList delayedAction = + ActionList.newBuilder() + .addActions( + Action.newBuilder() + .setDelay(Durations.fromMillis(500)) + .setResponse(VirtualRpcResponse.getDefaultInstance())) + .build(); + + sessionPool.start( + OpenFakeSessionRequest.newBuilder().putVrpcActions(0, delayedAction).build(), + new Metadata()); + + // Wait until initial 5 sessions are READY + long deadline = System.currentTimeMillis() + 5000; + ReentrantLock poolLock = extractPoolLock(sessionPool); + while (System.currentTimeMillis() < deadline) { + poolLock.lock(); + try { + if (sessionPool.sessions.getStats().getReadyCount() == 5) { + break; + } + } finally { + poolLock.unlock(); + } + Thread.sleep(20); + } + + // Occupy all 5 sessions with active calls + for (int i = 0; i < 5; i++) { + VRpc rpc = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + UnaryResponseFuture f = new UnaryResponseFuture<>(); + rpc.start( + SessionFakeScriptedRequest.newBuilder().setTag(0).build(), + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, vrpcTracer), + f); + } + + // Capture all 5 ready session handles before starting additional calls + List handles; + poolLock.lock(); + try { + handles = new ArrayList<>(sessionPool.sessions.getAllSessions()); + } finally { + poolLock.unlock(); + } + assertThat(handles).hasSize(5); + + // Start a 6th call that must queue in pendingRpcs because all 5 sessions are occupied + UnaryResponseFuture pendingFuture = new UnaryResponseFuture<>(); + VRpc queuedRpc = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + queuedRpc.start( + SessionFakeScriptedRequest.newBuilder().setTag(0).build(), + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, vrpcTracer), + pendingFuture); + + // Verify the 6th call could not be immediately dispatched and is queued + assertThat(pendingFuture.isDone()).isFalse(); + + // Force-close all 5 ready sessions (simulating simultaneous missed heartbeat) + for (SessionHandle handle : handles) { + handle + .getSession() + .forceClose( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_MISSED_HEARTBEAT) + .setDescription("missed heartbeat") + .build()); + } - // 1 replacement + 4 scale delta = 5 new sessions created - assertThat(fakeService.getOpenRequestCount().get()).isEqualTo(openCountBefore + 5); + // If consecutiveFailures incremented on READY force-closes, consecutiveFailures would hit 5 + // and popClosableRpcs() would have rejected pendingFuture with REJECTED. + // With the fix, consecutiveFailures does NOT increment for READY sessions. The pending vRPC + // remains queued, replacement sessions become READY, and it completes successfully. + SessionFakeScriptedResponse response = pendingFuture.get(5, TimeUnit.SECONDS); + assertThat(response).isNotNull(); } @Nested