Skip to content

Commit c7dc6e6

Browse files
olavloiteolavloite
authored
fix: UNAVAILABLE error on first query could cause transaction to get stuck (#807)
If the first query or read operation of a read/write transaction would return UNAVAILABLE for the first element of the result stream, the transaction could get stuck. This was caused by the internal retry mechanism that would wait for the initial attempt to return a transaction, which was never returned as the UNAVAILABLE exception was internally handled by the result stream iterator. Fixes #799
1 parent 557e761 commit c7dc6e6

File tree

5 files changed

+160
-13
lines changed

5 files changed

+160
-13
lines changed

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

+20-7
Original file line numberDiff line numberDiff line change
@@ -558,7 +558,7 @@ QueryOptions buildQueryOptions(QueryOptions requestOptions) {
558558
}
559559

560560
ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(
561-
Statement statement, QueryMode queryMode, Options options) {
561+
Statement statement, QueryMode queryMode, Options options, boolean withTransactionSelector) {
562562
ExecuteSqlRequest.Builder builder =
563563
ExecuteSqlRequest.newBuilder()
564564
.setSql(statement.getSql())
@@ -572,9 +572,11 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(
572572
builder.putParamTypes(param.getKey(), param.getValue().getType().toProto());
573573
}
574574
}
575-
TransactionSelector selector = getTransactionSelector();
576-
if (selector != null) {
577-
builder.setTransaction(selector);
575+
if (withTransactionSelector) {
576+
TransactionSelector selector = getTransactionSelector();
577+
if (selector != null) {
578+
builder.setTransaction(selector);
579+
}
578580
}
579581
builder.setSeqno(getSeqNo());
580582
builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions()));
@@ -619,18 +621,26 @@ ResultSet executeQueryInternalWithOptions(
619621
beforeReadOrQuery();
620622
final int prefetchChunks =
621623
options.hasPrefetchChunks() ? options.prefetchChunks() : defaultPrefetchChunks;
624+
final ExecuteSqlRequest.Builder request =
625+
getExecuteSqlRequestBuilder(
626+
statement, queryMode, options, /* withTransactionSelector = */ false);
622627
ResumableStreamIterator stream =
623628
new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, SpannerImpl.QUERY, span) {
624629
@Override
625630
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
626631
GrpcStreamIterator stream = new GrpcStreamIterator(statement, prefetchChunks);
627-
final ExecuteSqlRequest.Builder request =
628-
getExecuteSqlRequestBuilder(statement, queryMode, options);
629632
if (partitionToken != null) {
630633
request.setPartitionToken(partitionToken);
631634
}
635+
TransactionSelector selector = null;
632636
if (resumeToken != null) {
633637
request.setResumeToken(resumeToken);
638+
selector = getTransactionSelector();
639+
} else if (!request.hasTransaction()) {
640+
selector = getTransactionSelector();
641+
}
642+
if (selector != null) {
643+
request.setTransaction(selector);
634644
}
635645
SpannerRpc.StreamingCall call =
636646
rpc.executeQuery(request.build(), stream.consumer(), session.getOptions());
@@ -738,10 +748,13 @@ ResultSet readInternalWithOptions(
738748
@Override
739749
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
740750
GrpcStreamIterator stream = new GrpcStreamIterator(prefetchChunks);
751+
TransactionSelector selector = null;
741752
if (resumeToken != null) {
742753
builder.setResumeToken(resumeToken);
754+
selector = getTransactionSelector();
755+
} else if (!builder.hasTransaction()) {
756+
selector = getTransactionSelector();
743757
}
744-
TransactionSelector selector = getTransactionSelector();
745758
if (selector != null) {
746759
builder.setTransaction(selector);
747760
}

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

+1
Original file line numberDiff line numberDiff line change
@@ -1080,6 +1080,7 @@ protected PartialResultSet computeNext() {
10801080
backoffSleep(context, backOff);
10811081
}
10821082
}
1083+
10831084
continue;
10841085
}
10851086
span.addAnnotation("Stream broken. Not safe to retry");

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java

