add a flag for retry info and tests

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

@@ -107,6 +107,7 @@
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsUserCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.RowMergingCallable;
+import com.google.cloud.bigtable.gaxx.retrying.ApiResultRetryAlgorithm;
 import com.google.cloud.bigtable.gaxx.retrying.RetryInfoRetryAlgorithm;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
@@ -761,11 +762,18 @@ public Map<String, String> extract(MutateRowsRequest mutateRowsRequest) {
     ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> withBigtableTracer =
         new BigtableTracerStreamingCallable<>(convertException);
 
-    RetryAlgorithm<Void> retryAlgorithm =
-        new RetryAlgorithm<>(
-            new RetryInfoRetryAlgorithm<>(),
-            new ExponentialRetryAlgorithm(
-                settings.bulkMutateRowsSettings().getRetrySettings(), clientContext.getClock()));
+    RetryAlgorithm<Void> retryAlgorithm;
+    ExponentialRetryAlgorithm exponentialRetryAlgorithm =
+        new ExponentialRetryAlgorithm(
+            settings.bulkMutateRowsSettings().getRetrySettings(), clientContext.getClock());
+    if (settings.getEnableRetryInfo()) {
+      retryAlgorithm =
+          new RetryAlgorithm<>(new RetryInfoRetryAlgorithm<>(), exponentialRetryAlgorithm);
+    } else {
+      retryAlgorithm =
+          new RetryAlgorithm<>(new ApiResultRetryAlgorithm<>(), exponentialRetryAlgorithm);
+    }
+
     RetryingExecutorWithContext<Void> retryingExecutor =
         new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
 
@@ -1055,8 +1063,13 @@ public Map<String, String> extract(PingAndWarmRequest request) {
 
   private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> withRetries(
       UnaryCallable<RequestT, ResponseT> innerCallable, UnaryCallSettings<?, ?> unaryCallSettings) {
-    UnaryCallable<RequestT, ResponseT> retrying =
-        Callables.retrying(innerCallable, unaryCallSettings, clientContext);
+    UnaryCallable<RequestT, ResponseT> retrying;
+    if (settings.getEnableRetryInfo()) {
+      retrying = com.google.cloud.bigtable.gaxx.retrying.Callables.retrying(innerCallable, unaryCallSettings, clientContext);
+    } else {
+      retrying =
+              Callables.retrying(innerCallable, unaryCallSettings, clientContext);
+    }
     if (settings.getEnableRoutingCookie()) {
       return new CookiesUnaryCallable<>(retrying);
     }
@@ -1066,8 +1079,14 @@ private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> withRetries(
   private <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT> withRetries(
       ServerStreamingCallable<RequestT, ResponseT> innerCallable,
       ServerStreamingCallSettings<RequestT, ResponseT> serverStreamingCallSettings) {
-    ServerStreamingCallable<RequestT, ResponseT> retrying =
-        Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext);
+
+    ServerStreamingCallable<RequestT, ResponseT> retrying;
+    if (settings.getEnableRetryInfo()) {
+      retrying =
+              com.google.cloud.bigtable.gaxx.retrying.Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext);
+    } else {
+      retrying = Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext);
+    }
     if (settings.getEnableRoutingCookie()) {
       return new CookiesServerStreamingCallable<>(retrying);
     }

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

@@ -213,6 +213,7 @@ public class EnhancedBigtableStubSettings extends StubSettings<EnhancedBigtableS
   private ImmutableList<String> primedTableIds;
   private final Map<String, String> jwtAudienceMapping;
   private final boolean enableRoutingCookie;
+  private final boolean enableRetryInfo;
 
   private final ServerStreamingCallSettings<Query, Row> readRowsSettings;
   private final UnaryCallSettings<Query, Row> readRowSettings;
@@ -255,6 +256,7 @@ private EnhancedBigtableStubSettings(Builder builder) {
     primedTableIds = builder.primedTableIds;
     jwtAudienceMapping = builder.jwtAudienceMapping;
     enableRoutingCookie = builder.enableRoutingCookie;
+    enableRetryInfo = builder.enableRetryInfo;
 
     // Per method settings.
     readRowsSettings = builder.readRowsSettings.build();
@@ -325,6 +327,14 @@ public boolean getEnableRoutingCookie() {
     return enableRoutingCookie;
   }
 
+  /**
+   * 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}.
+   */
+  public boolean getEnableRetryInfo() {
+    return enableRetryInfo;
+  }
+
   /** Returns a builder for the default ChannelProvider for this service. */
   public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() {
     return BigtableStubSettings.defaultGrpcTransportProviderBuilder()
@@ -608,6 +618,7 @@ public static class Builder extends StubSettings.Builder<EnhancedBigtableStubSet
     private ImmutableList<String> primedTableIds;
     private Map<String, String> jwtAudienceMapping;
     private boolean enableRoutingCookie;
+    private boolean enableRetryInfo;
 
     private final ServerStreamingCallSettings.Builder<Query, Row> readRowsSettings;
     private final UnaryCallSettings.Builder<Query, Row> readRowSettings;
@@ -641,6 +652,7 @@ private Builder() {
       jwtAudienceMapping = DEFAULT_JWT_AUDIENCE_MAPPING;
       setCredentialsProvider(defaultCredentialsProviderBuilder().build());
       this.enableRoutingCookie = true;
+      this.enableRetryInfo = true;
 
       // Defaults provider
       BigtableStubSettings.Builder baseDefaults = BigtableStubSettings.newBuilder();
@@ -760,6 +772,7 @@ private Builder(EnhancedBigtableStubSettings settings) {
       primedTableIds = settings.primedTableIds;
       jwtAudienceMapping = settings.jwtAudienceMapping;
       enableRoutingCookie = settings.enableRoutingCookie;
+      enableRetryInfo = settings.enableRetryInfo;
 
       // Per method settings.
       readRowsSettings = settings.readRowsSettings.toBuilder();
@@ -918,6 +931,15 @@ 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.
@@ -927,6 +949,14 @@ public boolean getEnableRoutingCookie() {
       return enableRoutingCookie;
     }
 
+    /**
+     * 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}.
+     */
+    public boolean getEnableRetryInfo() {
+      return enableRetryInfo;
+    }
+
     /** Returns the builder for the settings used for calls to readRows. */
     public ServerStreamingCallSettings.Builder<Query, Row> readRowsSettings() {
       return readRowsSettings;
@@ -1054,6 +1084,7 @@ public String toString() {
         .add("primedTableIds", primedTableIds)
         .add("jwtAudienceMapping", jwtAudienceMapping)
         .add("enableRoutingCookie", enableRoutingCookie)
+        .add("enableRetryInfo", enableRetryInfo)
         .add("readRowsSettings", readRowsSettings)
         .add("readRowSettings", readRowSettings)
         .add("sampleRowKeysSettings", sampleRowKeysSettings)

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2018 Google LLC
+ * Copyright 2023 Google LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

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

@@ -1,3 +1,18 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 package com.google.cloud.bigtable.gaxx.retrying;
 
 import com.google.api.core.InternalApi;

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

@@ -78,6 +78,7 @@ public void settingsAreNotLostTest() {
     WatchdogProvider watchdogProvider = Mockito.mock(WatchdogProvider.class);
     Duration watchdogInterval = Duration.ofSeconds(12);
     boolean enableRoutingCookie = false;
+    boolean enableRetryInfo = false;
 
     EnhancedBigtableStubSettings.Builder builder =
         EnhancedBigtableStubSettings.newBuilder()
@@ -89,7 +90,8 @@ public void settingsAreNotLostTest() {
             .setCredentialsProvider(credentialsProvider)
             .setStreamWatchdogProvider(watchdogProvider)
             .setStreamWatchdogCheckInterval(watchdogInterval)
-            .setEnableRoutingCookie(enableRoutingCookie);
+            .setEnableRoutingCookie(enableRoutingCookie)
+            .setEnableRetryInfo(enableRetryInfo);
 
     verifyBuilder(
         builder,
@@ -101,7 +103,8 @@ public void settingsAreNotLostTest() {
         credentialsProvider,
         watchdogProvider,
         watchdogInterval,
-        enableRoutingCookie);
+        enableRoutingCookie,
+        enableRetryInfo);
     verifySettings(
         builder.build(),
         projectId,
@@ -112,7 +115,8 @@ public void settingsAreNotLostTest() {
         credentialsProvider,
         watchdogProvider,
         watchdogInterval,
-        enableRoutingCookie);
+        enableRoutingCookie,
+        enableRetryInfo);
     verifyBuilder(
         builder.build().toBuilder(),
         projectId,
@@ -123,7 +127,8 @@ public void settingsAreNotLostTest() {
         credentialsProvider,
         watchdogProvider,
         watchdogInterval,
-        enableRoutingCookie);
+        enableRoutingCookie,
+        enableRetryInfo);
   }
 
   private void verifyBuilder(
@@ -136,7 +141,8 @@ private void verifyBuilder(
       CredentialsProvider credentialsProvider,
       WatchdogProvider watchdogProvider,
       Duration watchdogInterval,
-      boolean enableRoutingCookie) {
+      boolean enableRoutingCookie,
+      boolean enableRetryInfo) {
     assertThat(builder.getProjectId()).isEqualTo(projectId);
     assertThat(builder.getInstanceId()).isEqualTo(instanceId);
     assertThat(builder.getAppProfileId()).isEqualTo(appProfileId);
@@ -146,6 +152,7 @@ private void verifyBuilder(
     assertThat(builder.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider);
     assertThat(builder.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval);
     assertThat(builder.getEnableRoutingCookie()).isEqualTo(enableRoutingCookie);
+    assertThat(builder.getEnableRetryInfo()).isEqualTo(enableRetryInfo);
   }
 
   private void verifySettings(
@@ -158,7 +165,8 @@ private void verifySettings(
       CredentialsProvider credentialsProvider,
       WatchdogProvider watchdogProvider,
       Duration watchdogInterval,
-      boolean enableRoutingCookie) {
+      boolean enableRoutingCookie,
+      boolean enableRetryInfo) {
     assertThat(settings.getProjectId()).isEqualTo(projectId);
     assertThat(settings.getInstanceId()).isEqualTo(instanceId);
     assertThat(settings.getAppProfileId()).isEqualTo(appProfileId);
@@ -168,6 +176,7 @@ private void verifySettings(
     assertThat(settings.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider);
     assertThat(settings.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval);
     assertThat(settings.getEnableRoutingCookie()).isEqualTo(enableRoutingCookie);
+    assertThat(settings.getEnableRetryInfo()).isEqualTo(enableRetryInfo);
   }
 
   @Test
@@ -797,17 +806,49 @@ public void routingCookieIsEnabled() throws IOException {
     CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
     Mockito.when(credentialsProvider.getCredentials()).thenReturn(new FakeCredentials());
     EnhancedBigtableStubSettings.Builder builder =
-        EnhancedBigtableStubSettings.newBuilder()
-            .setProjectId(dummyProjectId)
-            .setInstanceId(dummyInstanceId)
-            .setCredentialsProvider(credentialsProvider);
+            EnhancedBigtableStubSettings.newBuilder()
+                    .setProjectId(dummyProjectId)
+                    .setInstanceId(dummyInstanceId)
+                    .setCredentialsProvider(credentialsProvider);
     assertThat(builder.getEnableRoutingCookie()).isTrue();
     assertThat(builder.build().getEnableRoutingCookie()).isTrue();
     assertThat(builder.build().toBuilder().getEnableRoutingCookie()).isTrue();
   }
 
+  public void enableRetryInfoDefaultValueTest() throws IOException {
+    String dummyProjectId = "my-project";
+    String dummyInstanceId = "my-instance";
+    CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
+    Mockito.when(credentialsProvider.getCredentials()).thenReturn(new FakeCredentials());
+    EnhancedBigtableStubSettings.Builder builder =
+            EnhancedBigtableStubSettings.newBuilder()
+                    .setProjectId(dummyProjectId)
+                    .setInstanceId(dummyInstanceId)
+                    .setCredentialsProvider(credentialsProvider);
+    assertThat(builder.getEnableRetryInfo()).isTrue();
+    assertThat(builder.build().getEnableRetryInfo()).isTrue();
+    assertThat(builder.build().toBuilder().getEnableRetryInfo()).isTrue();
+  }
+
   @Test
   public void routingCookieFalseValueSet() throws IOException {
+    String dummyProjectId = "my-project";
+    String dummyInstanceId = "my-instance";
+    CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
+    Mockito.when(credentialsProvider.getCredentials()).thenReturn(new FakeCredentials());
+    EnhancedBigtableStubSettings.Builder builder =
+            EnhancedBigtableStubSettings.newBuilder()
+                    .setProjectId(dummyProjectId)
+                    .setInstanceId(dummyInstanceId)
+                    .setEnableRoutingCookie(false)
+                    .setCredentialsProvider(credentialsProvider);
+    assertThat(builder.getEnableRoutingCookie()).isFalse();
+    assertThat(builder.build().getEnableRoutingCookie()).isFalse();
+    assertThat(builder.build().toBuilder().getEnableRoutingCookie()).isFalse();
+  }
+
+  @Test
+  public void enableRetryInfoFalseValueTest() throws IOException {
     String dummyProjectId = "my-project";
     String dummyInstanceId = "my-instance";
     CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
@@ -816,11 +857,11 @@ public void routingCookieFalseValueSet() throws IOException {
         EnhancedBigtableStubSettings.newBuilder()
             .setProjectId(dummyProjectId)
             .setInstanceId(dummyInstanceId)
-            .setEnableRoutingCookie(false)
+            .setEnableRetryInfo(false)
             .setCredentialsProvider(credentialsProvider);
-    assertThat(builder.getEnableRoutingCookie()).isFalse();
-    assertThat(builder.build().getEnableRoutingCookie()).isFalse();
-    assertThat(builder.build().toBuilder().getEnableRoutingCookie()).isFalse();
+    assertThat(builder.getEnableRetryInfo()).isFalse();
+    assertThat(builder.build().getEnableRetryInfo()).isFalse();
+    assertThat(builder.build().toBuilder().getEnableRetryInfo()).isFalse();
   }
 
   static final String[] SETTINGS_LIST = {
@@ -831,6 +872,7 @@ public void routingCookieFalseValueSet() throws IOException {
     "primedTableIds",
     "jwtAudienceMapping",
     "enableRoutingCookie",
+    "enableRetryInfo",
     "readRowsSettings",
     "readRowSettings",
     "sampleRowKeysSettings",