Skip to content

Commit f78b838

Browse files
fix: use streaming read/query settings for stream retry (#2579)
* fix: use streaming read/query settings for stream retry Use the streaming read/query settings from the SpannerStubSettings to determine when and how to retry a failure halfway a streaming call. * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 051c530 commit f78b838

File tree

9 files changed

+279
-23
lines changed

9 files changed

+279
-23
lines changed

google-cloud-spanner/clirr-ignored-differences.xml

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -422,4 +422,25 @@
422422
<className>com/google/cloud/spanner/spi/v1/SpannerRpc$StreamingCall</className>
423423
<method>com.google.api.gax.rpc.ApiCallContext getCallContext()</method>
424424
</difference>
425+
<!-- (Internal change, propagate streaming retry settings) -->
426+
<difference>
427+
<differenceType>7012</differenceType>
428+
<className>com/google/cloud/spanner/spi/v1/SpannerRpc</className>
429+
<method>com.google.api.gax.retrying.RetrySettings getReadRetrySettings()</method>
430+
</difference>
431+
<difference>
432+
<differenceType>7012</differenceType>
433+
<className>com/google/cloud/spanner/spi/v1/SpannerRpc</className>
434+
<method>com.google.api.gax.retrying.RetrySettings getExecuteQueryRetrySettings()</method>
435+
</difference>
436+
<difference>
437+
<differenceType>7012</differenceType>
438+
<className>com/google/cloud/spanner/spi/v1/SpannerRpc</className>
439+
<method>java.util.Set getReadRetryableCodes()</method>
440+
</difference>
441+
<difference>
442+
<differenceType>7012</differenceType>
443+
<className>com/google/cloud/spanner/spi/v1/SpannerRpc</className>
444+
<method>java.util.Set getExecuteQueryRetryableCodes()</method>
445+
</difference>
425446
</differences>

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -664,7 +664,12 @@ ResultSet executeQueryInternalWithOptions(
664664
getExecuteSqlRequestBuilder(
665665
statement, queryMode, options, /* withTransactionSelector = */ false);
666666
ResumableStreamIterator stream =
667-
new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, SpannerImpl.QUERY, span) {
667+
new ResumableStreamIterator(
668+
MAX_BUFFERED_CHUNKS,
669+
SpannerImpl.QUERY,
670+
span,
671+
rpc.getExecuteQueryRetrySettings(),
672+
rpc.getExecuteQueryRetryableCodes()) {
668673
@Override
669674
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
670675
GrpcStreamIterator stream = new GrpcStreamIterator(statement, prefetchChunks);
@@ -798,7 +803,12 @@ ResultSet readInternalWithOptions(
798803
final int prefetchChunks =
799804
readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
800805
ResumableStreamIterator stream =
801-
new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, SpannerImpl.READ, span) {
806+
new ResumableStreamIterator(
807+
MAX_BUFFERED_CHUNKS,
808+
SpannerImpl.READ,
809+
span,
810+
rpc.getReadRetrySettings(),
811+
rpc.getReadRetryableCodes()) {
802812
@Override
803813
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
804814
GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java

Lines changed: 66 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,10 @@
2424

2525
import com.google.api.client.util.BackOff;
2626
import com.google.api.client.util.ExponentialBackOff;
27+
import com.google.api.gax.grpc.GrpcStatusCode;
2728
import com.google.api.gax.retrying.RetrySettings;
2829
import com.google.api.gax.rpc.ApiCallContext;
30+
import com.google.api.gax.rpc.StatusCode.Code;
2931
import com.google.cloud.ByteArray;
3032
import com.google.cloud.Date;
3133
import com.google.cloud.Timestamp;
@@ -65,6 +67,7 @@
6567
import java.util.LinkedList;
6668
import java.util.List;
6769
import java.util.Objects;
70+
import java.util.Set;
6871
import java.util.concurrent.BlockingQueue;
6972
import java.util.concurrent.CountDownLatch;
7073
import java.util.concurrent.Executor;
@@ -1082,10 +1085,12 @@ public void onError(SpannerException e) {
10821085
@VisibleForTesting
10831086
abstract static class ResumableStreamIterator extends AbstractIterator<PartialResultSet>
10841087
implements CloseableIterator<PartialResultSet> {
1085-
private static final RetrySettings STREAMING_RETRY_SETTINGS =
1088+
private static final RetrySettings DEFAULT_STREAMING_RETRY_SETTINGS =
10861089
SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings();
1090+
private final RetrySettings streamingRetrySettings;
1091+
private final Set<Code> retryableCodes;
10871092
private static final Logger logger = Logger.getLogger(ResumableStreamIterator.class.getName());
1088-
private final BackOff backOff = newBackOff();
1093+
private final BackOff backOff;
10891094
private final LinkedList<PartialResultSet> buffer = new LinkedList<>();
10901095
private final int maxBufferSize;
10911096
private final Span span;
@@ -1099,24 +1104,58 @@ abstract static class ResumableStreamIterator extends AbstractIterator<PartialRe
10991104
*/
11001105
private boolean safeToRetry = true;
11011106

1102-
protected ResumableStreamIterator(int maxBufferSize, String streamName, Span parent) {
1107+
protected ResumableStreamIterator(
1108+
int maxBufferSize,
1109+
String streamName,
1110+
Span parent,
1111+
RetrySettings streamingRetrySettings,
1112+
Set<Code> retryableCodes) {
11031113
checkArgument(maxBufferSize >= 0);
11041114
this.maxBufferSize = maxBufferSize;
11051115
this.span = tracer.spanBuilderWithExplicitParent(streamName, parent).startSpan();
1106-
}
1107-
1108-
private static ExponentialBackOff newBackOff() {
1116+
this.streamingRetrySettings = Preconditions.checkNotNull(streamingRetrySettings);
1117+
this.retryableCodes = Preconditions.checkNotNull(retryableCodes);
1118+
this.backOff = newBackOff();
1119+
}
1120+
1121+
private ExponentialBackOff newBackOff() {
1122+
if (Objects.equals(streamingRetrySettings, DEFAULT_STREAMING_RETRY_SETTINGS)) {
1123+
return new ExponentialBackOff.Builder()
1124+
.setMultiplier(streamingRetrySettings.getRetryDelayMultiplier())
1125+
.setInitialIntervalMillis(
1126+
Math.max(10, (int) streamingRetrySettings.getInitialRetryDelay().toMillis()))
1127+
.setMaxIntervalMillis(
1128+
Math.max(1000, (int) streamingRetrySettings.getMaxRetryDelay().toMillis()))
1129+
.setMaxElapsedTimeMillis(
1130+
Integer.MAX_VALUE) // Prevent Backoff.STOP from getting returned.
1131+
.build();
1132+
}
11091133
return new ExponentialBackOff.Builder()
1110-
.setMultiplier(STREAMING_RETRY_SETTINGS.getRetryDelayMultiplier())
1134+
.setMultiplier(streamingRetrySettings.getRetryDelayMultiplier())
1135+
// All of these values must be > 0.
11111136
.setInitialIntervalMillis(
1112-
Math.max(10, (int) STREAMING_RETRY_SETTINGS.getInitialRetryDelay().toMillis()))
1137+
Math.max(
1138+
1,
1139+
(int)
1140+
Math.min(
1141+
streamingRetrySettings.getInitialRetryDelay().toMillis(),
1142+
Integer.MAX_VALUE)))
11131143
.setMaxIntervalMillis(
1114-
Math.max(1000, (int) STREAMING_RETRY_SETTINGS.getMaxRetryDelay().toMillis()))
1115-
.setMaxElapsedTimeMillis(Integer.MAX_VALUE) // Prevent Backoff.STOP from getting returned.
1144+
Math.max(
1145+
1,
1146+
(int)
1147+
Math.min(
1148+
streamingRetrySettings.getMaxRetryDelay().toMillis(), Integer.MAX_VALUE)))
1149+
.setMaxElapsedTimeMillis(
1150+
Math.max(
1151+
1,
1152+
(int)
1153+
Math.min(
1154+
streamingRetrySettings.getTotalTimeout().toMillis(), Integer.MAX_VALUE)))
11161155
.build();
11171156
}
11181157

1119-
private static void backoffSleep(Context context, BackOff backoff) throws SpannerException {
1158+
private void backoffSleep(Context context, BackOff backoff) throws SpannerException {
11201159
backoffSleep(context, nextBackOffMillis(backoff));
11211160
}
11221161

@@ -1128,7 +1167,7 @@ private static long nextBackOffMillis(BackOff backoff) throws SpannerException {
11281167
}
11291168
}
11301169

1131-
private static void backoffSleep(Context context, long backoffMillis) throws SpannerException {
1170+
private void backoffSleep(Context context, long backoffMillis) throws SpannerException {
11321171
tracer
11331172
.getCurrentSpan()
11341173
.addAnnotation(
@@ -1145,7 +1184,7 @@ private static void backoffSleep(Context context, long backoffMillis) throws Spa
11451184
try {
11461185
if (backoffMillis == BackOff.STOP) {
11471186
// Highly unlikely but we handle it just in case.
1148-
backoffMillis = STREAMING_RETRY_SETTINGS.getMaxRetryDelay().toMillis();
1187+
backoffMillis = streamingRetrySettings.getMaxRetryDelay().toMillis();
11491188
}
11501189
if (latch.await(backoffMillis, TimeUnit.MILLISECONDS)) {
11511190
// Woken by context cancellation.
@@ -1233,19 +1272,20 @@ protected PartialResultSet computeNext() {
12331272
return null;
12341273
}
12351274
}
1236-
} catch (SpannerException e) {
1237-
if (safeToRetry && e.isRetryable()) {
1275+
} catch (SpannerException spannerException) {
1276+
if (safeToRetry && isRetryable(spannerException)) {
12381277
span.addAnnotation(
1239-
"Stream broken. Safe to retry", TraceUtil.getExceptionAnnotations(e));
1240-
logger.log(Level.FINE, "Retryable exception, will sleep and retry", e);
1278+
"Stream broken. Safe to retry",
1279+
TraceUtil.getExceptionAnnotations(spannerException));
1280+
logger.log(Level.FINE, "Retryable exception, will sleep and retry", spannerException);
12411281
// Truncate any items in the buffer before the last retry token.
12421282
while (!buffer.isEmpty() && buffer.getLast().getResumeToken().isEmpty()) {
12431283
buffer.removeLast();
12441284
}
12451285
assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken);
12461286
stream = null;
12471287
try (Scope s = tracer.withSpan(span)) {
1248-
long delay = e.getRetryDelayInMillis();
1288+
long delay = spannerException.getRetryDelayInMillis();
12491289
if (delay != -1) {
12501290
backoffSleep(context, delay);
12511291
} else {
@@ -1256,15 +1296,21 @@ protected PartialResultSet computeNext() {
12561296
continue;
12571297
}
12581298
span.addAnnotation("Stream broken. Not safe to retry");
1259-
TraceUtil.setWithFailure(span, e);
1260-
throw e;
1299+
TraceUtil.setWithFailure(span, spannerException);
1300+
throw spannerException;
12611301
} catch (RuntimeException e) {
12621302
span.addAnnotation("Stream broken. Not safe to retry");
12631303
TraceUtil.setWithFailure(span, e);
12641304
throw e;
12651305
}
12661306
}
12671307
}
1308+
1309+
boolean isRetryable(SpannerException spannerException) {
1310+
return spannerException.isRetryable()
1311+
|| retryableCodes.contains(
1312+
GrpcStatusCode.of(spannerException.getErrorCode().getGrpcStatusCode()).getCode());
1313+
}
12681314
}
12691315

12701316
static double valueProtoToFloat64(com.google.protobuf.Value proto) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import com.google.api.gax.rpc.ResponseObserver;
4646
import com.google.api.gax.rpc.ServerStream;
4747
import com.google.api.gax.rpc.StatusCode;
48+
import com.google.api.gax.rpc.StatusCode.Code;
4849
import com.google.api.gax.rpc.StreamController;
4950
import com.google.api.gax.rpc.TransportChannelProvider;
5051
import com.google.api.gax.rpc.UnaryCallSettings;
@@ -230,6 +231,10 @@ public class GapicSpannerRpc implements SpannerRpc {
230231

231232
private boolean rpcIsClosed;
232233
private final SpannerStub spannerStub;
234+
private final RetrySettings executeQueryRetrySettings;
235+
private final Set<Code> executeQueryRetryableCodes;
236+
private final RetrySettings readRetrySettings;
237+
private final Set<Code> readRetryableCodes;
233238
private final SpannerStub partitionedDmlStub;
234239
private final RetrySettings partitionedDmlRetrySettings;
235240
private final InstanceAdminStub instanceAdminStub;
@@ -368,6 +373,14 @@ public GapicSpannerRpc(final SpannerOptions options) {
368373
.setCredentialsProvider(credentialsProvider)
369374
.setStreamWatchdogProvider(watchdogProvider)
370375
.build());
376+
this.readRetrySettings =
377+
options.getSpannerStubSettings().streamingReadSettings().getRetrySettings();
378+
this.readRetryableCodes =
379+
options.getSpannerStubSettings().streamingReadSettings().getRetryableCodes();
380+
this.executeQueryRetrySettings =
381+
options.getSpannerStubSettings().executeStreamingSqlSettings().getRetrySettings();
382+
this.executeQueryRetryableCodes =
383+
options.getSpannerStubSettings().executeStreamingSqlSettings().getRetryableCodes();
371384
partitionedDmlRetrySettings =
372385
options
373386
.getSpannerStubSettings()
@@ -472,6 +485,10 @@ public <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUnaryCalla
472485
this.databaseAdminStub = null;
473486
this.instanceAdminStub = null;
474487
this.spannerStub = null;
488+
this.readRetrySettings = null;
489+
this.readRetryableCodes = null;
490+
this.executeQueryRetrySettings = null;
491+
this.executeQueryRetryableCodes = null;
475492
this.partitionedDmlStub = null;
476493
this.databaseAdminStubSettings = null;
477494
this.spannerWatchdog = null;
@@ -1585,6 +1602,16 @@ public ApiFuture<Empty> asyncDeleteSession(String sessionName, @Nullable Map<Opt
15851602
return spannerStub.deleteSessionCallable().futureCall(request, context);
15861603
}
15871604

1605+
@Override
1606+
public RetrySettings getReadRetrySettings() {
1607+
return readRetrySettings;
1608+
}
1609+
1610+
@Override
1611+
public Set<Code> getReadRetryableCodes() {
1612+
return readRetryableCodes;
1613+
}
1614+
15881615
@Override
15891616
public StreamingCall read(
15901617
ReadRequest request,
@@ -1599,6 +1626,16 @@ public StreamingCall read(
15991626
return new GrpcStreamingCall(context, responseObserver.getController());
16001627
}
16011628

1629+
@Override
1630+
public RetrySettings getExecuteQueryRetrySettings() {
1631+
return executeQueryRetrySettings;
1632+
}
1633+
1634+
@Override
1635+
public Set<Code> getExecuteQueryRetryableCodes() {
1636+
return executeQueryRetryableCodes;
1637+
}
1638+
16021639
@Override
16031640
public ResultSet executeQuery(
16041641
ExecuteSqlRequest request, @Nullable Map<Option, ?> options, boolean routeToLeader) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,14 @@
2222
import com.google.api.gax.retrying.RetrySettings;
2323
import com.google.api.gax.rpc.ApiCallContext;
2424
import com.google.api.gax.rpc.ServerStream;
25+
import com.google.api.gax.rpc.StatusCode.Code;
2526
import com.google.cloud.ServiceRpc;
2627
import com.google.cloud.spanner.BackupId;
2728
import com.google.cloud.spanner.Restore;
2829
import com.google.cloud.spanner.SpannerException;
2930
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
3031
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub;
32+
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
3133
import com.google.common.collect.ImmutableList;
3234
import com.google.iam.v1.GetPolicyOptions;
3335
import com.google.iam.v1.Policy;
@@ -53,6 +55,7 @@
5355
import com.google.spanner.v1.*;
5456
import java.util.List;
5557
import java.util.Map;
58+
import java.util.Set;
5659
import javax.annotation.Nullable;
5760
import org.threeten.bp.Duration;
5861

@@ -337,6 +340,16 @@ Session createSession(
337340
ApiFuture<Empty> asyncDeleteSession(String sessionName, @Nullable Map<Option, ?> options)
338341
throws SpannerException;
339342

343+
/** Returns the retry settings for streaming read operations. */
344+
default RetrySettings getReadRetrySettings() {
345+
return SpannerStubSettings.newBuilder().streamingReadSettings().getRetrySettings();
346+
}
347+
348+
/** Returns the retryable codes for streaming read operations. */
349+
default Set<Code> getReadRetryableCodes() {
350+
return SpannerStubSettings.newBuilder().streamingReadSettings().getRetryableCodes();
351+
}
352+
340353
/**
341354
* Performs a streaming read.
342355
*
@@ -351,6 +364,16 @@ StreamingCall read(
351364
@Nullable Map<Option, ?> options,
352365
boolean routeToLeader);
353366

367+
/** Returns the retry settings for streaming query operations. */
368+
default RetrySettings getExecuteQueryRetrySettings() {
369+
return SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings();
370+
}
371+
372+
/** Returns the retryable codes for streaming query operations. */
373+
default Set<Code> getExecuteQueryRetryableCodes() {
374+
return SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes();
375+
}
376+
354377
/**
355378
* Executes a query.
356379
*

0 commit comments

Comments
 (0)