PR feedback 4

Author: jinseopkim0

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/BigtableTableAdminClientV2.java

@@ -94,7 +94,9 @@ protected BigtableTableAdminClientV2(BaseBigtableTableAdminSettings settings) th
 
     this.awaitConsistencyCallable =
         createAwaitConsistencyCallable(
-            (BigtableTableAdminStubSettings) settings.getStubSettings(), lightweightContext);
+            (BigtableTableAdminStubSettings) settings.getStubSettings(),
+            settings.getStubSettings().getClock(),
+            this.backgroundExecutor);
     this.optimizeRestoredTableOperationBaseCallable =
         createOptimizeRestoredTableOperationBaseCallable(lightweightContext);
   }
@@ -119,8 +121,9 @@ protected BigtableTableAdminClientV2(BigtableTableAdminStub stub) {
   }
 
   private AwaitConsistencyCallable createAwaitConsistencyCallable(
-      BigtableTableAdminStubSettings settings, com.google.api.gax.rpc.ClientContext clientContext)
-      throws IOException {
+      BigtableTableAdminStubSettings settings,
+      com.google.api.core.ApiClock clock,
+      java.util.concurrent.ScheduledExecutorService executor) {
     // TODO(igorbernstein2): expose polling settings
     RetrySettings pollingSettings =
         AWAIT_CONSISTENCY_POLLING_SETTINGS_BASE.toBuilder()
@@ -131,7 +134,8 @@ private AwaitConsistencyCallable createAwaitConsistencyCallable(
     return AwaitConsistencyCallable.create(
         getStub().generateConsistencyTokenCallable(),
         getStub().checkConsistencyCallable(),
-        clientContext,
+        clock,
+        executor,
         pollingSettings);
   }
 

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/stub/AwaitConsistencyCallable.java

@@ -16,9 +16,11 @@
 package com.google.cloud.bigtable.admin.v2.stub;
 
 import com.google.api.core.ApiAsyncFunction;
+import com.google.api.core.ApiClock;
 import com.google.api.core.ApiFunction;
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
+import com.google.api.core.InternalApi;
 import com.google.api.gax.retrying.ExponentialPollAlgorithm;
 import com.google.api.gax.retrying.NonCancellableFuture;
 import com.google.api.gax.retrying.ResultRetryAlgorithmWithContext;
@@ -30,7 +32,6 @@
 import com.google.api.gax.retrying.ScheduledRetryingExecutor;
 import com.google.api.gax.retrying.TimedAttemptSettings;
 import com.google.api.gax.rpc.ApiCallContext;
-import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.bigtable.admin.v2.CheckConsistencyRequest;
 import com.google.bigtable.admin.v2.CheckConsistencyResponse;
@@ -42,6 +43,7 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CancellationException;
+import java.util.concurrent.ScheduledExecutorService;
 import javax.annotation.Nullable;
 
 /**
@@ -51,6 +53,7 @@
  * <p>This callable wraps GenerateConsistencyToken and CheckConsistency RPCs. It will generate a
  * token then poll until isConsistent is true.
  */
+@InternalApi
 public class AwaitConsistencyCallable extends UnaryCallable<ConsistencyRequest, Void> {
   private final UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
       generateCallable;
@@ -59,33 +62,36 @@ public class AwaitConsistencyCallable extends UnaryCallable<ConsistencyRequest,
 
   @Nullable private final TableAdminRequestContext requestContext;
 
+  @InternalApi
   public static AwaitConsistencyCallable create(
       UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
           generateCallable,
       UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable,
-      ClientContext clientContext,
+      ApiClock clock,
+      ScheduledExecutorService executor,
       RetrySettings pollingSettings,
       @Nullable TableAdminRequestContext requestContext) {
 
     RetryAlgorithm<CheckConsistencyResponse> retryAlgorithm =
         new RetryAlgorithm<>(
-            new PollResultAlgorithm(),
-            new ExponentialPollAlgorithm(pollingSettings, clientContext.getClock()));
+            new PollResultAlgorithm(), new ExponentialPollAlgorithm(pollingSettings, clock));
 
     RetryingExecutor<CheckConsistencyResponse> retryingExecutor =
-        new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
+        new ScheduledRetryingExecutor<>(retryAlgorithm, executor);
 
     return new AwaitConsistencyCallable(
         generateCallable, checkCallable, retryingExecutor, requestContext);
   }
 
+  @InternalApi
   public static AwaitConsistencyCallable create(
       UnaryCallable<GenerateConsistencyTokenRequest, GenerateConsistencyTokenResponse>
           generateCallable,
       UnaryCallable<CheckConsistencyRequest, CheckConsistencyResponse> checkCallable,
-      ClientContext clientContext,
+      ApiClock clock,
+      ScheduledExecutorService executor,
       RetrySettings pollingSettings) {
-    return create(generateCallable, checkCallable, clientContext, pollingSettings, null);
+    return create(generateCallable, checkCallable, clock, executor, pollingSettings, null);
   }
 
   @VisibleForTesting

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/admin/v2/stub/EnhancedBigtableTableAdminStub.java

@@ -136,7 +136,8 @@ private AwaitConsistencyCallable createAwaitConsistencyCallable() {
     return AwaitConsistencyCallable.create(
         generateConsistencyTokenCallable(),
         checkConsistencyCallable(),
-        clientContext,
+        clientContext.getClock(),
+        clientContext.getExecutor(),
         pollingSettings,
         requestContext);
   }

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/admin/v2/stub/AwaitConsistencyCallableTest.java

@@ -97,7 +97,8 @@ public void setUp() {
         AwaitConsistencyCallable.create(
             mockGenerateConsistencyTokenCallable,
             mockCheckConsistencyCallable,
-            clientContext,
+            clientContext.getClock(),
+            clientContext.getExecutor(),
             retrySettings,
             REQUEST_CONTEXT);
     awaitReplicationCallable = AwaitReplicationCallable.create(awaitConsistencyCallable);