fix(bigtable): count pre-establishment forceClose failures only

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,32 +174,28 @@ 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);
                 }
               },
               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);
     }
   }
 }

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<Listener> 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));
+  }
 }