diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizer.java index 01a83a661486..3e9d166d9ae0 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizer.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizer.java @@ -42,6 +42,8 @@ class PoolSizer { // the session is ready it takes 10-40ms to process all pending calls. private volatile int pendingVRpcsPerSession; + private volatile boolean disabled; + interface Sized { int getSize(); } @@ -52,11 +54,7 @@ interface Sized { SessionClientConfiguration.SessionPoolConfiguration poolConfig) { this.stats = stats; this.pendingRpcs = pendingRpcs; - - this.idlesSessionHeadRoom = poolConfig.getHeadroom(); - this.minIdleSessions = poolConfig.getMinSessionCount(); - this.maxIdleSessions = poolConfig.getMaxSessionCount(); - this.pendingVRpcsPerSession = poolConfig.getNewSessionQueueLength(); + updateConfig(poolConfig); } void updateConfig(SessionClientConfiguration.SessionPoolConfiguration poolConfiguration) { @@ -64,12 +62,23 @@ void updateConfig(SessionClientConfiguration.SessionPoolConfiguration poolConfig this.minIdleSessions = poolConfiguration.getMinSessionCount(); this.maxIdleSessions = poolConfiguration.getMaxSessionCount(); this.pendingVRpcsPerSession = poolConfiguration.getNewSessionQueueLength(); + this.disabled = + poolConfiguration.equals( + SessionClientConfiguration.SessionPoolConfiguration.getDefaultInstance()) + || this.maxIdleSessions <= 0; + } + + boolean isDisabled() { + return disabled; } public int getScaleDelta() { // Assume each session handles 1 RPC at a time. This should be revisited if sessions get // support for multiplexing. - int effectivePending = (int) Math.ceil((float) pendingRpcs.getSize() / pendingVRpcsPerSession); + int effectivePending = + pendingVRpcsPerSession > 0 + ? (int) Math.ceil((float) pendingRpcs.getSize() / pendingVRpcsPerSession) + : pendingRpcs.getSize(); int sessionsInUse = effectivePending + stats.getInUseCount(); int unboundedDesiredIdleSessions = (int) Math.ceil(sessionsInUse * idlesSessionHeadRoom); int desiredIdleSessions = @@ -82,6 +91,9 @@ public int getScaleDelta() { if (desiredCapacity < immediateCapacity) { return desiredCapacity - immediateCapacity; } else if (desiredCapacity > eventualCapacity) { + if (isDisabled()) { + return 0; + } return desiredCapacity - eventualCapacity; } else { return 0; @@ -91,13 +103,13 @@ public int getScaleDelta() { /** Returns true if the session should be replaced */ @CheckReturnValue public boolean handleGoAway(GoAwayResponse msg) { - return getScaleDelta() >= 0; + return getScaleDelta() > 0; } /** Returns true if the closed session should be replaced */ @CheckReturnValue boolean handleSessionClose(Status statusProto) { - return getScaleDelta() >= 0; + return getScaleDelta() > 0; } /** Returns true if a new session should be added. */ diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudget.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudget.java index bfa8dd2213a9..1b4829d82162 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudget.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudget.java @@ -69,11 +69,11 @@ Instant getNextAvailableBudget() { boolean tryReserveSession() { sanityCheck(); - if (concurrentRequests == maxConcurrentRequest) { + if (concurrentRequests >= maxConcurrentRequest) { drainCreationFailures(); } - if (concurrentRequests == maxConcurrentRequest) { + if (concurrentRequests >= maxConcurrentRequest) { return false; } @@ -116,9 +116,8 @@ private void sanityCheck() { if (concurrentRequests > maxConcurrentRequest) { DEFAULT_LOGGER.log( Level.FINE, - "Concurrent requests out of range: {0}. Resetting it to max.", - concurrentRequests); - concurrentRequests = maxConcurrentRequest; + "Concurrent requests ({0}) exceeds max ({1}).", + new Object[] {concurrentRequests, maxConcurrentRequest}); } } 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 543d41b1833c..79fc1109fe23 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 @@ -88,9 +88,9 @@ SessionHandle newHandle(Session session) { void removeUnstartedHandle(SessionHandle handle) { if (allSessions.remove(handle)) { poolStats.startingCount--; - if (handle.inExpectedCount) { + if (!handle.closing) { poolStats.expectedCapacity--; - handle.inExpectedCount = false; + handle.closing = true; } } } @@ -160,7 +160,7 @@ void prune() { @NotThreadSafe class SessionHandle { private final Session session; - private boolean inExpectedCount = true; + private boolean closing = false; @SuppressWarnings("OptionalUsedAsFieldOrParameterType") private Optional afe = Optional.empty(); @@ -180,15 +180,17 @@ void onSessionStarted() { AfeHandle afeHandle = afeHandles.computeIfAbsent(AfeId.extract(peerInfo), (ignored) -> new AfeHandle()); this.afe = Optional.of(afeHandle); - afeHandle.sessions.add(this); afeHandle.refCount++; afeHandle.lastConnected = Instant.now(); - if (afeHandle.sessions.size() == 1) { - afesWithReadySessions.add(afeHandle); - } poolStats.startingCount--; - poolStats.readyCount++; + if (!closing) { + afeHandle.sessions.add(this); + if (afeHandle.sessions.size() == 1) { + afesWithReadySessions.add(afeHandle); + } + poolStats.readyCount++; + } } /** @@ -221,7 +223,7 @@ private void releaseToPool() { poolStats.inUseCount--; inUseSessions.remove(this); - if (session.getState() == SessionState.READY) { + if (!closing && session.getState() == SessionState.READY) { poolStats.readyCount++; afeHandle.sessions.add(this); afeHandle.lastConnected = Instant.now(); @@ -256,8 +258,10 @@ void onSessionClosing() { poolStats.readyCount--; } // eagerly decrement expected count and make sure to avoid double counting in onSessionClosed - poolStats.expectedCapacity--; - inExpectedCount = false; + if (!closing) { + poolStats.expectedCapacity--; + closing = true; + } } void onSessionClosed(SessionState prevState) { @@ -271,9 +275,9 @@ void onSessionClosed(SessionState prevState) { // stranded handle blocks drainedFuture and pool.awaitTerminated hangs until the shared close // deadline elapses. try { - if (inExpectedCount) { + if (!closing) { poolStats.expectedCapacity--; - inExpectedCount = false; + closing = true; } // only update counts after the session started and has an afe associated afe.ifPresent(afeHandle -> afeHandle.refCount--); 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..73a536756acd 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 @@ -169,7 +169,7 @@ private enum PoolState { private final CompletableFuture drainedFuture = new CompletableFuture<>(); @GuardedBy("poolLock") - private BigtableTimer.Timeout retryCreateSessionFuture = null; + private BigtableTimer.Timeout retryCreateSessionTimeout = null; // TODO: get the max from ClientConfiguration @GuardedBy("poolLock") @@ -193,7 +193,7 @@ private enum PoolState { private final DebugTagTracer debugTagTracer; // @SuppressWarnings("GuardedBy"): error-prone flags writes to @GuardedBy("poolLock") fields - // (sessions, picker, poolSizer, pendingRpcs, budget, retryCreateSessionFuture) inside the + // (sessions, picker, poolSizer, pendingRpcs, budget, retryCreateSessionTimeout) inside the // constructor without holding the monitor. This is safe because the object is not yet published // to other threads — no external reference exists until the constructor returns. @SuppressWarnings("GuardedBy") @@ -281,14 +281,50 @@ public SessionPoolImpl( configManager.addListener( clientConfig -> clientConfig.getSessionConfiguration().getSessionPoolConfiguration(), newConfig -> { + List> toCancel = new ArrayList<>(); poolLock.lock(); try { budget.updateConfig(newConfig); poolSizer.updateConfig(newConfig); picker.updateConfig(newConfig); + if (poolSizer.isDisabled()) { + if (retryCreateSessionTimeout != null) { + retryCreateSessionTimeout.cancel(); + retryCreateSessionTimeout = null; + } + toCancel = popClosableRpcs(); + CloseSessionRequest closeReq = + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_DOWNSIZE) + .setDescription("SessionPool configuration disabled") + .build(); + for (SessionHandle handle : new ArrayList<>(sessions.getAllSessions())) { + handle.onSessionClosing(); + handle.getSession().close(closeReq); + } + } else if (poolState == PoolState.STARTED) { + while (poolSizer.getScaleDelta() > 0) { + if (!createSession(openParams)) { + break; + } + } + } } finally { poolLock.unlock(); } + + if (!toCancel.isEmpty()) { + VRpcResult cancelResult = + VRpcResult.createRejectedError( + Status.UNAVAILABLE.withDescription("SessionPool configuration disabled")); + for (PendingVRpc vrpc : toCancel) { + try { + vrpc.cancelWithResult(cancelResult); + } catch (Throwable t) { + logger.log(Level.WARNING, "Exception dispatching close to op executor", t); + } + } + } }); } @@ -350,9 +386,9 @@ public void close(CloseSessionRequest req) { afeListPruneTimeout.cancel(); afeListPruneTimeout = null; } - if (retryCreateSessionFuture != null) { - retryCreateSessionFuture.cancel(); - retryCreateSessionFuture = null; + if (retryCreateSessionTimeout != null) { + retryCreateSessionTimeout.cancel(); + retryCreateSessionTimeout = 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. @@ -537,7 +573,7 @@ public void onClose(SessionState prevState, Status status, Metadata trailers) { @GuardedBy("poolLock") private void maybeScheduleCreateSessionRetry() { - if (retryCreateSessionFuture != null) { + if (retryCreateSessionTimeout != null || poolSizer.isDisabled()) { return; } // This interval could be 0 or negative if nextAvailableBudget is now. In which case @@ -547,12 +583,12 @@ private void maybeScheduleCreateSessionRetry() { if (retryIntervalMs < 1) { retryIntervalMs = 1; } - retryCreateSessionFuture = + retryCreateSessionTimeout = timer.newTimeout( () -> { poolLock.lock(); try { - retryCreateSessionFuture = null; + retryCreateSessionTimeout = null; while (poolState != PoolState.CLOSED && poolSizer.getScaleDelta() > 0) { if (!createSession(openParams)) { break; @@ -646,7 +682,8 @@ private void onSessionGoAway(SessionHandle handle, GoAwayResponse msg) { // TODO: remove the check on if the open params is updated. In the future, server should // broadcast session refresh config to all the sessions that needs reconnect and client // just need to recreate sessions when pool sizer requries a replacement. - if (handle.getSession().isOpenParamsUpdated() || poolSizer.handleGoAway(msg)) { + if (!poolSizer.isDisabled() + && (handle.getSession().isOpenParamsUpdated() || poolSizer.handleGoAway(msg))) { logger.fine( String.format( "Adding new session to replace a going away session %s", @@ -885,7 +922,7 @@ class PendingVRpc implements VRpc realCall; - private BigtableTimer.Timeout deadlineMonitor; + private BigtableTimer.Timeout deadlineTimeout; public PendingVRpc(VRpcDescriptor desc) { this.desc = desc; @@ -931,10 +968,23 @@ public void start(ReqT req, VRpcCallContext ctx, VRpcListener listener) { ctx.getExecutor().execute(() -> listener.onClose(result)); return; } + if (poolSizer.isDisabled()) { + VRpcResult result = + VRpcResult.createRejectedError( + Status.UNAVAILABLE.withDescription("SessionPool configuration disabled")); + ctx.getExecutor() + .execute( + () -> { + if (isCancelled) return; + isCancelled = true; + listener.onClose(result); + }); + return; + } // Only arm the deadline monitor after we've committed to queueing; otherwise the // fast-fail early return above would leak a timer that fires later and emits a phantom // tracer.onAttemptFinish on the Active state's stale listener. - this.deadlineMonitor = monitorDeadline(ctx.getOperationInfo().getDeadline()); + this.deadlineTimeout = monitorDeadline(ctx.getOperationInfo().getDeadline()); pendingRpcs.add(this); if (logger.isLoggable(Level.FINE)) { @@ -974,6 +1024,9 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) { // region below. @SuppressWarnings("GuardedBy") private void cancel(Status status, boolean onlyCancelPendingCall) { + if (deadlineTimeout != null) { + deadlineTimeout.cancel(); + } // The eager removal is only an optimization to skip a wasted session pick in // tryDrainPendingRpcs. // Use the bounded hot-path acquisition instead of a blocking lock(): this runs on a callback @@ -1006,6 +1059,9 @@ private void cancel(Status status, boolean onlyCancelPendingCall) { } void cancelWithResult(VRpcResult result) { + if (deadlineTimeout != null) { + deadlineTimeout.cancel(); + } ctx.getExecutor() .execute( () -> { @@ -1032,8 +1088,8 @@ public void requestNext() { } private void drainTo(SessionHandle handle) { - if (deadlineMonitor != null) { - deadlineMonitor.cancel(); + if (deadlineTimeout != null) { + deadlineTimeout.cancel(); } ctx.getExecutor() .execute( diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizerTest.java index 136195dc3312..a343c7973ce8 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizerTest.java @@ -65,6 +65,12 @@ void testTestEquilibrium() { assertThat(sizer.getScaleDelta()).isEqualTo(0); assertThat(sizer.handleNewCall()).isFalse(); + assertThat(sizer.handleGoAway(GoAwayResponse.getDefaultInstance())).isFalse(); + assertThat(sizer.handleSessionClose(Status.getDefaultInstance())).isFalse(); + + // After onSessionClosing / onSessionClosed decrements expectedCapacity below MIN_IDLE_SESSIONS: + stats.getTestHelper().setReadyCount(MIN_IDLE_SESSIONS - 1); + stats.getTestHelper().setExpectedCapacity(MIN_IDLE_SESSIONS - 1); assertThat(sizer.handleGoAway(GoAwayResponse.getDefaultInstance())).isTrue(); assertThat(sizer.handleSessionClose(Status.getDefaultInstance())).isTrue(); } @@ -146,6 +152,12 @@ void testEquilibriumWhenStarting() { assertThat(sizer.getScaleDelta()).isEqualTo(0); // Remain steady in band assertThat(sizer.handleNewCall()).isFalse(); + assertThat(sizer.handleGoAway(GoAwayResponse.getDefaultInstance())).isFalse(); + assertThat(sizer.handleSessionClose(Status.getDefaultInstance())).isFalse(); + + // After onSessionClosing / onSessionClosed decrements expectedCapacity below MIN_IDLE_SESSIONS: + stats.getTestHelper().setStartingCount(MIN_IDLE_SESSIONS - 1); + stats.getTestHelper().setExpectedCapacity(MIN_IDLE_SESSIONS - 1); assertThat(sizer.handleGoAway(GoAwayResponse.getDefaultInstance())).isTrue(); assertThat(sizer.handleSessionClose(Status.getDefaultInstance())).isTrue(); } @@ -213,4 +225,28 @@ void testEquilibriumWhenStartingWithLoad() { // = 17 sessions needed, but we only have 15 assertThat(sizer.getScaleDelta()).isEqualTo(2); } + + @Test + void testEmptyConfig() { + PoolStats stats = new PoolStats(); + AtomicInteger pendingCount = new AtomicInteger(5); + PoolSizer sizer = + new PoolSizer( + stats, + pendingCount::get, + SessionClientConfiguration.SessionPoolConfiguration.getDefaultInstance()); + + // Even with pending RPCs and in-use sessions, pendingVRpcsPerSession == 0 must not divide by 0 + // or overflow int. + assertThat(sizer.getScaleDelta()).isEqualTo(0); + assertThat(sizer.handleNewCall()).isFalse(); + assertThat(sizer.handleGoAway(GoAwayResponse.getDefaultInstance())).isFalse(); + assertThat(sizer.handleSessionClose(Status.getDefaultInstance())).isFalse(); + + stats.getTestHelper().setInUseCount(2); + stats.getTestHelper().setExpectedCapacity(0); + assertThat(sizer.getScaleDelta()).isEqualTo(0); + assertThat(sizer.handleGoAway(GoAwayResponse.getDefaultInstance())).isFalse(); + assertThat(sizer.handleSessionClose(Status.getDefaultInstance())).isFalse(); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudgetTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudgetTest.java index 1eb9bbd467b9..a3f143f49b0b 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudgetTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudgetTest.java @@ -17,7 +17,9 @@ import static com.google.common.truth.Truth.assertThat; +import com.google.bigtable.v2.SessionClientConfiguration.SessionPoolConfiguration; import com.google.cloud.bigtable.data.v2.internal.session.fake.FakeClock; +import com.google.protobuf.util.Durations; import java.time.Duration; import java.time.Instant; import org.junit.jupiter.api.BeforeEach; @@ -56,4 +58,27 @@ public void test() { clock.increment(penalty.plus(Duration.ofNanos(1))); assertThat(budget.tryReserveSession()).isTrue(); } + + @Test + public void testUpdateConfigToEmptyAfterFailure() { + SessionCreationBudget budget = new SessionCreationBudget(2, penalty, clock); + assertThat(budget.tryReserveSession()).isTrue(); + budget.onSessionCreationFailure(); + + // Expire the failure token, then update config to empty (budget = 0). + clock.increment(penalty.plus(Duration.ofNanos(1))); + budget.updateConfig(SessionPoolConfiguration.getDefaultInstance()); + + // Must not allow reserving when budget is 0, even with an expired failure token. + assertThat(budget.tryReserveSession()).isFalse(); + + // Re-enabling budget should allow reservations again. + budget.updateConfig( + SessionPoolConfiguration.newBuilder() + .setNewSessionCreationBudget(1) + .setNewSessionCreationPenalty(Durations.fromMinutes(1)) + .build()); + assertThat(budget.tryReserveSession()).isTrue(); + assertThat(budget.tryReserveSession()).isFalse(); + } } 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..9313bde6be75 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 @@ -286,6 +286,56 @@ void testClosingInUseToAwaitCloseWhileRpcIsActive() { assertThat(stats.getExpectedCapacity()).isEqualTo(0); } + @Test + void testStartingClosingRaceWithSessionStarted() { + SessionList list = new SessionList(); + PoolStats stats = list.getStats(); + + fakeSession.state = SessionState.STARTING; + SessionHandle handle = list.newHandle(fakeSession); + + handle.onSessionClosing(); + assertThat(stats.getExpectedCapacity()).isEqualTo(0); + assertThat(stats.getStartingCount()).isEqualTo(1); + + // If onSessionStarted arrives after onSessionClosing, it must not mark the session ready. + fakeSession.state = SessionState.READY; + handle.onSessionStarted(); + + assertThat(list.getAfesWithReadySessions()).isEmpty(); + assertThat(stats.getStartingCount()).isEqualTo(0); + assertThat(stats.getReadyCount()).isEqualTo(0); + assertThat(stats.getExpectedCapacity()).isEqualTo(0); + + fakeSession.state = SessionState.CLOSED; + handle.onSessionClosed(SessionState.WAIT_SERVER_CLOSE); + assertThat(list.getAllSessions()).isEmpty(); + } + + @Test + void testInUseClosingRaceWithVRpcFinishBeforeSessionStateUpdates() { + SessionList list = new SessionList(); + PoolStats stats = list.getStats(); + + fakeSession.state = SessionState.STARTING; + SessionHandle handle = list.newHandle(fakeSession); + + fakeSession.state = SessionState.READY; + handle.onSessionStarted(); + + handle = list.checkoutSession(list.getAfesWithReadySessions().get(0)).get(); + + // Pool initiates close (onSessionClosing), but SessionImpl.close() has not yet transitioned + // session.getState() out of READY on sessionSyncContext when onVRpcFinish runs. + handle.onSessionClosing(); + handle.onVRpcFinish(FAKE_VRPC_E2E_LATENCY, OK_RESULT); + + assertThat(list.getAfesWithReadySessions()).isEmpty(); + assertThat(stats.getReadyCount()).isEqualTo(0); + assertThat(stats.getInUseCount()).isEqualTo(0); + assertThat(stats.getExpectedCapacity()).isEqualTo(0); + } + @Test void testStartingToClosed() { 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 207265693c93..a38f7c0e463c 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 @@ -34,6 +34,7 @@ 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.SessionClientConfiguration.SessionPoolConfiguration; import com.google.bigtable.v2.SessionFakeScriptedRequest; import com.google.bigtable.v2.SessionFakeScriptedResponse; import com.google.bigtable.v2.SessionRefreshConfig; @@ -1325,6 +1326,85 @@ void healthyPool_preExistingInterrupt_doesNotSpuriouslyDegrade() throws Exceptio assertPoolServesVRpc(sessionPool, "after an interrupted caller (pool must remain healthy)"); } + @Test + @SuppressWarnings({"unchecked", "GuardedBy"}) + void configUpdateToEmptyClosesSessionsAndScalesBackUpWhenEnabled() throws Exception { + ArgumentCaptor> + listenerCaptor = ArgumentCaptor.forClass(ClientConfigurationManager.ConfigListener.class); + verify(configManager).addListener(any(), listenerCaptor.capture()); + ClientConfigurationManager.ConfigListener configListener = + listenerCaptor.getValue(); + + sessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); + assertPoolServesVRpc(sessionPool, "initial started pool"); + waitForOpenRequestCount(5, Duration.ofSeconds(5)); + assertThat(sessionPool.hasSession()).isTrue(); + + // Queue a pending vRPC while no ready sessions are available, then update config to empty and + // verify the pending vRPC is rejected immediately instead of hanging until deadline. + ReentrantLock poolLock = extractPoolLock(sessionPool); + UnaryResponseFuture pendingFuture = new UnaryResponseFuture<>(); + poolLock.lock(); + try { + for (SessionHandle handle : new ArrayList<>(sessionPool.sessions.getAllSessions())) { + handle.onSessionClosing(); + } + VRpc queuedRpc = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + queuedRpc.start( + SessionFakeScriptedRequest.newBuilder().setTag(0).build(), + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, vrpcTracer), + pendingFuture); + assertThat(pendingFuture.isDone()).isFalse(); + } finally { + poolLock.unlock(); + } + + // Update config to empty (session_load == 0 clears session_configuration). + configListener.onChange(SessionPoolConfiguration.getDefaultInstance()); + assertThat(sessionPool.hasSession()).isFalse(); + ExecutionException pendingError = + assertThrows(ExecutionException.class, () -> pendingFuture.get(5, TimeUnit.SECONDS)); + assertThat(pendingError).hasCauseThat().isInstanceOf(VRpcException.class); + VRpcException vrpcError = (VRpcException) pendingError.getCause(); + assertThat(vrpcError.getResult().getRejected()).isTrue(); + assertThat(vrpcError.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + + // Wait for all sessions to finish closing. + long deadline = System.currentTimeMillis() + 5000; + while (System.currentTimeMillis() < deadline + && !sessionPool.sessions.getAllSessions().isEmpty()) { + Thread.sleep(10); + } + assertThat(sessionPool.sessions.getAllSessions()).isEmpty(); + assertThat(sessionPool.sessions.getStats().getExpectedCapacity()).isEqualTo(0); + + // Starting a new call while config is disabled must also fail immediately without queueing. + UnaryResponseFuture disabledCallFuture = + new UnaryResponseFuture<>(); + VRpc disabledRpc = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + disabledRpc.start( + SessionFakeScriptedRequest.newBuilder().setTag(0).build(), + VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, vrpcTracer), + disabledCallFuture); + ExecutionException disabledCallError = + assertThrows(ExecutionException.class, () -> disabledCallFuture.get(5, TimeUnit.SECONDS)); + assertThat(disabledCallError).hasCauseThat().isInstanceOf(VRpcException.class); + assertThat(((VRpcException) disabledCallError.getCause()).getResult().getRejected()).isTrue(); + + // Re-enable SessionPoolConfiguration and verify pool scales back up. + SessionPoolConfiguration enabledConfig = + ClientConfigurationManager.loadDefault() + .getSessionConfiguration() + .getSessionPoolConfiguration(); + configListener.onChange(enabledConfig); + + waitForOpenRequestCount(10, Duration.ofSeconds(5)); + assertPoolServesVRpc(sessionPool, "after re-enabling session pool configuration"); + assertThat(sessionPool.hasSession()).isTrue(); + } + private void waitForOpenRequestCount(int minCount, Duration timeout) throws InterruptedException { long deadline = System.currentTimeMillis() + timeout.toMillis(); while (System.currentTimeMillis() < deadline) {