From 0e31fc41c5e6db624a1dd14a7ffc9b0ee56c953a Mon Sep 17 00:00:00 2001 From: Yuri Golobokov Date: Thu, 12 Nov 2026 03:50:48 +0000 Subject: [PATCH] fix(bigtable): count pre-establishment forceClose failures only --- .../channels/FallbackChannelPool.java | 71 +++++++----- .../channels/FallbackChannelPoolTest.java | 109 ++++++++++++++++++ 2 files changed, 152 insertions(+), 28 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/FallbackChannelPool.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/FallbackChannelPool.java index ea8cd22d1292..f4f2ca28d143 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/FallbackChannelPool.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/FallbackChannelPool.java @@ -31,13 +31,16 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; -import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; public class FallbackChannelPool implements ChannelPool { private volatile FallbackConfiguration config; + // Written under synchronized (this) in close(); volatile for lock-free reads in + // newStream callbacks. + private volatile boolean closed = false; private final ChannelPool primary; private final ChannelPool secondary; private final PoolFallbackListener poolFallbackListener; @@ -75,7 +78,7 @@ public void start() { @Override public void updateConfig(ChannelPoolConfiguration newConfig) { - if (!newConfig.hasDirectAccessWithFallback()) { + if (closed || !newConfig.hasDirectAccessWithFallback()) { return; } @@ -98,6 +101,9 @@ void updateConfig(FallbackConfiguration newConfig) { boolean triggerCallback = false; FallbackConfiguration localConfig; synchronized (this) { + if (closed) { + return; + } config = config.toBuilder() .setEnabled(newConfig.isEnabled()) @@ -120,9 +126,16 @@ void updateConfig(FallbackConfiguration newConfig) { } @Override - public synchronized void close() { - if (checkTask != null) { - checkTask.cancel(false); + public void close() { + synchronized (this) { + if (closed) { + return; + } + closed = true; + if (checkTask != null) { + checkTask.cancel(false); + checkTask = null; + } } primary.close(); secondary.close(); @@ -135,7 +148,7 @@ public SessionStream newStream( final SessionStream stream = current.newStream(desc, callOptions); return new ForwardingSessionStream(stream) { - private volatile boolean ignoreError; + private final AtomicBoolean outcomeRecorded = new AtomicBoolean(false); @Override public void start(Listener responseListener, Metadata headers) { @@ -144,12 +157,11 @@ public void start(Listener responseListener, Metadata headers) { new SessionStream.ForwardingListener(responseListener) { @Override public void onBeforeSessionStart(PeerInfo peerInfo) { - if (current == primary) { + if (current == primary && !closed && outcomeRecorded.compareAndSet(false, true)) { // Count successful session establishment as success. - successCount.incrementAndGet(); // peerInfo received successfully, session is established, // this is a successful start, thus ignoring any following errors. - ignoreError = true; + successCount.incrementAndGet(); } super.onBeforeSessionStart(peerInfo); } @@ -162,7 +174,7 @@ public void onClose(Status status, Metadata trailers) { return; } - if (!status.isOk() && !ignoreError) { + if (!closed && !status.isOk() && outcomeRecorded.compareAndSet(false, true)) { failureCount.incrementAndGet(); } super.onClose(status, trailers); @@ -170,24 +182,20 @@ public void onClose(Status status, Metadata trailers) { }, headers); } catch (Exception e) { - if (current == primary) { + if (current == primary && !closed && outcomeRecorded.compareAndSet(false, true)) { // Any exception on newStream start is a failure. failureCount.incrementAndGet(); } throw e; } } - - @Override - public void forceClose(@Nullable String message, @Nullable Throwable cause) { - // Force closing the session, thus ignoring the error caused by that. - ignoreError = true; - super.forceClose(message, cause); - } }; } private synchronized void scheduleCheckRates() { + if (closed) { + return; + } checkTask = exec.schedule( this::checkRatesAndReschedule, @@ -204,19 +212,26 @@ private void checkRates() { int successes = successCount.getAndSet(0); int failures = failureCount.getAndSet(0); int total = successes + failures; - FallbackConfiguration localConfig = config; + FallbackConfiguration localConfig; + boolean triggerCallback = false; - if (!localConfig.isEnabled()) { - return; + synchronized (this) { + if (closed) { + return; + } + localConfig = config; + if (localConfig.isEnabled() + && total > 0 + && ((double) failures) / total >= localConfig.getErrorRate()) { + triggerCallback = currentPool.compareAndSet(primary, secondary); + } } - if (total > 0 && ((double) failures) / total >= localConfig.getErrorRate()) { - if (currentPool.compareAndSet(primary, secondary)) { - poolFallbackListener.onFallback( - localConfig.getPrimaryName(), - localConfig.getFallbackName(), - ChannelFallbackReason.ERROR_RATE); - } + if (triggerCallback) { + poolFallbackListener.onFallback( + localConfig.getPrimaryName(), + localConfig.getFallbackName(), + ChannelFallbackReason.ERROR_RATE); } } } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/FallbackChannelPoolTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/FallbackChannelPoolTest.java index 08aa7b3afe1f..768de8253f2d 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/FallbackChannelPoolTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/FallbackChannelPoolTest.java @@ -18,8 +18,10 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; 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; @@ -32,6 +34,8 @@ 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.csm.tracers.PoolFallbackListener; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.PoolFallbackListener.ChannelFallbackReason; import io.grpc.CallOptions; import io.grpc.Metadata; import io.grpc.MethodDescriptor; @@ -42,6 +46,7 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -362,4 +367,108 @@ void testRescheduleAfterCheck() { checkTask.run(); verify(exec, times(3)).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); } + + @Test + void testPreEstablishmentForceCloseCountsAsFailure() { + PoolFallbackListener fallbackListener = mock(PoolFallbackListener.class); + pool = new FallbackChannelPool(config, fallbackListener, exec); + pool.start(); + + SessionStream stream = pool.newStream(methodDescriptor, CallOptions.DEFAULT); + stream.start(listener, new Metadata()); + + verify(primaryStream).start(listenerCaptor.capture(), any()); + Listener primaryListener = listenerCaptor.getValue(); + + // Force close before onBeforeSessionStart (e.g. SessionImpl.checkOpenTimeout()). + stream.forceClose("session open timeout", null); + verify(primaryStream).forceClose("session open timeout", null); + primaryListener.onClose(Status.CANCELLED, new Metadata()); + + checkTask.run(); + + verify(fallbackListener).onFallback("Primary", "Fallback", ChannelFallbackReason.ERROR_RATE); + pool.newStream(methodDescriptor, CallOptions.DEFAULT); + verify(secondary).newStream(any(), any()); + } + + @Test + void testStartExceptionFollowedByForceCloseDoesNotDoubleCount() { + pool.start(); + + // 2 successful sessions on primary + openSessionSuccessfully(primaryStream); + openSessionSuccessfully(primaryStream); + + // 3rd session throws on start(), followed by forceClose() and onClose(CANCELLED) + AtomicReference failingListener = new AtomicReference<>(); + doAnswer( + invocation -> { + failingListener.set(invocation.getArgument(0)); + throw new RuntimeException("boom"); + }) + .when(primaryStream) + .start(any(), any()); + + SessionStream failingStream = pool.newStream(methodDescriptor, CallOptions.DEFAULT); + try { + failingStream.start(listener, new Metadata()); + } catch (RuntimeException e) { + // expected + } + failingStream.forceClose("Session aborted due to uncaught exception", null); + failingListener.get().onClose(Status.CANCELLED, new Metadata()); + + // 2 successes, 1 failure (not 2). Error rate is 1/3 < 0.5, so should NOT fallback. + checkTask.run(); + + pool.newStream(methodDescriptor, CallOptions.DEFAULT); + verify(secondary, never()).newStream(any(), any()); + } + + @Test + void testConcurrentCloseDuringCheckRatesDoesNotReschedule() { + PoolFallbackListener fallbackListener = + (from, to, reason) -> { + // Simulate close() occurring while checkRatesAndReschedule() is in flight before + // scheduleCheckRates(). + pool.close(); + }; + pool = new FallbackChannelPool(config, fallbackListener, exec); + pool.start(); + verify(exec, times(1)).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + + openSessionWithError(primaryStream, Status.UNAVAILABLE); + + checkTask.run(); + + // Verify schedule() was NOT called a second time after pool.close() + verify(exec, times(1)).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + } + + @Test + void testCloseIsIdempotentAndSuppressesPostCloseCallbacks() { + PoolFallbackListener fallbackListener = mock(PoolFallbackListener.class); + pool = new FallbackChannelPool(config, fallbackListener, exec); + pool.start(); + + Listener inFlightListener = openSessionGetListener(primaryStream); + + pool.close(); + pool.close(); + + // Verify close() is idempotent + verify(primary, times(1)).close(); + verify(secondary, times(1)).close(); + verify(future, times(1)).cancel(false); + + // Post-close stream failure, checkTask execution, start(), and updateConfig() should be no-ops + inFlightListener.onClose(Status.UNAVAILABLE, new Metadata()); + checkTask.run(); + pool.start(); + pool.updateConfig(config.toBuilder().setEnabled(false).build()); + + verify(fallbackListener, never()).onFallback(any(), any(), any()); + verify(exec, times(1)).schedule(any(Runnable.class), anyLong(), any(TimeUnit.class)); + } }