fix: feedback

Change-Id: Ib7f24524b6050ec11ca18e5dc624ca8513f99ec8

Author: nicholsl

google-cloud-bigtable/clirr-ignored-differences.xml

@@ -419,4 +419,11 @@
         <method>*create*</method>
         <to>*</to>
     </difference>
+    <!-- change method return type is ok because BigtableChannelPrimer is InternalApi -->
+    <difference>
+        <differenceType>7006</differenceType>
+        <className>com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer</className>
+        <method>*sendPrimeRequestsAsync*</method>
+        <to>com.google.api.core.ApiFuture</to>
+    </difference>
 </differences>

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java

@@ -64,11 +64,11 @@ public class BigtableChannelPool extends ManagedChannel {
   private final BigtableChannelPoolSettings settings;
   private final ChannelFactory channelFactory;
 
-  private ChannelPrimer channelPrimer;
+  private final ChannelPrimer channelPrimer;
   private final ScheduledExecutorService executor;
   private final Object entryWriteLock = new Object();
   @VisibleForTesting final AtomicReference<ImmutableList<Entry>> entries = new AtomicReference<>();
-  private ChannelPoolHealthChecker channelPoolHealthChecker;
+  private final ChannelPoolHealthChecker channelPoolHealthChecker;
   private final AtomicInteger indexTicker = new AtomicInteger();
   private final String authority;
 

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/ChannelPoolHealthChecker.java

@@ -21,6 +21,7 @@
 import com.google.cloud.bigtable.data.v2.stub.BigtableChannelPrimer;
 import com.google.cloud.bigtable.gaxx.grpc.BigtableChannelPool.Entry;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.time.Clock;
@@ -98,7 +99,9 @@ public ChannelPoolHealthChecker(
   }
 
   void start() {
-    if (channelPrimer instanceof BigtableChannelPrimer) {
+    if (!(channelPrimer instanceof BigtableChannelPrimer)) {
+      logger.log(Level.WARNING, "Provided channelPrimer not an instance of BigtableChannelPrimer, not checking channel health.");
+    } else {
       Duration initialDelayProbe =
           Duration.ofMillis(ThreadLocalRandom.current().nextLong(PROBE_INTERVAL.toMillis()));
       this.probeTaskScheduledFuture =
@@ -115,34 +118,30 @@ void start() {
               initialDelayDetect.toMillis(),
               PROBE_INTERVAL.toMillis(),
               TimeUnit.MILLISECONDS);
-    } else {
-      logger.log(Level.WARNING, "NoOpChannelPrimer was provided, not checking channel health.");
     }
   }
 
   /** Stop running health checking */
   public void stop() {
     if (probeTaskScheduledFuture != null) {
-      probeTaskScheduledFuture.cancel(true);
+      probeTaskScheduledFuture.cancel(false);
     }
     if (detectAndRemoveTaskScheduledFuture != null) {
-      detectAndRemoveTaskScheduledFuture.cancel(true);
+      detectAndRemoveTaskScheduledFuture.cancel(false);
     }
   }
 
   /** Runs probes on all the channels in the pool. */
   @VisibleForTesting
   void runProbes() {
+    Preconditions.checkState(channelPrimer instanceof BigtableChannelPrimer, "Health checking can only be enabled with BigtableChannelPrimer, found %s", channelPrimer);
     for (Entry entry : this.entrySupplier.get()) {
       final Instant startTime = clock.instant();
       final ApiFuture<PingAndWarmResponse> probeFuture;
 
-      if (channelPrimer instanceof BigtableChannelPrimer) {
-        BigtableChannelPrimer primer = (BigtableChannelPrimer) channelPrimer;
-        probeFuture = primer.sendPrimeRequestsAsync(entry.getManagedChannel());
-      } else {
-        continue;
-      }
+      BigtableChannelPrimer primer = (BigtableChannelPrimer) channelPrimer;
+      probeFuture = primer.sendPrimeRequestsAsync(entry.getManagedChannel());
+
       probeFuture.addListener(
           () -> onComplete(entry, startTime, probeFuture), MoreExecutors.directExecutor());
     }

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactoryTest.java

@@ -40,6 +40,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.io.BaseEncoding;
 import io.grpc.Attributes;
+import io.grpc.Grpc;
 import io.grpc.Metadata;
 import io.grpc.Server;
 import io.grpc.ServerCall;
@@ -50,9 +51,12 @@
 import io.grpc.stub.StreamObserver;
 import java.io.IOException;
 import java.lang.reflect.Method;
+import java.net.SocketAddress;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.LinkedBlockingDeque;
 import org.junit.After;
 import org.junit.Before;
@@ -87,6 +91,7 @@ public class BigtableDataClientFactoryTest {
   private final BlockingQueue<Attributes> setUpAttributes = new LinkedBlockingDeque<>();
   private final BlockingQueue<Attributes> terminateAttributes = new LinkedBlockingDeque<>();
   private final BlockingQueue<Metadata> requestMetadata = new LinkedBlockingDeque<>();
+  private final ConcurrentMap<SocketAddress, Boolean> warmedChannels = new ConcurrentHashMap<>();
 
   @Before
   public void setUp() throws IOException {
@@ -101,6 +106,15 @@ public <ReqT, RespT> Listener<ReqT> interceptCall(
                       Metadata headers,
                       ServerCallHandler<ReqT, RespT> next) {
                     requestMetadata.add(headers);
+
+                    // Check if the call is PingAndWarm and mark the channel address as warmed up.
+                    if (BigtableGrpc.getPingAndWarmMethod().equals(call.getMethodDescriptor())) {
+                      SocketAddress remoteAddr =
+                          call.getAttributes().get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR);
+                      if (remoteAddr != null) {
+                        warmedChannels.put(remoteAddr, true);
+                      }
+                    }
                     return next.startCall(call, headers);
                   }
                 })
@@ -249,7 +263,7 @@ public void testCreateForInstanceWithAppProfileHasCorrectSettings() throws Excep
   @Test
   public void testCreateWithRefreshingChannel() throws Exception {
     int poolSize = 3;
-    BigtableDataSettings.Builder builder =
+     BigtableDataSettings.Builder builder =
         BigtableDataSettings.newBuilderForEmulator(server.getPort())
             .setProjectId(DEFAULT_PROJECT_ID)
             .setInstanceId(DEFAULT_INSTANCE_ID)
@@ -278,21 +292,8 @@ public void testCreateWithRefreshingChannel() throws Exception {
     Mockito.verify(executorProvider, Mockito.times(1)).getExecutor();
     Mockito.verify(watchdogProvider, Mockito.times(1)).getWatchdog();
 
-    // Make sure that the clients are sharing the same ChannelPool
-    assertThat(setUpAttributes).hasSize(poolSize);
-
-    // Make sure that prime requests were sent only once per table per connection
-    assertThat(service.pingAndWarmRequests).hasSize(poolSize);
-    List<PingAndWarmRequest> expectedRequests = new LinkedList<>();
-    for (int i = 0; i < poolSize; i++) {
-      expectedRequests.add(
-          PingAndWarmRequest.newBuilder()
-              .setName(InstanceName.format(DEFAULT_PROJECT_ID, DEFAULT_INSTANCE_ID))
-              .setAppProfileId(DEFAULT_APP_PROFILE_ID)
-              .build());
-    }
-
-    assertThat(service.pingAndWarmRequests).containsExactly(expectedRequests.toArray());
+    assertThat(warmedChannels).hasSize(poolSize);
+    assertThat(warmedChannels.values()).doesNotContain(false);
 
     // Wait for all the connections to close asynchronously
     factory.close();