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..e1b2f1d4f027 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 @@ -69,7 +69,10 @@ void updateConfig(SessionClientConfiguration.SessionPoolConfiguration poolConfig 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) + : 0; int sessionsInUse = effectivePending + stats.getInUseCount(); int unboundedDesiredIdleSessions = (int) Math.ceil(sessionsInUse * idlesSessionHeadRoom); int desiredIdleSessions = @@ -91,13 +94,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/SessionList.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java index 543d41b1833c..439b22c76689 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 @@ -256,8 +256,10 @@ void onSessionClosing() { poolStats.readyCount--; } // eagerly decrement expected count and make sure to avoid double counting in onSessionClosed - poolStats.expectedCapacity--; - inExpectedCount = false; + if (inExpectedCount) { + poolStats.expectedCapacity--; + inExpectedCount = false; + } } void onSessionClosed(SessionState 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 502cdcacb637..9fa8ba0c5abf 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 @@ -286,6 +286,28 @@ public SessionPoolImpl( budget.updateConfig(newConfig); poolSizer.updateConfig(newConfig); picker.updateConfig(newConfig); + if (newConfig.equals( + SessionClientConfiguration.SessionPoolConfiguration.getDefaultInstance())) { + if (retryCreateSessionFuture != null) { + retryCreateSessionFuture.cancel(); + retryCreateSessionFuture = null; + } + 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(); } @@ -537,7 +559,7 @@ public void onClose(SessionState prevState, Status status, Metadata trailers) { @GuardedBy("poolLock") private void maybeScheduleCreateSessionRetry() { - if (retryCreateSessionFuture != null) { + if (retryCreateSessionFuture != null || budget.getMaxConcurrentRequest() <= 0) { return; } // This interval could be 0 or negative if nextAvailableBudget is now. In which case 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..3eef4df004dc 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(2); + 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/SessionPoolImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java index 207265693c93..e5ba95a113c6 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,46 @@ 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(); + + // Update config to empty (session_load == 0 clears session_configuration). + configListener.onChange(SessionPoolConfiguration.getDefaultInstance()); + assertThat(sessionPool.hasSession()).isFalse(); + + // 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); + + // 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) {