From a6c0062636db67652bef2f9662659fb5b1234c28 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 3 Nov 2026 13:21:55 -0500 Subject: [PATCH 1/2] fix(bigtable): disable SwitchingChannelPool when ChannelPoolConfiguration is empty --- .../channels/SwitchingChannelPool.java | 63 ++++-- .../data/v2/internal/api/ClientTest.java | 30 +-- .../channels/SwitchingChannelPoolTest.java | 190 ++++++++++++++++++ 3 files changed, 252 insertions(+), 31 deletions(-) create mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPoolTest.java diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPool.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPool.java index c34fa30ac1d9..65f0e82da91a 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPool.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPool.java @@ -30,6 +30,7 @@ import io.grpc.ClientInterceptor; import io.grpc.ManagedChannel; import io.grpc.MethodDescriptor; +import io.grpc.Status; import java.time.Duration; import java.util.concurrent.ScheduledExecutorService; import java.util.function.Supplier; @@ -69,19 +70,11 @@ public SwitchingChannelPool( currentConfiguration = configManager.getClientConfiguration().getSessionConfiguration().getChannelConfiguration(); - if (channelProvider.getFallback().isPresent()) { - fallbackAvailable = true; - delegate = createChannelPoolForCurrentMode(); - } else { - // Depending on whether DirectAccess is avaialbe, the channelProvider - // is either a DirectAccess with fallback to Cloudpath - // or a basic Cloudpath. - // If it's a basic cloudpath we can't do any fallback, - // and we can't do directaccess without fallback. - // Essentially, fallback configuration is not applicable here. - fallbackAvailable = false; - delegate = newChannelPoolFromProvider(channelProvider); - } + // Depending on whether DirectAccess is available, the channelProvider is either a DirectAccess + // with fallback to Cloudpath or a basic Cloudpath. If it's a basic cloudpath we can't do any + // fallback, and we can't do directaccess without fallback. + fallbackAvailable = channelProvider.getFallback().isPresent(); + delegate = createDelegate(); // Subscribe to config changes. configListener = @@ -95,10 +88,14 @@ public synchronized void updateConfig(ChannelPoolConfiguration newConfig) { return; } - if (fallbackAvailable && currentConfiguration.getModeCase() != newConfig.getModeCase()) { - // Mode has changed and we can change modes. + boolean wasDisabled = isDisabled(currentConfiguration); + boolean nowDisabled = isDisabled(newConfig); + if (wasDisabled != nowDisabled + || (!nowDisabled + && fallbackAvailable + && currentConfiguration.getModeCase() != newConfig.getModeCase())) { currentConfiguration = newConfig; - ChannelPool cp = createChannelPoolForCurrentMode(); + ChannelPool cp = createDelegate(); replaceChannelPool(cp); return; } @@ -107,6 +104,20 @@ public synchronized void updateConfig(ChannelPoolConfiguration newConfig) { delegate.updateConfig(currentConfiguration); } + private static boolean isDisabled(ChannelPoolConfiguration config) { + return config.equals(ChannelPoolConfiguration.getDefaultInstance()); + } + + private ChannelPool createDelegate() { + if (isDisabled(currentConfiguration)) { + return NoopChannelPool.INSTANCE; + } + if (fallbackAvailable) { + return createChannelPoolForCurrentMode(); + } + return newChannelPoolFromProvider(channelProvider); + } + private ChannelPool createChannelPoolForCurrentMode() { switch (currentConfiguration.getModeCase()) { default: @@ -218,4 +229,24 @@ private Supplier channelSupplier( ChannelProvider channelProvider, ClientInterceptor... interceptors) { return () -> channelProvider.newChannelBuilder().intercept(interceptors).build(); } + + private static class NoopChannelPool implements ChannelPool { + private static final NoopChannelPool INSTANCE = new NoopChannelPool(); + + @Override + public SessionStream newStream( + MethodDescriptor desc, CallOptions callOptions) { + return new FailingSessionStream( + Status.UNAVAILABLE.withDescription("ChannelPool is disabled")); + } + + @Override + public void start() {} + + @Override + public void updateConfig(ChannelPoolConfiguration config) {} + + @Override + public void close() {} + } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/ClientTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/ClientTest.java index 1abbf1274d88..a1739f70d45a 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/ClientTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/ClientTest.java @@ -73,21 +73,21 @@ public class ClientTest { @BeforeEach void setUp() throws IOException { ClientConfiguration.Builder builder = ClientConfigurationManager.loadDefault().toBuilder(); - builder.setSessionConfiguration( - SessionClientConfiguration.newBuilder() - .setSessionDiversionConfiguration( - SessionDiversionConfiguration.newBuilder() - .setGlobalDiversion( - SessionScopeDiversionConfiguration.newBuilder().setSessionLoad(1))) - .setSessionPoolConfiguration( - SessionClientConfiguration.SessionPoolConfiguration.newBuilder() - .setNewSessionCreationBudget(10) - .setNewSessionCreationPenalty(Durations.fromMinutes(1)) - .setHeadroom(0.5f) - .setMinSessionCount(5) - .setMaxSessionCount(20) - .setNewSessionQueueLength(10) - .setConsecutiveSessionFailureThreshold(10))); + builder + .getSessionConfigurationBuilder() + .setSessionDiversionConfiguration( + SessionDiversionConfiguration.newBuilder() + .setGlobalDiversion( + SessionScopeDiversionConfiguration.newBuilder().setSessionLoad(1))) + .setSessionPoolConfiguration( + SessionClientConfiguration.SessionPoolConfiguration.newBuilder() + .setNewSessionCreationBudget(10) + .setNewSessionCreationPenalty(Durations.fromMinutes(1)) + .setHeadroom(0.5f) + .setMinSessionCount(5) + .setMaxSessionCount(20) + .setNewSessionQueueLength(10) + .setConsecutiveSessionFailureThreshold(10)); defaultConfig = builder.build(); executor = Executors.newScheduledThreadPool(4); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPoolTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPoolTest.java new file mode 100644 index 000000000000..3b35349928d5 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPoolTest.java @@ -0,0 +1,190 @@ +/* + * 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.channels; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.bigtable.v2.ClientConfiguration; +import com.google.bigtable.v2.FakeSessionGrpc; +import com.google.bigtable.v2.SessionClientConfiguration; +import com.google.bigtable.v2.SessionClientConfiguration.ChannelPoolConfiguration; +import com.google.bigtable.v2.SessionClientConfiguration.ChannelPoolConfiguration.CloudPathOnly; +import com.google.bigtable.v2.SessionClientConfiguration.ChannelPoolConfiguration.DirectAccessOnly; +import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ChannelProvider; +import com.google.cloud.bigtable.data.v2.internal.channels.SessionStream.Listener; +import com.google.cloud.bigtable.data.v2.internal.csm.Metrics; +import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; +import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; +import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager.ListenerHandle; +import io.grpc.CallOptions; +import io.grpc.ClientInterceptor; +import io.grpc.ManagedChannel; +import io.grpc.ManagedChannelBuilder; +import io.grpc.Metadata; +import io.grpc.Status; +import java.util.Optional; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +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.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class SwitchingChannelPoolTest { + @Mock private ChannelProvider primaryProvider; + @Mock private ChannelProvider fallbackProvider; + @Mock private ManagedChannelBuilder primaryBuilder; + @Mock private ManagedChannelBuilder fallbackBuilder; + @Mock private ManagedChannel primaryChannel; + @Mock private ManagedChannel fallbackChannel; + @Mock private ClientConfigurationManager configManager; + @Mock private ListenerHandle listenerHandle; + @Mock private ScheduledExecutorService backgroundExecutor; + @Mock private ScheduledFuture scheduledFuture; + @Mock private Listener streamListener; + + private Metrics metrics; + + private static final ChannelPoolConfiguration ENABLED_CLOUD_CONFIG = + ChannelPoolConfiguration.newBuilder() + .setMinServerCount(2) + .setMaxServerCount(10) + .setPerServerSessionCount(5) + .setCloudPathOnly(CloudPathOnly.getDefaultInstance()) + .build(); + + private static final ChannelPoolConfiguration ENABLED_DIRECT_CONFIG = + ChannelPoolConfiguration.newBuilder() + .setMinServerCount(2) + .setMaxServerCount(10) + .setPerServerSessionCount(5) + .setDirectAccessOnly(DirectAccessOnly.getDefaultInstance()) + .build(); + + @BeforeEach + @SuppressWarnings({"unchecked", "rawtypes"}) + void setUp() { + metrics = new NoopMetrics(); + lenient().when(configManager.addListener(any(), any())).thenReturn(listenerHandle); + lenient() + .when( + backgroundExecutor.scheduleWithFixedDelay( + any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class))) + .thenReturn((ScheduledFuture) scheduledFuture); + + lenient() + .when(primaryProvider.newChannelBuilder()) + .thenReturn((ManagedChannelBuilder) primaryBuilder); + lenient() + .when(primaryBuilder.intercept(any(ClientInterceptor[].class))) + .thenReturn((ManagedChannelBuilder) primaryBuilder); + lenient().when(primaryBuilder.build()).thenReturn(primaryChannel); + + lenient() + .when(fallbackProvider.newChannelBuilder()) + .thenReturn((ManagedChannelBuilder) fallbackBuilder); + lenient() + .when(fallbackBuilder.intercept(any(ClientInterceptor[].class))) + .thenReturn((ManagedChannelBuilder) fallbackBuilder); + lenient().when(fallbackBuilder.build()).thenReturn(fallbackChannel); + } + + private static ClientConfiguration clientConfigWithChannelConfig( + ChannelPoolConfiguration channelConfig) { + return ClientConfiguration.newBuilder() + .setSessionConfiguration( + SessionClientConfiguration.newBuilder().setChannelConfiguration(channelConfig)) + .build(); + } + + @Test + void startWithEmptyConfigUsesNoopPoolAndTransitionsToEnabledAndBack() { + when(primaryProvider.getFallback()).thenReturn(Optional.empty()); + when(configManager.getClientConfiguration()) + .thenReturn(ClientConfiguration.getDefaultInstance()); + + SwitchingChannelPool pool = + new SwitchingChannelPool(primaryProvider, configManager, metrics, backgroundExecutor); + pool.start(); + + // No channels created when config is default/empty. + verify(primaryProvider, never()).newChannelBuilder(); + + SessionStream disabledStream = + pool.newStream(FakeSessionGrpc.getOpenSessionMethod(), CallOptions.DEFAULT); + disabledStream.start(streamListener, new Metadata()); + ArgumentCaptor statusCaptor = ArgumentCaptor.forClass(Status.class); + verify(streamListener).onClose(statusCaptor.capture(), any(Metadata.class)); + assertThat(statusCaptor.getValue().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + + // Enable channel pool via updateConfig. + pool.updateConfig(ENABLED_CLOUD_CONFIG); + verify(primaryProvider, times(2)).newChannelBuilder(); + verify(primaryChannel, never()).shutdown(); + + // Disable channel pool via updateConfig(getDefaultInstance()) -> closes active pool. + pool.updateConfig(ChannelPoolConfiguration.getDefaultInstance()); + verify(primaryChannel, times(2)).shutdown(); + + Listener secondListener = mock(Listener.class); + SessionStream disabledAgainStream = + pool.newStream(FakeSessionGrpc.getOpenSessionMethod(), CallOptions.DEFAULT); + disabledAgainStream.start(secondListener, new Metadata()); + verify(secondListener).onClose(statusCaptor.capture(), any(Metadata.class)); + assertThat(statusCaptor.getValue().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + + pool.close(); + } + + @Test + void fallbackAvailableTransitionsBetweenEnabledAndDisabled() { + when(primaryProvider.getFallback()).thenReturn(Optional.of(fallbackProvider)); + when(configManager.getClientConfiguration()) + .thenReturn(clientConfigWithChannelConfig(ENABLED_DIRECT_CONFIG)); + + SwitchingChannelPool pool = + new SwitchingChannelPool(primaryProvider, configManager, metrics, backgroundExecutor); + pool.start(); + + verify(primaryProvider, times(2)).newChannelBuilder(); + verify(primaryChannel, never()).shutdown(); + + // Transition to disabled (e.g. session_load = 0 clears session_configuration). + pool.updateConfig(ChannelPoolConfiguration.getDefaultInstance()); + verify(primaryChannel, times(2)).shutdown(); + verify(fallbackProvider, never()).newChannelBuilder(); + + // Transition back to enabled with CloudPathOnly. + pool.updateConfig(ENABLED_CLOUD_CONFIG); + verify(fallbackProvider, times(2)).newChannelBuilder(); + verify(fallbackChannel, never()).shutdown(); + + pool.close(); + verify(fallbackChannel, times(2)).shutdown(); + } +} From f718e1b263d4050359b8bc24cc7c44bfd0a3dca1 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 3 Nov 2026 13:47:41 -0500 Subject: [PATCH 2/2] test(bigtable): use separate ArgumentCaptor in SwitchingChannelPoolTest --- .../data/v2/internal/channels/SwitchingChannelPoolTest.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPoolTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPoolTest.java index 3b35349928d5..cc89ab1212b6 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPoolTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/SwitchingChannelPoolTest.java @@ -155,8 +155,9 @@ void startWithEmptyConfigUsesNoopPoolAndTransitionsToEnabledAndBack() { SessionStream disabledAgainStream = pool.newStream(FakeSessionGrpc.getOpenSessionMethod(), CallOptions.DEFAULT); disabledAgainStream.start(secondListener, new Metadata()); - verify(secondListener).onClose(statusCaptor.capture(), any(Metadata.class)); - assertThat(statusCaptor.getValue().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + ArgumentCaptor secondStatusCaptor = ArgumentCaptor.forClass(Status.class); + verify(secondListener).onClose(secondStatusCaptor.capture(), any(Metadata.class)); + assertThat(secondStatusCaptor.getValue().getCode()).isEqualTo(Status.Code.UNAVAILABLE); pool.close(); }