Skip to content

Commit f1b7fc7

Browse files
mutianfmutianf
authored
feat: handle retry info so client respect the delay server sets (#2026)
* feat: add a flag to add / remove routing cookie from callable chain
1 parent 201e631 commit f1b7fc7

9 files changed

Lines changed: 780 additions & 45 deletions

File tree

‎google-cloud-bigtable/pom.xml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,10 @@
161161
<artifactId>grpc-alts</artifactId>
162162
<scope>runtime</scope>
163163
</dependency>
164-
<!-- Runtime dependencies for credentials -->
164+
<dependency>
165+
<groupId>org.checkerframework</groupId>
166+
<artifactId>checker-qual</artifactId>
167+
</dependency><!-- Runtime dependencies for credentials -->
165168
<dependency>
166169
<groupId>com.google.http-client</groupId>
167170
<artifactId>google-http-client</artifactId>

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

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.google.api.gax.grpc.GrpcCallSettings;
2929
import com.google.api.gax.grpc.GrpcRawCallableFactory;
3030
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
31+
import com.google.api.gax.retrying.BasicResultRetryAlgorithm;
3132
import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
3233
import com.google.api.gax.retrying.RetryAlgorithm;
3334
import com.google.api.gax.retrying.RetryingExecutorWithContext;
@@ -108,6 +109,7 @@
108109
import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsUserCallable;
109110
import com.google.cloud.bigtable.data.v2.stub.readrows.RowMergingCallable;
110111
import com.google.cloud.bigtable.gaxx.retrying.ApiResultRetryAlgorithm;
112+
import com.google.cloud.bigtable.gaxx.retrying.RetryInfoRetryAlgorithm;
111113
import com.google.common.base.MoreObjects;
112114
import com.google.common.base.Preconditions;
113115
import com.google.common.collect.ImmutableList;
@@ -762,11 +764,19 @@ public Map<String, String> extract(MutateRowsRequest mutateRowsRequest) {
762764
ServerStreamingCallable<MutateRowsRequest, MutateRowsResponse> withBigtableTracer =
763765
new BigtableTracerStreamingCallable<>(convertException);
764766

767+
BasicResultRetryAlgorithm<Void> resultRetryAlgorithm;
768+
if (settings.getEnableRetryInfo()) {
769+
resultRetryAlgorithm = new RetryInfoRetryAlgorithm<>();
770+
} else {
771+
resultRetryAlgorithm = new ApiResultRetryAlgorithm<>();
772+
}
773+
765774
RetryAlgorithm<Void> retryAlgorithm =
766775
new RetryAlgorithm<>(
767-
new ApiResultRetryAlgorithm<Void>(),
776+
resultRetryAlgorithm,
768777
new ExponentialRetryAlgorithm(
769778
settings.bulkMutateRowsSettings().getRetrySettings(), clientContext.getClock()));
779+
770780
RetryingExecutorWithContext<Void> retryingExecutor =
771781
new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
772782

@@ -1056,8 +1066,14 @@ public Map<String, String> extract(PingAndWarmRequest request) {
10561066

10571067
private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> withRetries(
10581068
UnaryCallable<RequestT, ResponseT> innerCallable, UnaryCallSettings<?, ?> unaryCallSettings) {
1059-
UnaryCallable<RequestT, ResponseT> retrying =
1060-
Callables.retrying(innerCallable, unaryCallSettings, clientContext);
1069+
UnaryCallable<RequestT, ResponseT> retrying;
1070+
if (settings.getEnableRetryInfo()) {
1071+
retrying =
1072+
com.google.cloud.bigtable.gaxx.retrying.Callables.retrying(
1073+
innerCallable, unaryCallSettings, clientContext);
1074+
} else {
1075+
retrying = Callables.retrying(innerCallable, unaryCallSettings, clientContext);
1076+
}
10611077
if (settings.getEnableRoutingCookie()) {
10621078
return new CookiesUnaryCallable<>(retrying);
10631079
}
@@ -1067,8 +1083,15 @@ private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> withRetries(
10671083
private <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT> withRetries(
10681084
ServerStreamingCallable<RequestT, ResponseT> innerCallable,
10691085
ServerStreamingCallSettings<RequestT, ResponseT> serverStreamingCallSettings) {
1070-
ServerStreamingCallable<RequestT, ResponseT> retrying =
1071-
Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext);
1086+
1087+
ServerStreamingCallable<RequestT, ResponseT> retrying;
1088+
if (settings.getEnableRetryInfo()) {
1089+
retrying =
1090+
com.google.cloud.bigtable.gaxx.retrying.Callables.retrying(
1091+
innerCallable, serverStreamingCallSettings, clientContext);
1092+
} else {
1093+
retrying = Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext);
1094+
}
10721095
if (settings.getEnableRoutingCookie()) {
10731096
return new CookiesServerStreamingCallable<>(retrying);
10741097
}

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

Lines changed: 34 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,15 @@ 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+
@BetaApi("RetryInfo is not currently stable and may change in the future")
335+
public boolean getEnableRetryInfo() {
336+
return enableRetryInfo;
337+
}
338+
328339
/** Returns a builder for the default ChannelProvider for this service. */
329340
public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() {
330341
return BigtableStubSettings.defaultGrpcTransportProviderBuilder()
@@ -608,6 +619,7 @@ public static class Builder extends StubSettings.Builder<EnhancedBigtableStubSet
608619
private ImmutableList<String> primedTableIds;
609620
private Map<String, String> jwtAudienceMapping;
610621
private boolean enableRoutingCookie;
622+
private boolean enableRetryInfo;
611623

612624
private final ServerStreamingCallSettings.Builder<Query, Row> readRowsSettings;
613625
private final UnaryCallSettings.Builder<Query, Row> readRowSettings;
@@ -641,6 +653,7 @@ private Builder() {
641653
jwtAudienceMapping = DEFAULT_JWT_AUDIENCE_MAPPING;
642654
setCredentialsProvider(defaultCredentialsProviderBuilder().build());
643655
this.enableRoutingCookie = true;
656+
this.enableRetryInfo = true;
644657

645658
// Defaults provider
646659
BigtableStubSettings.Builder baseDefaults = BigtableStubSettings.newBuilder();
@@ -760,6 +773,7 @@ private Builder(EnhancedBigtableStubSettings settings) {
760773
primedTableIds = settings.primedTableIds;
761774
jwtAudienceMapping = settings.jwtAudienceMapping;
762775
enableRoutingCookie = settings.enableRoutingCookie;
776+
enableRetryInfo = settings.enableRetryInfo;
763777

764778
// Per method settings.
765779
readRowsSettings = settings.readRowsSettings.toBuilder();
@@ -927,6 +941,25 @@ public boolean getEnableRoutingCookie() {
927941
return enableRoutingCookie;
928942
}
929943

944+
/**
945+
* Sets if RetryInfo is enabled. If true, client bases retry decision and back off time on
946+
* server returned RetryInfo value. Otherwise, client uses {@link RetrySettings}.
947+
*/
948+
@BetaApi("RetryInfo is not currently stable and may change in the future")
949+
public Builder setEnableRetryInfo(boolean enableRetryInfo) {
950+
this.enableRetryInfo = enableRetryInfo;
951+
return this;
952+
}
953+
954+
/**
955+
* Gets if RetryInfo is enabled. If true, client bases retry decision and back off time on
956+
* server returned RetryInfo value. Otherwise, client uses {@link RetrySettings}.
957+
*/
958+
@BetaApi("RetryInfo is not currently stable and may change in the future")
959+
public boolean getEnableRetryInfo() {
960+
return enableRetryInfo;
961+
}
962+
930963
/** Returns the builder for the settings used for calls to readRows. */
931964
public ServerStreamingCallSettings.Builder<Query, Row> readRowsSettings() {
932965
return readRowsSettings;
@@ -1054,6 +1087,7 @@ public String toString() {
10541087
.add("primedTableIds", primedTableIds)
10551088
.add("jwtAudienceMapping", jwtAudienceMapping)
10561089
.add("enableRoutingCookie", enableRoutingCookie)
1090+
.add("enableRetryInfo", enableRetryInfo)
10571091
.add("readRowsSettings", readRowsSettings)
10581092
.add("readRowSettings", readRowSettings)
10591093
.add("sampleRowKeysSettings", sampleRowKeysSettings)

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import com.google.bigtable.v2.MutateRowsResponse.Entry;
3232
import com.google.cloud.bigtable.data.v2.models.MutateRowsException;
3333
import com.google.cloud.bigtable.data.v2.models.MutateRowsException.FailedMutation;
34+
import com.google.cloud.bigtable.gaxx.retrying.ApiExceptions;
3435
import com.google.cloud.bigtable.gaxx.retrying.NonCancellableFuture;
3536
import com.google.common.base.Preconditions;
3637
import com.google.common.collect.ImmutableList;
@@ -235,7 +236,8 @@ private void handleAttemptError(Throwable rpcError) {
235236
FailedMutation failedMutation = FailedMutation.create(origIndex, entryError);
236237
allFailures.add(failedMutation);
237238

238-
if (!failedMutation.getError().isRetryable()) {
239+
if (!ApiExceptions.isRetryable2(failedMutation.getError())
240+
&& !failedMutation.getError().isRetryable()) {
239241
permanentFailures.add(failedMutation);
240242
} else {
241243
// Schedule the mutation entry for the next RPC by adding it to the request builder and

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
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+
*/
16+
package com.google.cloud.bigtable.gaxx.retrying;
17+
18+
import com.google.api.core.InternalApi;
19+
20+
// TODO: move this to gax later
21+
@InternalApi
22+
public class ApiExceptions {
23+
24+
private ApiExceptions() {}
25+
26+
// TODO: this should replace the existing ApiException#isRetryable() method,
27+
// but that cant be done in bigtable, so this lives here for now.
28+
public static boolean isRetryable2(Throwable e) {
29+
if (RetryInfoRetryAlgorithm.extractRetryDelay(e) != null) {
30+
return true;
31+
}
32+
return false;
33+
}
34+
}

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

Lines changed: 5 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -18,16 +18,13 @@
1818
import com.google.api.core.InternalApi;
1919
import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
2020
import com.google.api.gax.retrying.RetryAlgorithm;
21-
import com.google.api.gax.retrying.RetrySettings;
2221
import com.google.api.gax.retrying.ScheduledRetryingExecutor;
2322
import com.google.api.gax.retrying.StreamingRetryAlgorithm;
2423
import com.google.api.gax.rpc.ClientContext;
2524
import com.google.api.gax.rpc.ServerStreamingCallSettings;
2625
import com.google.api.gax.rpc.ServerStreamingCallable;
27-
import com.google.api.gax.rpc.StatusCode;
2826
import com.google.api.gax.rpc.UnaryCallSettings;
2927
import com.google.api.gax.rpc.UnaryCallable;
30-
import java.util.Collection;
3128

3229
// TODO: remove this once ApiResultRetryAlgorithm is added to gax.
3330
/**
@@ -48,23 +45,14 @@ public static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> retrying(
4845

4946
UnaryCallSettings<?, ?> settings = callSettings;
5047

51-
if (areRetriesDisabled(settings.getRetryableCodes(), settings.getRetrySettings())) {
52-
// When retries are disabled, the total timeout can be treated as the rpc timeout.
53-
settings =
54-
settings
55-
.toBuilder()
56-
.setSimpleTimeoutNoRetries(settings.getRetrySettings().getTotalTimeout())
57-
.build();
58-
}
59-
6048
RetryAlgorithm<ResponseT> retryAlgorithm =
6149
new RetryAlgorithm<>(
62-
new ApiResultRetryAlgorithm<ResponseT>(),
50+
new RetryInfoRetryAlgorithm<>(),
6351
new ExponentialRetryAlgorithm(settings.getRetrySettings(), clientContext.getClock()));
64-
ScheduledRetryingExecutor<ResponseT> retryingExecutor =
52+
ScheduledRetryingExecutor<ResponseT> executor =
6553
new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
66-
return new RetryingCallable<>(
67-
clientContext.getDefaultCallContext(), innerCallable, retryingExecutor);
54+
55+
return new RetryingCallable<>(clientContext.getDefaultCallContext(), innerCallable, executor);
6856
}
6957

7058
public static <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT> retrying(
@@ -73,18 +61,10 @@ public static <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT>
7361
ClientContext clientContext) {
7462

7563
ServerStreamingCallSettings<RequestT, ResponseT> settings = callSettings;
76-
if (areRetriesDisabled(settings.getRetryableCodes(), settings.getRetrySettings())) {
77-
// When retries are disabled, the total timeout can be treated as the rpc timeout.
78-
settings =
79-
settings
80-
.toBuilder()
81-
.setSimpleTimeoutNoRetries(settings.getRetrySettings().getTotalTimeout())
82-
.build();
83-
}
8464

8565
StreamingRetryAlgorithm<Void> retryAlgorithm =
8666
new StreamingRetryAlgorithm<>(
87-
new ApiResultRetryAlgorithm<Void>(),
67+
new RetryInfoRetryAlgorithm<>(),
8868
new ExponentialRetryAlgorithm(settings.getRetrySettings(), clientContext.getClock()));
8969

9070
ScheduledRetryingExecutor<Void> retryingExecutor =
@@ -93,11 +73,4 @@ public static <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT>
9373
return new RetryingServerStreamingCallable<>(
9474
innerCallable, retryingExecutor, settings.getResumptionStrategy());
9575
}
96-
97-
private static boolean areRetriesDisabled(
98-
Collection<StatusCode.Code> retryableCodes, RetrySettings retrySettings) {
99-
return retrySettings.getMaxAttempts() == 1
100-
|| retryableCodes.isEmpty()
101-
|| (retrySettings.getMaxAttempts() == 0 && retrySettings.getTotalTimeout().isZero());
102-
}
10376
}

0 commit comments

Comments
 (0)