+8-2
Original file line numberDiff line numberDiff line change
@@ -583,7 +583,10 @@ public long executeUpdate(Statement statement, UpdateOption... options) {
583583
beforeReadOrQuery();
584584
final ExecuteSqlRequest.Builder builder =
585585
getExecuteSqlRequestBuilder(
586-
statement, QueryMode.NORMAL, Options.fromUpdateOptions(options));
586+
statement,
587+
QueryMode.NORMAL,
588+
Options.fromUpdateOptions(options),
589+
/* withTransactionSelector = */ true);
587590
try {
588591
com.google.spanner.v1.ResultSet resultSet =
589592
rpc.executeQuery(builder.build(), session.getOptions());
@@ -608,7 +611,10 @@ public ApiFuture<Long> executeUpdateAsync(Statement statement, UpdateOption... o
608611
beforeReadOrQuery();
609612
final ExecuteSqlRequest.Builder builder =
610613
getExecuteSqlRequestBuilder(
611-
statement, QueryMode.NORMAL, Options.fromUpdateOptions(options));
614+
statement,
615+
QueryMode.NORMAL,
616+
Options.fromUpdateOptions(options),
617+
/* withTransactionSelector = */ true);
612618
final ApiFuture<com.google.spanner.v1.ResultSet> resultSet;
613619
try {
614620
// Register the update as an async operation that must finish before the transaction may

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java

+6-2
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,10 @@ public void executeSqlRequestBuilderWithoutQueryOptions() {
9090
ExecuteSqlRequest request =
9191
context
9292
.getExecuteSqlRequestBuilder(
93-
Statement.of("SELECT FOO FROM BAR"), QueryMode.NORMAL, Options.fromQueryOptions())
93+
Statement.of("SELECT FOO FROM BAR"),
94+
QueryMode.NORMAL,
95+
Options.fromQueryOptions(),
96+
true)
9497
.build();
9598
assertThat(request.getSql()).isEqualTo("SELECT FOO FROM BAR");
9699
assertThat(request.getQueryOptions()).isEqualTo(defaultQueryOptions);
@@ -105,7 +108,8 @@ public void executeSqlRequestBuilderWithQueryOptions() {
105108
.withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("2.0").build())
106109
.build(),
107110
QueryMode.NORMAL,
108-
Options.fromQueryOptions())
111+
Options.fromQueryOptions(),
112+
true)
109113
.build();
110114
assertThat(request.getSql()).isEqualTo("SELECT FOO FROM BAR");
111115
assertThat(request.getQueryOptions().getOptimizerVersion()).isEqualTo("2.0");

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java

+125-2
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,130 @@ public Long run(TransactionContext transaction) throws Exception {
257257
assertThat(countTransactionsStarted()).isEqualTo(2);
258258
}
259259

