Unit test.

Author: kannanjgithub

rls/src/main/java/io/grpc/rls/CachingRlsLbClient.java

@@ -134,6 +134,7 @@ final class CachingRlsLbClient {
   private final ChannelLogger logger;
   private final ChildPolicyWrapper fallbackChildPolicyWrapper;
   private ConnectivityState lastRlsServerConnectivityState;
+  private boolean wasReady;
   private boolean wasInTransientFailure;
 
   static {
@@ -262,7 +263,11 @@ public void accept(BatchRecorder recorder) {
 
   private void rlsServerConnectionStateChanged() {
     ConnectivityState currentState = rlsChannel.getState(false);
-    if (wasInTransientFailure && currentState == ConnectivityState.READY) {
+    if (!wasInTransientFailure && currentState == ConnectivityState.READY) {
+      wasReady = true;
+    } else if (wasReady && currentState == ConnectivityState.TRANSIENT_FAILURE) {
+      wasInTransientFailure = true;
+    } else if (wasInTransientFailure && currentState == ConnectivityState.READY) {
       wasInTransientFailure = false;
       synchronized (lock) {
         for (CacheEntry value : linkedHashLruCache.values()) {
@@ -271,8 +276,6 @@ private void rlsServerConnectionStateChanged() {
           }
         }
       }
-    } else if (lastRlsServerConnectivityState == ConnectivityState.READY && currentState == ConnectivityState.TRANSIENT_FAILURE) {
-      wasInTransientFailure = true;
     }
     rlsChannel.notifyWhenStateChanged(currentState, () -> rlsServerConnectionStateChanged());
     lastRlsServerConnectivityState = currentState;

rls/src/test/java/io/grpc/rls/CachingRlsLbClientTest.java

@@ -68,7 +68,10 @@
 import io.grpc.inprocess.InProcessServerBuilder;
 import io.grpc.internal.BackoffPolicy;
 import io.grpc.internal.FakeClock;
+import io.grpc.internal.GrpcUtil;
+import io.grpc.internal.ObjectPool;
 import io.grpc.internal.PickSubchannelArgsImpl;
+import io.grpc.internal.SharedResourcePool;
 import io.grpc.lookup.v1.RouteLookupServiceGrpc;
 import io.grpc.rls.CachingRlsLbClient.CacheEntry;
 import io.grpc.rls.CachingRlsLbClient.CachedRouteLookupResponse;
@@ -98,10 +101,12 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import javax.annotation.Nonnull;
 import org.junit.After;
 import org.junit.Before;
@@ -341,28 +346,80 @@ public void get_throttledAndRecover() throws Exception {
   @Test
   public void controlPlaneTransientToReady_backOffEntriesRemovedAndPickerUpdated() throws Exception {
     setUpRlsLbClient();
+    final ConnectivityState[] rlsChannelState = new ConnectivityState[1];
+    Runnable channelStateListener = new Runnable() {
+      @Override
+      public void run() {
+        rlsChannelState[0] = fakeHelper.oobChannel.getState(false);
+        fakeHelper.oobChannel.notifyWhenStateChanged(rlsChannelState[0], this);
+        synchronized(this) {
+          notify();
+        }
+      }
+    };
+    fakeHelper.oobChannel.notifyWhenStateChanged(fakeHelper.oobChannel.getState(false),
+        channelStateListener);
     RouteLookupRequest routeLookupRequest = RouteLookupRequest.create(ImmutableMap.of(
         "server", "bigtable.googleapis.com", "service-key", "foo", "method-key", "bar"));
     rlsServerImpl.setLookupTable(
         ImmutableMap.of(
             routeLookupRequest,
             RouteLookupResponse.create(ImmutableList.of("target"), "header")));
 
+    CachedRouteLookupResponse resp = getInSyncContext(routeLookupRequest);
+    assertThat(resp.isPending()).isTrue();
+    // server response
+    fakeClock.forwardTime(SERVER_LATENCY_MILLIS, TimeUnit.MILLISECONDS);
+    resp = getInSyncContext(routeLookupRequest);
+    assertThat(resp.hasData()).isTrue();
+
     fakeHelper.server.shutdown();
+    // Channel goes to IDLE state from the shutdown listener handling.
+    try {
+      if (!fakeHelper.server.awaitTermination(10, TimeUnit.SECONDS)) {
+        fakeHelper.server.shutdownNow(); // Forceful shutdown if graceful timeout expires
+      }
+    } catch (InterruptedException e) {
+      fakeHelper.server.shutdownNow();
+    }
+    // Use a different key to cause a cache miss and trigger a RPC.
+    RouteLookupRequest routeLookupRequest2 = RouteLookupRequest.create(ImmutableMap.of(
+        "server", "bigtable.googleapis.com", "service-key", "foo2", "method-key", "bar"));
+    // Rls channel will go to TRANSIENT_FAILURE (back-off) because the picker notices the
+    // subchannel state IDLE and the server transport listener is null.
+    resp = getInSyncContext(routeLookupRequest2);
+    assertThat(resp.isPending()).isTrue();
+    assertThat(rlsChannelState[0]).isEqualTo(ConnectivityState.TRANSIENT_FAILURE);
+    // Throttle the next rpc call.
     fakeThrottler.nextResult = true;
     fakeBackoffProvider.nextPolicy = createBackoffPolicy(10, TimeUnit.MILLISECONDS);
 
-    CachedRouteLookupResponse resp = getInSyncContext(routeLookupRequest);
+    // Cause another cache miss by using a new request key. This will create a back-off Rls
+    // cache entry.
+    RouteLookupRequest routeLookupRequest3 = RouteLookupRequest.create(ImmutableMap.of(
+        "server", "bigtable.googleapis.com", "service-key", "foo3", "method-key", "bar"));
+    resp = getInSyncContext(routeLookupRequest3);
 
     assertThat(resp.hasError()).isTrue();
 
-    fakeHelper.server.start();
-    Thread.sleep(10000);
-    // Backoff entry evicted from cache.
-    verify(evictionListener)
-        .onEviction(eq(routeLookupRequest), any(CacheEntry.class), eq(EvictionType.EXPLICIT));
-    // Assert that Rls LB policy picker was updated.
-    assertThat(fakeHelper.lastPicker.toString()).isEqualTo("RlsPicker{target=service1}");
+    fakeHelper.createServerAndRegister("service1");
+    // Wait for Rls subchannel back-off expiry and its moving to READY
+    synchronized(channelStateListener) {
+      channelStateListener.wait(5000);
+    }
+    assertThat(rlsChannelState[0]).isEqualTo(ConnectivityState.READY);
+    final ObjectPool<? extends Executor> defaultExecutorPool =
+        SharedResourcePool.forResource(GrpcUtil.SHARED_CHANNEL_EXECUTOR);
+    AtomicBoolean isSuccess = new AtomicBoolean(false);
+    defaultExecutorPool.getObject().execute(() -> {
+      // Backoff entry evicted from cache.
+      verify(evictionListener)
+          .onEviction(eq(routeLookupRequest3), any(CacheEntry.class), eq(EvictionType.EXPLICIT));
+      // Assert that Rls LB policy picker was updated.
+      assertThat(fakeHelper.lastPicker.toString()).isEqualTo("RlsPicker{target=service1}");
+      isSuccess.set(true);
+    });
+    // No need to wait because direct executor is used.
   }
 
   @Test
@@ -912,16 +969,21 @@ private final class FakeHelper extends Helper {
     SubchannelPicker lastPicker;
     Server server;
     ManagedChannel oobChannel;
+
+    void createServerAndRegister(String target) throws IOException {
+      server = InProcessServerBuilder.forName(target)
+          .addService(rlsServerImpl)
+          .directExecutor()
+          .build()
+          .start();
+      grpcCleanupRule.register(server);
+    }
+
     @Override
     public ManagedChannelBuilder<?> createResolvingOobChannelBuilder(
         String target, ChannelCredentials creds) {
        try {
-         server = InProcessServerBuilder.forName(target)
-             .addService(rlsServerImpl)
-             .directExecutor()
-             .build()
-             .start();
-        grpcCleanupRule.register(server);
+         createServerAndRegister(target);
       } catch (IOException e) {
         throw new RuntimeException("cannot create server: " + target, e);
       }
@@ -937,7 +999,8 @@ protected ManagedChannelBuilder<?> delegate() {
 
         @Override
         public ManagedChannel build() {
-          return grpcCleanupRule.register(super.build());
+          oobChannel = super.build();
+          return grpcCleanupRule.register(oobChannel);
         }
 
         @Override
@@ -955,9 +1018,7 @@ public CleaningChannelBuilder overrideAuthority(String authority) {
         }
       }
 
-      CleaningChannelBuilder oobChannelBuilder = new CleaningChannelBuilder();
-      oobChannel = oobChannelBuilder.build();
-      return oobChannelBuilder;
+      return new CleaningChannelBuilder();
     }
 
     @Override