fix(grpc-gcp): correct channel lifecycle bookkeeping (#14196)

Use a copy-on-write test metric recorder because synchronized-list
iteration raced with asynchronous exports.

Author: rahul2393

grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpClientCall.java

@@ -125,7 +125,10 @@ public void sendMessage(ReqT message) {
         delegateChannelRef.activeStreamsCountIncr();
 
         // Create the client call and do the previous operations.
-        delegateCall = delegateChannelRef.getChannel().newCall(methodDescriptor, callOptions);
+        CallOptions callOptionsWithChannelId =
+            callOptions.withOption(GcpManagedChannel.CHANNEL_ID_KEY, delegateChannelRef.getId());
+        delegateCall =
+            delegateChannelRef.getChannel().newCall(methodDescriptor, callOptionsWithChannelId);
         for (Runnable call : calls) {
           call.run();
         }

grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java

@@ -269,6 +269,16 @@ private static int stateFromChannelId(int channelId) {
   // Clock supplier for nanoTime, injectable for testing.
   private Supplier<Long> nanoClock = System::nanoTime;
 
+  @VisibleForTesting
+  Map<Integer, Map<String, Integer>> fallbackMapForTest() {
+    return fallbackMap;
+  }
+
+  @VisibleForTesting
+  int readyChannelCountForTest() {
+    return readyChannels.get();
+  }
+
   @VisibleForTesting
   void setNanoClock(Supplier<Long> nanoClock) {
     this.nanoClock = nanoClock;
@@ -408,10 +418,7 @@ private void removeOldestChannels(int num) {
 
     for (ChannelRef channelRef : channelsToRemove) {
       channelRef.resetAffinityCount();
-      channelRef.deactivate();
-      if (channelRef.getState() == ConnectivityState.READY) {
-        decReadyChannels(false);
-      }
+      channelRef.deactivateAndAccountReadiness();
     }
 
     // Remove affinity keys mapping for the channels.
@@ -1482,6 +1489,9 @@ private class ChannelStateMonitor implements Runnable {
     private long connectingStartNanos;
     private long connectedSinceNanos;
 
+    @GuardedBy("channelRef")
+    private boolean readyAccounted;
+
     private ChannelStateMonitor(ManagedChannel channel, ChannelRef channelRef) {
       this.channelRef = channelRef;
       this.channel = channel;
@@ -1496,15 +1506,26 @@ public ConnectivityState getCurrentState() {
       return currentState;
     }
 
+    private void accountReadyIfNeeded() {
+      if (currentState == ConnectivityState.READY && !readyAccounted) {
+        readyAccounted = true;
+        incReadyChannels(false);
+      }
+    }
+
+    private void unaccountReadyIfNeeded() {
+      if (readyAccounted) {
+        readyAccounted = false;
+        decReadyChannels(false);
+      }
+    }
+
     @Override
     public void run() {
       if (channel == null) {
         return;
       }
 
-      // Is the channel in the pool?
-      boolean isActive = channelRefs.contains(this.channelRef);
-
       // Keep minSize channels always connected.
       boolean requestConnection =
           channelRefs.size() < minSize
@@ -1515,35 +1536,39 @@ public void run() {
                   .anyMatch(id -> (id == channelRef.getId()));
 
       ConnectivityState newState = channel.getState(requestConnection);
-      if (logger.isLoggable(Level.FINER)) {
-        logger.finer(
-            log(
-                "Channel %d state change detected: %s -> %s",
-                channelRef.getId(), currentState, newState));
-      }
-      if (newState == ConnectivityState.READY && currentState != ConnectivityState.READY) {
-        connectedSinceNanos = System.nanoTime();
-        if (isActive) {
-          incReadyChannels(true);
-          if (connectingStartNanos > 0) {
-            saveReadinessTime(System.nanoTime() - connectingStartNanos);
+      boolean isActive;
+      synchronized (channelRef) {
+        isActive = channelRef.isActive() && channelRefs.contains(channelRef);
+        if (logger.isLoggable(Level.FINER)) {
+          logger.finer(
+              log(
+                  "Channel %d state change detected: %s -> %s",
+                  channelRef.getId(), currentState, newState));
+        }
+        if (newState == ConnectivityState.READY && currentState != ConnectivityState.READY) {
+          connectedSinceNanos = System.nanoTime();
+          if (isActive && !readyAccounted) {
+            readyAccounted = true;
+            incReadyChannels(true);
+            if (connectingStartNanos > 0) {
+              saveReadinessTime(System.nanoTime() - connectingStartNanos);
+            }
           }
+          connectingStartNanos = 0;
         }
-        connectingStartNanos = 0;
-      }
-      if (isActive
-          && newState != ConnectivityState.READY
-          && currentState == ConnectivityState.READY) {
-        decReadyChannels(true);
-      }
-      if (newState == ConnectivityState.CONNECTING
-          && currentState != ConnectivityState.CONNECTING) {
-        connectingStartNanos = System.nanoTime();
-      }
-      if (newState != ConnectivityState.READY) {
-        connectedSinceNanos = 0;
+        if (newState != ConnectivityState.READY && readyAccounted) {
+          readyAccounted = false;
+          decReadyChannels(true);
+        }
+        if (newState == ConnectivityState.CONNECTING
+            && currentState != ConnectivityState.CONNECTING) {
+          connectingStartNanos = System.nanoTime();
+        }
+        if (newState != ConnectivityState.READY) {
+          connectedSinceNanos = 0;
+        }
+        currentState = newState;
       }
-      currentState = newState;
 
       processChannelStateChange(channelRef.getId(), newState);
       if (isActive) {
@@ -1686,8 +1711,12 @@ protected ChannelRef getChannelRef(@Nullable String key) {
       if (logger.isLoggable(Level.FINEST)) {
         logger.finest(log("Using fallback channel: %d -> %d", mappedChannel.getId(), channelId));
       }
-      fallbacksSucceeded.incrementAndGet();
-      return channelRefs.get(channelId);
+      ChannelRef fallbackChannel = channelIdToChannelRef.get(channelId);
+      if (fallbackChannel != null && fallbackChannel.isActive()) {
+        fallbacksSucceeded.incrementAndGet();
+        return fallbackChannel;
+      }
+      tempMap.remove(key, channelId);
     }
     // No temp mapping for this key or fallback channel is also broken.
     ChannelRef channelRef = pickLeastBusyChannel(/* forFallback= */ true);
@@ -1710,7 +1739,10 @@ protected ChannelRef getChannelRef(@Nullable String key) {
     fallbacksFailed.incrementAndGet();
     if (channelId != null) {
       // Stick with previous mapping if fallback has failed.
-      return channelRefs.get(channelId);
+      ChannelRef fallbackChannel = channelIdToChannelRef.get(channelId);
+      if (fallbackChannel != null && fallbackChannel.isActive()) {
+        return fallbackChannel;
+      }
     }
     return mappedChannel;
   }
@@ -1779,16 +1811,16 @@ ChannelRef createNewChannel() {
       channelRefs.add(chRef);
       removedChannelRefs.remove(chRef);
       channelIdToChannelRef.put(chRef.getId(), chRef);
-      chRef.activate();
+      chRef.activateAndAccountReadiness();
       logger.finer(log("Channel %d reused.", chRef.getId()));
-      incReadyChannels(false);
       maxChannels.accumulateAndGet(getNumberOfChannels(), Math::max);
       return chRef;
     }
 
     ChannelRef channelRef = new ChannelRef(delegateChannelBuilder.build());
     channelRefs.add(channelRef);
     channelIdToChannelRef.put(channelRef.getId(), channelRef);
+    channelRef.activateAndAccountReadiness();
     logger.finer(log("Channel %d created.", channelRef.getId()));
     maxChannels.accumulateAndGet(getNumberOfChannels(), Math::max);
     return channelRef;
@@ -2447,12 +2479,27 @@ protected boolean isActive() {
       return active;
     }
 
-    private void activate() {
-      active = true;
+    private void activateAndAccountReadiness() {
+      synchronized (this) {
+        active = true;
+        channelStateMonitor.accountReadyIfNeeded();
+      }
+    }
+
+    private void deactivateAndAccountReadiness() {
+      synchronized (this) {
+        channelStateMonitor.unaccountReadyIfNeeded();
+        active = false;
+      }
     }
 
     private void deactivate() {
-      active = false;
+      deactivateAndAccountReadiness();
+    }
+
+    @VisibleForTesting
+    void deactivateForTest() {
+      deactivateAndAccountReadiness();
     }
 
     protected void affinityCountIncr() {

grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelTest.java

@@ -364,6 +364,72 @@ public void testChannelAffinityRefInvalidChannelIdPicksAvailableChannel() throws
     }
   }
 
+  @Test
+  public void fallbackUsesChannelIdMapAfterPoolHasIndexGap() {
+    resetGcpChannel();
+    ExecutorService executorService = Executors.newSingleThreadExecutor();
+    try {
+      List<FakeManagedChannel> channels = new ArrayList<>();
+      for (int i = 0; i < 3; i++) {
+        FakeManagedChannel channel = new FakeManagedChannel(executorService);
+        channel.setState(ConnectivityState.READY);
+        channels.add(channel);
+      }
+      gcpChannel =
+          (GcpManagedChannel)
+              GcpManagedChannelBuilder.forDelegateBuilder(new FakeManagedChannelBuilder(channels))
+                  .withOptions(
+                      GcpManagedChannelOptions.newBuilder()
+                          .withChannelPoolOptions(
+                              GcpChannelPoolOptions.newBuilder()
+                                  .setMinSize(3)
+                                  .setMaxSize(3)
+                                  .build())
+                          .withResiliencyOptions(
+                              GcpResiliencyOptions.newBuilder().setNotReadyFallback(true).build())
+                          .build())
+                  .build();
+      ChannelRef removed = gcpChannel.channelRefs.get(0);
+      ChannelRef mapped = gcpChannel.channelRefs.get(1);
+      ChannelRef fallback = gcpChannel.channelRefs.get(2);
+      String key = "session";
+      gcpChannel.bind(mapped, Collections.singletonList(key));
+      gcpChannel.processChannelStateChange(mapped.getId(), ConnectivityState.TRANSIENT_FAILURE);
+      gcpChannel.fallbackMapForTest().get(mapped.getId()).put(key, fallback.getId());
+      gcpChannel.channelRefs.remove(removed);
+
+      assertThat(gcpChannel.getChannelRef(key)).isSameInstanceAs(fallback);
+    } finally {
+      gcpChannel.shutdownNow();
+      executorService.shutdownNow();
+    }
+  }
+
+  @Test
+  public void readyAccountingRemainsExactWhenReadyChannelIsReused() {
+    resetGcpChannel();
+    ExecutorService executorService = Executors.newSingleThreadExecutor();
+    try {
+      gcpChannel = createPoolWithFakeReadyChannels(executorService, 2);
+      assertThat(gcpChannel.readyChannelCountForTest()).isEqualTo(2);
+
+      ChannelRef reused = gcpChannel.channelRefs.get(0);
+      gcpChannel.channelRefs.remove(reused);
+      reused.deactivateForTest();
+      gcpChannel.removedChannelRefs.add(reused);
+      assertThat(gcpChannel.readyChannelCountForTest()).isEqualTo(1);
+
+      assertThat(gcpChannel.createNewChannel()).isSameInstanceAs(reused);
+      assertThat(gcpChannel.readyChannelCountForTest()).isEqualTo(2);
+      reused.deactivateForTest();
+      reused.deactivateForTest();
+      assertThat(gcpChannel.readyChannelCountForTest()).isEqualTo(1);
+    } finally {
+      gcpChannel.shutdownNow();
+      executorService.shutdownNow();
+    }
+  }
+
   @Test
   public void testChannelAffinityRefRemovedChannelPicksAvailableChannel() throws Exception {
     resetGcpChannel();

grpc-gcp-java/src/test/java/com/google/cloud/grpc/fallback/GcpFallbackChannelTest.java

@@ -75,11 +75,10 @@
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.EnumSet;
 import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
@@ -352,7 +351,7 @@ private void simulateCanceledCall(boolean expectFallbackRouting) {
   }
 
   private class TestMetricExporter implements MetricExporter {
-    public final List<MetricData> exportedMetrics = Collections.synchronizedList(new ArrayList<>());
+    public final List<MetricData> exportedMetrics = new CopyOnWriteArrayList<>();
 
     @Override
     public CompletableResultCode export(@Nonnull Collection<MetricData> metrics) {

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RequestIdMockServerTest.java

@@ -680,9 +680,9 @@ public void testOtherClientId() {
         ImmutableList.of(
             XGoogSpannerRequestId.of(getClientId(), -1, 1, 1),
             // The CreateSession RPC from the initialization of the second client is included in
-            // the requests that we see. This request does not include a channel hint, hence the
-            // zero value for the channel number in the request ID.
-            XGoogSpannerRequestId.of(otherClientId, 0, 1, 1),
+            // the requests that we see. Its request ID carries the channel that grpc-gcp selected
+            // for the affinity-key path.
+            XGoogSpannerRequestId.of(otherClientId, -1, 1, 1),
             XGoogSpannerRequestId.of(otherClientId, -1, 2, 1),
             XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)),
         actual);