address comments

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

@@ -28,6 +28,7 @@
 import com.google.api.gax.grpc.GrpcCallSettings;
 import com.google.api.gax.grpc.GrpcRawCallableFactory;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.api.gax.retrying.BasicResultRetryAlgorithm;
 import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
 import com.google.api.gax.retrying.RetryAlgorithm;
 import com.google.api.gax.retrying.RetryingExecutorWithContext;
@@ -763,18 +764,19 @@ public Map<String, String> extract(MutateRowsRequest mutateRowsRequest) {
     ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> withBigtableTracer =
         new BigtableTracerStreamingCallable<>(convertException);
 
-    RetryAlgorithm<Void> retryAlgorithm;
-    ExponentialRetryAlgorithm exponentialRetryAlgorithm =
-        new ExponentialRetryAlgorithm(
-            settings.bulkMutateRowsSettings().getRetrySettings(), clientContext.getClock());
+    BasicResultRetryAlgorithm<Void> resultRetryAlgorithm;
     if (settings.getEnableRetryInfo()) {
-      retryAlgorithm =
-          new RetryAlgorithm<>(new RetryInfoRetryAlgorithm<>(), exponentialRetryAlgorithm);
+      resultRetryAlgorithm = new RetryInfoRetryAlgorithm<>();
     } else {
-      retryAlgorithm =
-          new RetryAlgorithm<>(new ApiResultRetryAlgorithm<>(), exponentialRetryAlgorithm);
+      resultRetryAlgorithm = new ApiResultRetryAlgorithm<>();
     }
 
+    RetryAlgorithm<Void> retryAlgorithm =
+        new RetryAlgorithm<>(
+            resultRetryAlgorithm,
+            new ExponentialRetryAlgorithm(
+                settings.bulkMutateRowsSettings().getRetrySettings(), clientContext.getClock()));
+
     RetryingExecutorWithContext<Void> retryingExecutor =
         new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
 

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

@@ -331,6 +331,7 @@ public boolean getEnableRoutingCookie() {
    * Gets if RetryInfo is enabled. If true, client bases retry decision and back off time on server
    * returned RetryInfo value. Otherwise, client uses {@link RetrySettings}.
    */
+  @BetaApi("RetryInfo is not currently stable and may change in the future")
   public boolean getEnableRetryInfo() {
     return enableRetryInfo;
   }
@@ -931,15 +932,6 @@ public Builder setEnableRoutingCookie(boolean enableRoutingCookie) {
       return this;
     }
 
-    /**
-     * Sets if RetryInfo is enabled. If true, client bases retry decision and back off time on
-     * server returned RetryInfo value. Otherwise, client uses {@link RetrySettings}.
-     */
-    public Builder setEnableRetryInfo(boolean enableRetryInfo) {
-      this.enableRetryInfo = enableRetryInfo;
-      return this;
-    }
-
     /**
      * Gets if routing cookie is enabled. If true, client will retry a request with extra metadata
      * server sent back.
@@ -949,10 +941,21 @@ public boolean getEnableRoutingCookie() {
       return enableRoutingCookie;
     }
 
+    /**
+     * Sets if RetryInfo is enabled. If true, client bases retry decision and back off time on
+     * server returned RetryInfo value. Otherwise, client uses {@link RetrySettings}.
+     */
+    @BetaApi("RetryInfo is not currently stable and may change in the future")
+    public Builder setEnableRetryInfo(boolean enableRetryInfo) {
+      this.enableRetryInfo = enableRetryInfo;
+      return this;
+    }
+
     /**
      * Gets if RetryInfo is enabled. If true, client bases retry decision and back off time on
      * server returned RetryInfo value. Otherwise, client uses {@link RetrySettings}.
      */
+    @BetaApi("RetryInfo is not currently stable and may change in the future")
     public boolean getEnableRetryInfo() {
       return enableRetryInfo;
     }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/ApiException.java

@@ -21,6 +21,8 @@
 @InternalApi
 public class ApiException {
 
+  private ApiException() {}
+
   // TODO: this should replace the existing ApiException#isRetryable() method,
   //  but that cant be done in bigtable, so this lives here for now.
   public static boolean isRetryable2(Throwable e) {

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/RetryInfoRetryAlgorithm.java

@@ -26,6 +26,7 @@
 import io.grpc.Metadata;
 import io.grpc.Status;
 import io.grpc.protobuf.ProtoUtils;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.threeten.bp.Duration;
 
 // TODO move this algorithm to gax
@@ -57,13 +58,7 @@ public TimedAttemptSettings createNextAttempt(
   /** Returns true if previousThrowable is an {@link ApiException} that is retryable. */
   @Override
   public boolean shouldRetry(Throwable previousThrowable, ResponseT previousResponse) {
-    if (extractRetryDelay(previousThrowable) != null) {
-      // First check if server wants us to retry
-      return true;
-    }
-    // Server didn't have retry information, use the local status code config.
-    return (previousThrowable instanceof ApiException
-        && ((ApiException) previousThrowable).isRetryable());
+    return shouldRetry(null, previousThrowable, previousResponse);
   }
 
   /**
@@ -74,23 +69,27 @@ public boolean shouldRetry(Throwable previousThrowable, ResponseT previousRespon
    */
   @Override
   public boolean shouldRetry(
-      RetryingContext context, Throwable previousThrowable, ResponseT previousResponse) {
+      @Nullable RetryingContext context, Throwable previousThrowable, ResponseT previousResponse) {
     if (extractRetryDelay(previousThrowable) != null) {
       // First check if server wants us to retry
       return true;
     }
-    if (context.getRetryableCodes() != null) {
+    if (context != null && context.getRetryableCodes() != null) {
       // Ignore the isRetryable() value of the throwable if the RetryingContext has a specific list
       // of codes that should be retried.
       return ((previousThrowable instanceof ApiException)
           && context
               .getRetryableCodes()
               .contains(((ApiException) previousThrowable).getStatusCode().getCode()));
     }
-    return shouldRetry(previousThrowable, previousResponse);
+    // Server didn't have retry information and there's no retry context, use the local status
+    // code config.
+    return previousThrowable instanceof ApiException
+        && ((ApiException) previousThrowable).isRetryable();
   }
 
-  static Duration extractRetryDelay(Throwable throwable) {
+  @Nullable
+  static Duration extractRetryDelay(@Nullable Throwable throwable) {
     if (throwable == null) {
       return null;
     }