Skip to content

fix(bigtable): handle empty SessionPoolConfiguration in SessionPoolImpl and PoolSizer #14912

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
igorbernstein2 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(bigtable): handle empty SessionPoolConfiguration in SessionPoolImpl and PoolSizer #14912

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions ...-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

@nimf nimf Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should it be : pendingRpcs.getSize(); instead? I.e., when pendingVRpcsPerSession is not defined we assume it's 1:1. I know it is not applicable for when config is empty / sessions are disabled, but we are not specifically checking it here.

int sessionsInUse = effectivePending + stats.getInUseCount();
int unboundedDesiredIdleSessions = (int) Math.ceil(sessionsInUse * idlesSessionHeadRoom);
int desiredIdleSessions =
Expand All @@ -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. */
Expand Down
6 changes: 4 additions & 2 deletions ...igtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
24 changes: 23 additions & 1 deletion ...ble/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -286,6 +286,28 @@ public SessionPoolImpl(
budget.updateConfig(newConfig);
poolSizer.updateConfig(newConfig);
picker.updateConfig(newConfig);
if (newConfig.equals(
SessionClientConfiguration.SessionPoolConfiguration.getDefaultInstance())) {

@nimf nimf Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we also cancel all pending vRPCs when config becomes empty? Otherwise they will be hanging out there until their deadlines.

if (retryCreateSessionFuture != null) {
retryCreateSessionFuture.cancel();
Comment thread
nimf marked this conversation as resolved.
Resolved
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();

@nimf nimf Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is subject to two races

  1. Race with a starting session becoming ready.
  2. Race with an in-use session returning to the pool.

in both cases it will keep AFE in afesWithReadySessions and readyCount > 0.

We can protect from this by guarding against inExpectedCount in the SessionHandle.onSessionStarted() and SessionHandle.releaseToPool()

// In SessionHandle.onSessionStarted():
poolStats.startingCount--;
if (inExpectedCount) {
  afeHandle.sessions.add(this);
  if (afeHandle.sessions.size() == 1) {
    afesWithReadySessions.add(afeHandle);
  }
  poolStats.readyCount++;
}
// In SessionHandle.releaseToPool():
if (inExpectedCount && session.getState() == SessionState.READY) {

handle.getSession().close(closeReq);
}
} else if (poolState == PoolState.STARTED) {
while (poolSizer.getScaleDelta() > 0) {
if (!createSession(openParams)) {
break;
}
}
}
} finally {
poolLock.unlock();
}
Expand Down Expand Up @@ -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
Expand Down
36 changes: 36 additions & 0 deletions ...table/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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();
}
}
41 changes: 41 additions & 0 deletions ...src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<ClientConfigurationManager.ConfigListener<SessionPoolConfiguration>>
listenerCaptor =
ArgumentCaptor.forClass(ClientConfigurationManager.ConfigListener.class);
verify(configManager).addListener(any(), listenerCaptor.capture());
ClientConfigurationManager.ConfigListener<SessionPoolConfiguration> 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) {
Expand Down
Loading