Skip to content
This repository was archived by the owner on May 8, 2026. It is now read-only.

Commit a3ef1d1

Browse files
mutianfmutianf
committed
add a flag for retry info and tests
1 parent 5f6d583 commit a3ef1d1

6 files changed

Lines changed: 427 additions & 26 deletions

File tree

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

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@
107107
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable;
108108
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsUserCallable;
109109
import com.google.cloud.bigtable.data.v2.stub.readrows.RowMergingCallable;
110+
import com.google.cloud.bigtable.gaxx.retrying.ApiResultRetryAlgorithm;
110111
import com.google.cloud.bigtable.gaxx.retrying.RetryInfoRetryAlgorithm;
111112
import com.google.common.base.MoreObjects;
112113
import com.google.common.base.Preconditions;
@@ -761,11 +762,18 @@ public Map<String, String> extract(MutateRowsRequest mutateRowsRequest) {
761762
ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> withBigtableTracer =
762763
new BigtableTracerStreamingCallable<>(convertException);
763764

764-
RetryAlgorithm<Void> retryAlgorithm =
765-
new RetryAlgorithm<>(
766-
new RetryInfoRetryAlgorithm<>(),
767-
new ExponentialRetryAlgorithm(
768-
settings.bulkMutateRowsSettings().getRetrySettings(), clientContext.getClock()));
765+
RetryAlgorithm<Void> retryAlgorithm;
766+
ExponentialRetryAlgorithm exponentialRetryAlgorithm =
767+
new ExponentialRetryAlgorithm(
768+
settings.bulkMutateRowsSettings().getRetrySettings(), clientContext.getClock());
769+
if (settings.getEnableRetryInfo()) {
770+
retryAlgorithm =
771+
new RetryAlgorithm<>(new RetryInfoRetryAlgorithm<>(), exponentialRetryAlgorithm);
772+
} else {
773+
retryAlgorithm =
774+
new RetryAlgorithm<>(new ApiResultRetryAlgorithm<>(), exponentialRetryAlgorithm);
775+
}
776+
769777
RetryingExecutorWithContext<Void> retryingExecutor =
770778
new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
771779

@@ -1055,8 +1063,13 @@ public Map<String, String> extract(PingAndWarmRequest request) {
10551063

10561064
private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> withRetries(
10571065
UnaryCallable<RequestT, ResponseT> innerCallable, UnaryCallSettings<?, ?> unaryCallSettings) {
1058-
UnaryCallable<RequestT, ResponseT> retrying =
1059-
Callables.retrying(innerCallable, unaryCallSettings, clientContext);
1066+
UnaryCallable<RequestT, ResponseT> retrying;
1067+
if (settings.getEnableRetryInfo()) {
1068+
retrying = com.google.cloud.bigtable.gaxx.retrying.Callables.retrying(innerCallable, unaryCallSettings, clientContext);
1069+
} else {
1070+
retrying =
1071+
Callables.retrying(innerCallable, unaryCallSettings, clientContext);
1072+
}
10601073
if (settings.getEnableRoutingCookie()) {
10611074
return new CookiesUnaryCallable<>(retrying);
10621075
}
@@ -1066,8 +1079,14 @@ private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> withRetries(
10661079
private <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT> withRetries(
10671080
ServerStreamingCallable<RequestT, ResponseT> innerCallable,
10681081
ServerStreamingCallSettings<RequestT, ResponseT> serverStreamingCallSettings) {
1069-
ServerStreamingCallable<RequestT, ResponseT> retrying =
1070-
Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext);
1082+
1083+
ServerStreamingCallable<RequestT, ResponseT> retrying;
1084+
if (settings.getEnableRetryInfo()) {
1085+
retrying =
1086+
com.google.cloud.bigtable.gaxx.retrying.Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext);
1087+
} else {
1088+
retrying = Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext);
1089+
}
10711090
if (settings.getEnableRoutingCookie()) {
10721091
return new CookiesServerStreamingCallable<>(retrying);
10731092
}

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,7 @@ public class EnhancedBigtableStubSettings extends StubSettings<EnhancedBigtableS
213213
private ImmutableList<String> primedTableIds;
214214
private final Map<String, String> jwtAudienceMapping;
215215
private final boolean enableRoutingCookie;
216+
private final boolean enableRetryInfo;
216217

217218
private final ServerStreamingCallSettings<Query, Row> readRowsSettings;
218219
private final UnaryCallSettings<Query, Row> readRowSettings;
@@ -255,6 +256,7 @@ private EnhancedBigtableStubSettings(Builder builder) {
255256
primedTableIds = builder.primedTableIds;
256257
jwtAudienceMapping = builder.jwtAudienceMapping;
257258
enableRoutingCookie = builder.enableRoutingCookie;
259+
enableRetryInfo = builder.enableRetryInfo;
258260

259261
// Per method settings.
260262
readRowsSettings = builder.readRowsSettings.build();
@@ -325,6 +327,14 @@ public boolean getEnableRoutingCookie() {
325327
return enableRoutingCookie;
326328
}
327329

330+
/**
331+
* Gets if RetryInfo is enabled. If true, client bases retry decision and back off time on server
332+
* returned RetryInfo value. Otherwise, client uses {@link RetrySettings}.
333+
*/
334+
public boolean getEnableRetryInfo() {
335+
return enableRetryInfo;
336+
}
337+
328338
/** Returns a builder for the default ChannelProvider for this service. */
329339
public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() {
330340
return BigtableStubSettings.defaultGrpcTransportProviderBuilder()
@@ -608,6 +618,7 @@ public static class Builder extends StubSettings.Builder<EnhancedBigtableStubSet
608618
private ImmutableList<String> primedTableIds;
609619
private Map<String, String> jwtAudienceMapping;
610620
private boolean enableRoutingCookie;
621+
private boolean enableRetryInfo;
611622

612623
private final ServerStreamingCallSettings.Builder<Query, Row> readRowsSettings;
613624
private final UnaryCallSettings.Builder<Query, Row> readRowSettings;
@@ -641,6 +652,7 @@ private Builder() {
641652
jwtAudienceMapping = DEFAULT_JWT_AUDIENCE_MAPPING;
642653
setCredentialsProvider(defaultCredentialsProviderBuilder().build());
643654
this.enableRoutingCookie = true;
655+
this.enableRetryInfo = true;
644656

645657
// Defaults provider
646658
BigtableStubSettings.Builder baseDefaults = BigtableStubSettings.newBuilder();
@@ -760,6 +772,7 @@ private Builder(EnhancedBigtableStubSettings settings) {
760772
primedTableIds = settings.primedTableIds;
761773
jwtAudienceMapping = settings.jwtAudienceMapping;
762774
enableRoutingCookie = settings.enableRoutingCookie;
775+
enableRetryInfo = settings.enableRetryInfo;
763776

764777
// Per method settings.
765778
readRowsSettings = settings.readRowsSettings.toBuilder();
@@ -918,6 +931,15 @@ public Builder setEnableRoutingCookie(boolean enableRoutingCookie) {
918931
return this;
919932
}
920933

934+
/**
935+
* Sets if RetryInfo is enabled. If true, client bases retry decision and back off time on
936+
* server returned RetryInfo value. Otherwise, client uses {@link RetrySettings}.
937+
*/
938+
public Builder setEnableRetryInfo(boolean enableRetryInfo) {
939+
this.enableRetryInfo = enableRetryInfo;
940+
return this;
941+
}
942+
921943
/**
922944
* Gets if routing cookie is enabled. If true, client will retry a request with extra metadata
923945
* server sent back.
@@ -927,6 +949,14 @@ public boolean getEnableRoutingCookie() {
927949
return enableRoutingCookie;
928950
}
929951

952+
/**
953+
* Gets if RetryInfo is enabled. If true, client bases retry decision and back off time on
954+
* server returned RetryInfo value. Otherwise, client uses {@link RetrySettings}.
955+
*/
956+
public boolean getEnableRetryInfo() {
957+
return enableRetryInfo;
958+
}
959+
930960
/** Returns the builder for the settings used for calls to readRows. */
931961
public ServerStreamingCallSettings.Builder<Query, Row> readRowsSettings() {
932962
return readRowsSettings;
@@ -1054,6 +1084,7 @@ public String toString() {
10541084
.add("primedTableIds", primedTableIds)
10551085
.add("jwtAudienceMapping", jwtAudienceMapping)
10561086
.add("enableRoutingCookie", enableRoutingCookie)
1087+
.add("enableRetryInfo", enableRetryInfo)
10571088
.add("readRowsSettings", readRowsSettings)
10581089
.add("readRowSettings", readRowSettings)
10591090
.add("sampleRowKeysSettings", sampleRowKeysSettings)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2018 Google LLC
2+
* Copyright 2023 Google LLC
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* 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‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,18 @@
1+
/*
2+
* Copyright 2023 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
116
package com.google.cloud.bigtable.gaxx.retrying;
217

318
import com.google.api.core.InternalApi;

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

Lines changed: 56 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ public void settingsAreNotLostTest() {
7878
WatchdogProvider watchdogProvider = Mockito.mock(WatchdogProvider.class);
7979
Duration watchdogInterval = Duration.ofSeconds(12);
8080
boolean enableRoutingCookie = false;
81+
boolean enableRetryInfo = false;
8182

8283
EnhancedBigtableStubSettings.Builder builder =
8384
EnhancedBigtableStubSettings.newBuilder()
@@ -89,7 +90,8 @@ public void settingsAreNotLostTest() {
8990
.setCredentialsProvider(credentialsProvider)
9091
.setStreamWatchdogProvider(watchdogProvider)
9192
.setStreamWatchdogCheckInterval(watchdogInterval)
92-
.setEnableRoutingCookie(enableRoutingCookie);
93+
.setEnableRoutingCookie(enableRoutingCookie)
94+
.setEnableRetryInfo(enableRetryInfo);
9395

9496
verifyBuilder(
9597
builder,
@@ -101,7 +103,8 @@ public void settingsAreNotLostTest() {
101103
credentialsProvider,
102104
watchdogProvider,
103105
watchdogInterval,
104-
enableRoutingCookie);
106+
enableRoutingCookie,
107+
enableRetryInfo);
105108
verifySettings(
106109
builder.build(),
107110
projectId,
@@ -112,7 +115,8 @@ public void settingsAreNotLostTest() {
112115
credentialsProvider,
113116
watchdogProvider,
114117
watchdogInterval,
115-
enableRoutingCookie);
118+
enableRoutingCookie,
119+
enableRetryInfo);
116120
verifyBuilder(
117121
builder.build().toBuilder(),
118122
projectId,
@@ -123,7 +127,8 @@ public void settingsAreNotLostTest() {
123127
credentialsProvider,
124128
watchdogProvider,
125129
watchdogInterval,
126-
enableRoutingCookie);
130+
enableRoutingCookie,
131+
enableRetryInfo);
127132
}
128133

129134
private void verifyBuilder(
@@ -136,7 +141,8 @@ private void verifyBuilder(
136141
CredentialsProvider credentialsProvider,
137142
WatchdogProvider watchdogProvider,
138143
Duration watchdogInterval,
139-
boolean enableRoutingCookie) {
144+
boolean enableRoutingCookie,
145+
boolean enableRetryInfo) {
140146
assertThat(builder.getProjectId()).isEqualTo(projectId);
141147
assertThat(builder.getInstanceId()).isEqualTo(instanceId);
142148
assertThat(builder.getAppProfileId()).isEqualTo(appProfileId);
@@ -146,6 +152,7 @@ private void verifyBuilder(
146152
assertThat(builder.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider);
147153
assertThat(builder.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval);
148154
assertThat(builder.getEnableRoutingCookie()).isEqualTo(enableRoutingCookie);
155+
assertThat(builder.getEnableRetryInfo()).isEqualTo(enableRetryInfo);
149156
}
150157

151158
private void verifySettings(
@@ -158,7 +165,8 @@ private void verifySettings(
158165
CredentialsProvider credentialsProvider,
159166
WatchdogProvider watchdogProvider,
160167
Duration watchdogInterval,
161-
boolean enableRoutingCookie) {
168+
boolean enableRoutingCookie,
169+
boolean enableRetryInfo) {
162170
assertThat(settings.getProjectId()).isEqualTo(projectId);
163171
assertThat(settings.getInstanceId()).isEqualTo(instanceId);
164172
assertThat(settings.getAppProfileId()).isEqualTo(appProfileId);
@@ -168,6 +176,7 @@ private void verifySettings(
168176
assertThat(settings.getStreamWatchdogProvider()).isSameInstanceAs(watchdogProvider);
169177
assertThat(settings.getStreamWatchdogCheckInterval()).isEqualTo(watchdogInterval);
170178
assertThat(settings.getEnableRoutingCookie()).isEqualTo(enableRoutingCookie);
179+
assertThat(settings.getEnableRetryInfo()).isEqualTo(enableRetryInfo);
171180
}
172181

173182
@Test
@@ -797,17 +806,49 @@ public void routingCookieIsEnabled() throws IOException {
797806
CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
798807
Mockito.when(credentialsProvider.getCredentials()).thenReturn(new FakeCredentials());
799808
EnhancedBigtableStubSettings.Builder builder =
800-
EnhancedBigtableStubSettings.newBuilder()
801-
.setProjectId(dummyProjectId)
802-
.setInstanceId(dummyInstanceId)
803-
.setCredentialsProvider(credentialsProvider);
809+
EnhancedBigtableStubSettings.newBuilder()
810+
.setProjectId(dummyProjectId)
811+
.setInstanceId(dummyInstanceId)
812+
.setCredentialsProvider(credentialsProvider);
804813
assertThat(builder.getEnableRoutingCookie()).isTrue();
805814
assertThat(builder.build().getEnableRoutingCookie()).isTrue();
806815
assertThat(builder.build().toBuilder().getEnableRoutingCookie()).isTrue();
807816
}
808817

818+
public void enableRetryInfoDefaultValueTest() throws IOException {
819+
String dummyProjectId = "my-project";
820+
String dummyInstanceId = "my-instance";
821+
CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
822+
Mockito.when(credentialsProvider.getCredentials()).thenReturn(new FakeCredentials());
823+
EnhancedBigtableStubSettings.Builder builder =
824+
EnhancedBigtableStubSettings.newBuilder()
825+
.setProjectId(dummyProjectId)
826+
.setInstanceId(dummyInstanceId)
827+
.setCredentialsProvider(credentialsProvider);
828+
assertThat(builder.getEnableRetryInfo()).isTrue();
829+
assertThat(builder.build().getEnableRetryInfo()).isTrue();
830+
assertThat(builder.build().toBuilder().getEnableRetryInfo()).isTrue();
831+
}
832+
809833
@Test
810834
public void routingCookieFalseValueSet() throws IOException {
835+
String dummyProjectId = "my-project";
836+
String dummyInstanceId = "my-instance";
837+
CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
838+
Mockito.when(credentialsProvider.getCredentials()).thenReturn(new FakeCredentials());
839+
EnhancedBigtableStubSettings.Builder builder =
840+
EnhancedBigtableStubSettings.newBuilder()
841+
.setProjectId(dummyProjectId)
842+
.setInstanceId(dummyInstanceId)
843+
.setEnableRoutingCookie(false)
844+
.setCredentialsProvider(credentialsProvider);
845+
assertThat(builder.getEnableRoutingCookie()).isFalse();
846+
assertThat(builder.build().getEnableRoutingCookie()).isFalse();
847+
assertThat(builder.build().toBuilder().getEnableRoutingCookie()).isFalse();
848+
}
849+
850+
@Test
851+
public void enableRetryInfoFalseValueTest() throws IOException {
811852
String dummyProjectId = "my-project";
812853
String dummyInstanceId = "my-instance";
813854
CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
@@ -816,11 +857,11 @@ public void routingCookieFalseValueSet() throws IOException {
816857
EnhancedBigtableStubSettings.newBuilder()
817858
.setProjectId(dummyProjectId)
818859
.setInstanceId(dummyInstanceId)
819-
.setEnableRoutingCookie(false)
860+
.setEnableRetryInfo(false)
820861
.setCredentialsProvider(credentialsProvider);
821-
assertThat(builder.getEnableRoutingCookie()).isFalse();
822-
assertThat(builder.build().getEnableRoutingCookie()).isFalse();
823-
assertThat(builder.build().toBuilder().getEnableRoutingCookie()).isFalse();
862+
assertThat(builder.getEnableRetryInfo()).isFalse();
863+
assertThat(builder.build().getEnableRetryInfo()).isFalse();
864+
assertThat(builder.build().toBuilder().getEnableRetryInfo()).isFalse();
824865
}
825866

826867
static final String[] SETTINGS_LIST = {
@@ -831,6 +872,7 @@ public void routingCookieFalseValueSet() throws IOException {
831872
"primedTableIds",
832873
"jwtAudienceMapping",
833874
"enableRoutingCookie",
875+
"enableRetryInfo",
834876
"readRowsSettings",
835877
"readRowSettings",
836878
"sampleRowKeysSettings",

0 commit comments

Comments
 (0)