fix: pr feedback

Change-Id: Ic5f18450c5f60bcf5473e6890b41b006821ddf9d

Author: nicholsl

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.bigtable.data.v2.stub;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.core.InternalApi;
 import com.google.api.core.SettableApiFuture;
 import com.google.auth.Credentials;
@@ -110,8 +111,7 @@ private void sendPrimeRequestsBlocking(ManagedChannel managedChannel) {
     }
   }
 
-  public SettableApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(
-      ManagedChannel managedChannel) {
+  public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel managedChannel) {
     ClientCall<PingAndWarmRequest, PingAndWarmResponse> clientCall =
         managedChannel.newCall(
             BigtableGrpc.getPingAndWarmMethod(),

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/NoOpChannelPrimer.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.bigtable.data.v2.stub;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.core.InternalApi;
 import com.google.api.core.SettableApiFuture;
 import com.google.bigtable.v2.PingAndWarmResponse;
@@ -35,7 +36,7 @@ public void primeChannel(ManagedChannel channel) {
   }
 
   @Override
-  public SettableApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel channel) {
+  public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel channel) {
     SettableApiFuture future = SettableApiFuture.create();
     future.set(PingAndWarmResponse.getDefaultInstance());
     return future;

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

@@ -100,7 +100,7 @@ public static BigtableChannelPool create(
     this.channelPrimer = channelPrimer;
     Clock systemClock = Clock.systemUTC();
     this.channelPoolHealthChecker =
-        new ChannelPoolHealthChecker(() -> entries.get(), channelPrimer, executor, systemClock);
+        new ChannelPoolHealthChecker(entries::get, channelPrimer, executor, systemClock);
     this.channelPoolHealthChecker.start();
 
     ImmutableList.Builder<Entry> initialListBuilder = ImmutableList.builder();

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

@@ -15,28 +15,34 @@
  */
 package com.google.cloud.bigtable.gaxx.grpc;
 
-import com.google.api.core.SettableApiFuture;
+import com.google.api.core.ApiFuture;
 import com.google.auto.value.AutoValue;
 import com.google.bigtable.v2.PingAndWarmResponse;
 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.collect.ImmutableList;
+import com.google.common.util.concurrent.MoreExecutors;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.Comparator;
 import java.util.List;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /** Class that manages the health checking in the BigtableChannelPool */
 class ChannelPoolHealthChecker {
 
+  private static final Logger logger = Logger.getLogger(ChannelPoolHealthChecker.class.getName());
+
   // Configuration constants
   // Window_Duration is the duration over which we keep probe results
   private static final Duration WINDOW_DURATION = Duration.ofMinutes(5);
@@ -69,9 +75,12 @@ static ProbeResult create(Instant startTime, boolean success) {
 
   private final Supplier<ImmutableList<Entry>> entrySupplier;
   private volatile Instant lastEviction;
-  private ScheduledExecutorService executor;
+  private final ScheduledExecutorService executor;
+
+  private final ChannelPrimer channelPrimer;
 
-  private ChannelPrimer channelPrimer;
+  private ScheduledFuture<?> probeTaskScheduledFuture;
+  private ScheduledFuture<?> detectAndRemoveTaskScheduledFuture;
 
   private final Clock clock;
 
@@ -89,54 +98,66 @@ public ChannelPoolHealthChecker(
   }
 
   void start() {
-    Duration initialDelayProbe =
-        Duration.ofMillis(ThreadLocalRandom.current().nextLong(PROBE_INTERVAL.toMillis()));
-    executor.scheduleAtFixedRate(
-        this::runProbes,
-        initialDelayProbe.toMillis(),
-        PROBE_INTERVAL.toMillis(),
-        TimeUnit.MILLISECONDS);
-    Duration initialDelayDetect =
-        Duration.ofMillis(ThreadLocalRandom.current().nextLong(PROBE_INTERVAL.toMillis()));
-    executor.scheduleAtFixedRate(
-        this::detectAndRemoveOutlierEntries,
-        initialDelayDetect.toMillis(),
-        PROBE_INTERVAL.toMillis(),
-        TimeUnit.MILLISECONDS);
+    if (channelPrimer instanceof BigtableChannelPrimer) {
+      Duration initialDelayProbe =
+          Duration.ofMillis(ThreadLocalRandom.current().nextLong(PROBE_INTERVAL.toMillis()));
+      this.probeTaskScheduledFuture =
+          executor.scheduleAtFixedRate(
+              this::runProbes,
+              initialDelayProbe.toMillis(),
+              PROBE_INTERVAL.toMillis(),
+              TimeUnit.MILLISECONDS);
+      Duration initialDelayDetect =
+          Duration.ofMillis(ThreadLocalRandom.current().nextLong(PROBE_INTERVAL.toMillis()));
+      this.detectAndRemoveTaskScheduledFuture =
+          executor.scheduleAtFixedRate(
+              this::detectAndRemoveOutlierEntries,
+              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() {
-    executor.shutdownNow();
+    if (probeTaskScheduledFuture != null) {
+      probeTaskScheduledFuture.cancel(true);
+    }
+    if (detectAndRemoveTaskScheduledFuture != null) {
+      detectAndRemoveTaskScheduledFuture.cancel(true);
+    }
   }
 
   /** Runs probes on all the channels in the pool. */
   @VisibleForTesting
   void runProbes() {
     for (Entry entry : this.entrySupplier.get()) {
       final Instant startTime = clock.instant();
-      final SettableApiFuture<PingAndWarmResponse> probeFuture;
+      final ApiFuture<PingAndWarmResponse> probeFuture;
 
       if (channelPrimer instanceof BigtableChannelPrimer) {
         BigtableChannelPrimer primer = (BigtableChannelPrimer) channelPrimer;
         probeFuture = primer.sendPrimeRequestsAsync(entry.getManagedChannel());
       } else {
         continue;
       }
-      probeFuture.addListener(() -> onComplete(entry, startTime, probeFuture), executor);
+      probeFuture.addListener(
+          () -> onComplete(entry, startTime, probeFuture), MoreExecutors.directExecutor());
     }
   }
 
   /** Callback that will update Entry data on probe complete. */
   @VisibleForTesting
-  void onComplete(
-      Entry entry, Instant startTime, SettableApiFuture<PingAndWarmResponse> probeFuture) {
+  void onComplete(Entry entry, Instant startTime, ApiFuture<PingAndWarmResponse> probeFuture) {
     boolean success;
     try {
       probeFuture.get(PROBE_DEADLINE.toMillis(), TimeUnit.MILLISECONDS);
       success = true;
     } catch (Exception e) {
       success = false;
+      logger.log(Level.WARNING, "Probe failed");
     }
     addProbeResult(entry, ProbeResult.create(startTime, success));
   }
@@ -219,6 +240,9 @@ void detectAndRemoveOutlierEntries() {
     Entry outlier = findOutlierEntry();
     if (outlier != null) {
       this.lastEviction = clock.instant();
+      outlier.failedProbesInWindow.set(0);
+      outlier.successfulProbesInWindow.set(0);
+      outlier.probeHistory.clear();
       outlier.getManagedChannel().enterIdle();
     }
   }

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

@@ -15,14 +15,14 @@
  */
 package com.google.cloud.bigtable.gaxx.grpc;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.core.InternalApi;
-import com.google.api.core.SettableApiFuture;
 import com.google.bigtable.v2.PingAndWarmResponse;
 import io.grpc.ManagedChannel;
 
 @InternalApi("For internal use by google-cloud-java clients only")
 public interface ChannelPrimer {
   void primeChannel(ManagedChannel var1);
 
-  SettableApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel var1);
+  ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel var1);
 }

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

@@ -19,7 +19,7 @@
 import static org.junit.Assert.assertThrows;
 
 import com.google.api.core.ApiFunction;
-import com.google.api.core.SettableApiFuture;
+import com.google.api.core.ApiFuture;
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.bigtable.v2.BigtableGrpc.BigtableImplBase;
@@ -173,7 +173,7 @@ public void testHeadersAreSent() {
   // New test for the async success path
   @Test
   public void testAsyncSuccess() throws Exception {
-    SettableApiFuture<PingAndWarmResponse> future = primer.sendPrimeRequestsAsync(channel);
+    ApiFuture<PingAndWarmResponse> future = primer.sendPrimeRequestsAsync(channel);
 
     PingAndWarmResponse response = future.get(1, TimeUnit.SECONDS);
     assertThat(response).isNotNull();
@@ -192,7 +192,7 @@ public PingAndWarmResponse apply(PingAndWarmRequest pingAndWarmRequest) {
           }
         };
 
-    SettableApiFuture<PingAndWarmResponse> future = primer.sendPrimeRequestsAsync(channel);
+    ApiFuture<PingAndWarmResponse> future = primer.sendPrimeRequestsAsync(channel);
 
     ExecutionException e =
         assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS));

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/HealthChecker.java

@@ -15,6 +15,9 @@
  */
 package com.google.cloud.bigtable.gaxx.grpc;
 
+import com.google.api.core.InternalApi;
+
+@InternalApi
 public interface HealthChecker {
   void start();