260+
@Test
261+
public void testInlinedBeginFirstUpdateAborts() {
262+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
263+
long updateCount =
264+
client
265+
.readWriteTransaction()
266+
.run(
267+
new TransactionCallable<Long>() {
268+
boolean firstAttempt = true;
269+
270+
@Override
271+
public Long run(TransactionContext transaction) throws Exception {
272+
if (firstAttempt) {
273+
firstAttempt = false;
274+
mockSpanner.putStatementResult(
275+
StatementResult.exception(
276+
UPDATE_STATEMENT,
277+
mockSpanner.createAbortedException(
278+
ByteString.copyFromUtf8("some-tx"))));
279+
} else {
280+
mockSpanner.putStatementResult(
281+
StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT));
282+
}
283+
return transaction.executeUpdate(UPDATE_STATEMENT);
284+
}
285+
});
286+
assertThat(updateCount).isEqualTo(UPDATE_COUNT);
287+
assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(1);
288+
assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(2);
289+
assertThat(countRequests(CommitRequest.class)).isEqualTo(1);
290+
}
291+
292+
@Test
293+
public void testInlinedBeginFirstQueryAborts() {
294+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
295+
long updateCount =
296+
client
297+
.readWriteTransaction()
298+
.run(
299+
new TransactionCallable<Long>() {
300+
boolean firstAttempt = true;
301+
302+
@Override
303+
public Long run(TransactionContext transaction) throws Exception {
304+
if (firstAttempt) {
305+
firstAttempt = false;
306+
mockSpanner.putStatementResult(
307+
StatementResult.exception(
308+
SELECT1,
309+
mockSpanner.createAbortedException(
310+
ByteString.copyFromUtf8("some-tx"))));
311+
} else {
312+
mockSpanner.putStatementResult(
313+
StatementResult.query(SELECT1, SELECT1_RESULTSET));
314+
}
315+
try (ResultSet rs = transaction.executeQuery(SELECT1)) {
316+
while (rs.next()) {
317+
return rs.getLong(0);
318+
}
319+
}
320+
return 0L;
321+
}
322+
});
323+
assertThat(updateCount).isEqualTo(1L);
324+
assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(1);
325+
assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(2);
326+
assertThat(countRequests(CommitRequest.class)).isEqualTo(1);
327+
}
328+
329+
@Test
330+
public void testInlinedBeginFirstQueryReturnsUnavailable() {
331+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
332+
mockSpanner.setExecuteStreamingSqlExecutionTime(
333+
SimulatedExecutionTime.ofStreamException(Status.UNAVAILABLE.asRuntimeException(), 0));
334+
long value =
335+
client
336+
.readWriteTransaction()
337+
.run(
338+
new TransactionCallable<Long>() {
339+
@Override
340+
public Long run(TransactionContext transaction) throws Exception {
341+
// The first attempt will return UNAVAILABLE and retry internally.
342+
try (ResultSet rs = transaction.executeQuery(SELECT1)) {
343+
while (rs.next()) {
344+
return rs.getLong(0);
345+
}
346+
}
347+
return 0L;
348+
}
349+
});
350+
assertThat(value).isEqualTo(1L);
351+
assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
352+
assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(2);
353+
assertThat(countRequests(CommitRequest.class)).isEqualTo(1);
354+
}
355+
356+
@Test
357+
public void testInlinedBeginFirstReadReturnsUnavailable() {
358+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
359+
mockSpanner.setStreamingReadExecutionTime(
360+
SimulatedExecutionTime.ofStreamException(Status.UNAVAILABLE.asRuntimeException(), 0));
361+
long value =
362+
client
363+
.readWriteTransaction()
364+
.run(
365+
new TransactionCallable<Long>() {
366+
@Override
367+
public Long run(TransactionContext transaction) throws Exception {
368+
// The first attempt will return UNAVAILABLE and retry internally.
369+
try (ResultSet rs =
370+
transaction.read("FOO", KeySet.all(), Arrays.asList("ID"))) {
371+
while (rs.next()) {
372+
return rs.getLong(0);
373+
}
374+
}
375+
return 0L;
376+
}
377+
});
378+
assertThat(value).isEqualTo(1L);
379+
assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
380+
assertThat(countRequests(ReadRequest.class)).isEqualTo(2);
381+
assertThat(countRequests(CommitRequest.class)).isEqualTo(1);
382+
}
383+
260384
@Test
261385
public void testInlinedBeginTxWithQuery() {
262386
DatabaseClient client =
@@ -285,8 +409,7 @@ public Long run(TransactionContext transaction) throws Exception {
285409

286410
@Test
287411
public void testInlinedBeginTxWithRead() {
288-
DatabaseClient client =
289-
spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
412+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
290413
long updateCount =
291414
client
292415
.readWriteTransaction()

0 commit comments

Comments
 (0)