Skip to content

Commit 48f92e3

Browse files
olavloiteolavloite
authored
fix: query could hang transaction if ResultSet#next() is not called (#643)
If the first statement of a read/write transaction was a query or a read operation, and the application would not call ResultSet#next() on the return result, the transaction would hang indefinetely as the query would be marked as the one that should initiate the transaction (inline the BeginTransaction option). The query would however never be executed, as the actual query execution is deferred until the first call to ResultSet#next(). Fixes #641
1 parent 7584baa commit 48f92e3

File tree

7 files changed

+171
-31
lines changed

7 files changed

+171
-31
lines changed

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

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -608,33 +608,33 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder(Iterable<Stateme
608608

609609
ResultSet executeQueryInternalWithOptions(
610610
final Statement statement,
611-
com.google.spanner.v1.ExecuteSqlRequest.QueryMode queryMode,
611+
final com.google.spanner.v1.ExecuteSqlRequest.QueryMode queryMode,
612612
Options options,
613-
ByteString partitionToken) {
613+
final ByteString partitionToken) {
614614
beforeReadOrQuery();
615-
final ExecuteSqlRequest.Builder request = getExecuteSqlRequestBuilder(statement, queryMode);
616-
if (partitionToken != null) {
617-
request.setPartitionToken(partitionToken);
618-
}
619615
final int prefetchChunks =
620616
options.hasPrefetchChunks() ? options.prefetchChunks() : defaultPrefetchChunks;
621617
ResumableStreamIterator stream =
622618
new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, SpannerImpl.QUERY, span) {
623619
@Override
624620
CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
625621
GrpcStreamIterator stream = new GrpcStreamIterator(statement, prefetchChunks);
622+
final ExecuteSqlRequest.Builder request =
623+
getExecuteSqlRequestBuilder(statement, queryMode);
624+
if (partitionToken != null) {
625+
request.setPartitionToken(partitionToken);
626+
}
626627
if (resumeToken != null) {
627628
request.setResumeToken(resumeToken);
628629
}
629630
SpannerRpc.StreamingCall call =
630631
rpc.executeQuery(request.build(), stream.consumer(), session.getOptions());
631632
call.request(prefetchChunks);
632-
stream.setCall(call);
633+
stream.setCall(call, request.hasTransaction() && request.getTransaction().hasBegin());
633634
return stream;
634635
}
635636
};
636-
return new GrpcResultSet(
637-
stream, this, request.hasTransaction() && request.getTransaction().hasBegin());
637+
return new GrpcResultSet(stream, this);
638638
}
639639

640640
/**
@@ -723,10 +723,6 @@ ResultSet readInternalWithOptions(
723723
if (index != null) {
724724
builder.setIndex(index);
725725
}
726-
TransactionSelector selector = getTransactionSelector();
727-
if (selector != null) {
728-
builder.setTransaction(selector);
729-
}
730726
if (partitionToken != null) {
731727
builder.setPartitionToken(partitionToken);
732728
}
@@ -740,15 +736,18 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
740736
if (resumeToken != null) {
741737
builder.setResumeToken(resumeToken);
742738
}
739+
TransactionSelector selector = getTransactionSelector();
740+
if (selector != null) {
741+
builder.setTransaction(selector);
742+
}
743743
SpannerRpc.StreamingCall call =
744744
rpc.read(builder.build(), stream.consumer(), session.getOptions());
745745
call.request(prefetchChunks);
746-
stream.setCall(call);
746+
stream.setCall(call, selector != null && selector.hasBegin());
747747
return stream;
748748
}
749749
};
750-
GrpcResultSet resultSet =
751-
new GrpcResultSet(stream, this, selector != null && selector.hasBegin());
750+
GrpcResultSet resultSet = new GrpcResultSet(stream, this);
752751
return resultSet;
753752
}
754753

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

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -91,17 +91,14 @@ interface Listener {
9191
static class GrpcResultSet extends AbstractResultSet<List<Object>> {
9292
private final GrpcValueIterator iterator;
9393
private final Listener listener;
94-
private final boolean beginTransaction;
9594
private GrpcStruct currRow;
9695
private SpannerException error;
9796
private ResultSetStats statistics;
9897
private boolean closed;
9998

100-
GrpcResultSet(
101-
CloseableIterator<PartialResultSet> iterator, Listener listener, boolean beginTransaction) {
99+
GrpcResultSet(CloseableIterator<PartialResultSet> iterator, Listener listener) {
102100
this.iterator = new GrpcValueIterator(iterator);
103101
this.listener = listener;
104-
this.beginTransaction = beginTransaction;
105102
}
106103

107104
@Override
@@ -130,7 +127,7 @@ public boolean next() throws SpannerException {
130127
}
131128
return hasNext;
132129
} catch (SpannerException e) {
133-
throw yieldError(e, beginTransaction && currRow == null);
130+
throw yieldError(e, iterator.isWithBeginTransaction() && currRow == null);
134131
}
135132
}
136133

@@ -297,6 +294,10 @@ void close(@Nullable String message) {
297294
stream.close(message);
298295
}
299296

297+
boolean isWithBeginTransaction() {
298+
return stream.isWithBeginTransaction();
299+
}
300+
300301
/** @param a is a mutable list and b will be concatenated into a. */
301302
private void concatLists(List<com.google.protobuf.Value> a, List<com.google.protobuf.Value> b) {
302303
if (a.size() == 0 || b.size() == 0) {
@@ -760,6 +761,8 @@ interface CloseableIterator<T> extends Iterator<T> {
760761
* @param message a message to include in the final RPC status
761762
*/
762763
void close(@Nullable String message);
764+
765+
boolean isWithBeginTransaction();
763766
}
764767

765768
/** Adapts a streaming read/query call into an iterator over partial result sets. */
@@ -774,6 +777,7 @@ static class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
774777
private final Statement statement;
775778

776779
private SpannerRpc.StreamingCall call;
780+
private boolean withBeginTransaction;
777781
private SpannerException error;
778782

779783
@VisibleForTesting
@@ -792,8 +796,9 @@ protected final SpannerRpc.ResultStreamConsumer consumer() {
792796
return consumer;
793797
}
794798

795-
public void setCall(SpannerRpc.StreamingCall call) {
799+
public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) {
796800
this.call = call;
801+
this.withBeginTransaction = withBeginTransaction;
797802
}
798803

799804
@Override
@@ -803,6 +808,11 @@ public void close(@Nullable String message) {
803808
}
804809
}
805810

811+
@Override
812+
public boolean isWithBeginTransaction() {
813+
return withBeginTransaction;
814+
}
815+
806816
@Override
807817
protected final PartialResultSet computeNext() {
808818
PartialResultSet next;
@@ -873,8 +883,8 @@ public void onError(SpannerException e) {
873883

874884
// Visible only for testing.
875885
@VisibleForTesting
876-
void setCall(SpannerRpc.StreamingCall call) {
877-
GrpcStreamIterator.this.setCall(call);
886+
void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) {
887+
GrpcStreamIterator.this.setCall(call, withBeginTransaction);
878888
}
879889
}
880890
}
@@ -987,6 +997,11 @@ public void close(@Nullable String message) {
987997
}
988998
}
989999

1000+
@Override
1001+
public boolean isWithBeginTransaction() {
1002+
return stream != null && stream.isWithBeginTransaction();
1003+
}
1004+
9901005
@Override
9911006
protected PartialResultSet computeNext() {
9921007
Context context = Context.current();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -267,7 +267,7 @@ ApiFuture<Timestamp> commitAsync() {
267267
final SettableApiFuture<Void> finishOps;
268268
CommitRequest.Builder builder = CommitRequest.newBuilder().setSession(session.getName());
269269
synchronized (lock) {
270-
if (transactionIdFuture == null && transactionId == null) {
270+
if (transactionIdFuture == null && transactionId == null && runningAsyncOperations == 0) {
271271
finishOps = SettableApiFuture.create();
272272
createTxnAsync(finishOps);
273273
} else {

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -75,13 +75,14 @@ public void cancel(@Nullable String message) {}
7575

7676
@Override
7777
public void request(int numMessages) {}
78-
});
78+
},
79+
false);
7980
consumer = stream.consumer();
80-
resultSet = new AbstractResultSet.GrpcResultSet(stream, new NoOpListener(), false);
81+
resultSet = new AbstractResultSet.GrpcResultSet(stream, new NoOpListener());
8182
}
8283

8384
public AbstractResultSet.GrpcResultSet resultSetWithMode(QueryMode queryMode) {
84-
return new AbstractResultSet.GrpcResultSet(stream, new NoOpListener(), false);
85+
return new AbstractResultSet.GrpcResultSet(stream, new NoOpListener());
8586
}
8687

8788
@Test
@@ -642,7 +643,7 @@ public com.google.protobuf.Value apply(@Nullable Value input) {
642643

643644
private void verifySerialization(
644645
Function<Value, com.google.protobuf.Value> protoFn, Value... values) {
645-
resultSet = new AbstractResultSet.GrpcResultSet(stream, new NoOpListener(), false);
646+
resultSet = new AbstractResultSet.GrpcResultSet(stream, new NoOpListener());
646647
PartialResultSet.Builder builder = PartialResultSet.newBuilder();
647648
List<Type.StructField> types = new ArrayList<>();
648649
for (Value value : values) {

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

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.google.cloud.spanner;
1818

19+
import static com.google.cloud.spanner.MockSpannerTestUtil.SELECT1;
1920
import static com.google.common.truth.Truth.assertThat;
2021
import static org.junit.Assert.fail;
2122

@@ -65,6 +66,7 @@
6566
import java.util.concurrent.Future;
6667
import java.util.concurrent.ScheduledExecutorService;
6768
import java.util.concurrent.ScheduledThreadPoolExecutor;
69+
import java.util.concurrent.TimeoutException;
6870
import java.util.concurrent.atomic.AtomicBoolean;
6971
import org.junit.After;
7072
import org.junit.AfterClass;
@@ -1139,6 +1141,123 @@ public ApiFuture<Long> apply(TransactionContext txn, Long input)
11391141
assertThat(countTransactionsStarted()).isEqualTo(1);
11401142
}
11411143

1144+
@Test
1145+
public void queryWithoutNext() {
1146+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
1147+
assertThat(
1148+
client
1149+
.readWriteTransaction()
1150+
.run(
1151+
new TransactionCallable<Long>() {
1152+
@Override
1153+
public Long run(TransactionContext transaction) throws Exception {
1154+
// This will not actually send an RPC, so it will also not request a
1155+
// transaction.
1156+
transaction.executeQuery(SELECT1);
1157+
return transaction.executeUpdate(UPDATE_STATEMENT);
1158+
}
1159+
}))
1160+
.isEqualTo(UPDATE_COUNT);
1161+
assertThat(mockSpanner.countRequestsOfType(BeginTransactionRequest.class)).isEqualTo(0L);
1162+
assertThat(mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)).isEqualTo(1L);
1163+
assertThat(countTransactionsStarted()).isEqualTo(1);
1164+
}
1165+
1166+
@Test
1167+
public void queryAsyncWithoutCallback() {
1168+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
1169+
assertThat(
1170+
client
1171+
.readWriteTransaction()
1172+
.run(
1173+
new TransactionCallable<Long>() {
1174+
@Override
1175+
public Long run(TransactionContext transaction) throws Exception {
1176+
transaction.executeQueryAsync(SELECT1);
1177+
return transaction.executeUpdate(UPDATE_STATEMENT);
1178+
}
1179+
}))
1180+
.isEqualTo(UPDATE_COUNT);
1181+
assertThat(mockSpanner.countRequestsOfType(BeginTransactionRequest.class)).isEqualTo(0L);
1182+
assertThat(mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)).isEqualTo(1L);
1183+
assertThat(countTransactionsStarted()).isEqualTo(1);
1184+
}
1185+
1186+
@Test
1187+
public void readWithoutNext() {
1188+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
1189+
assertThat(
1190+
client
1191+
.readWriteTransaction()
1192+
.run(
1193+
new TransactionCallable<Long>() {
1194+
@Override
1195+
public Long run(TransactionContext transaction) throws Exception {
1196+
transaction.read("FOO", KeySet.all(), Arrays.asList("ID"));
1197+
return transaction.executeUpdate(UPDATE_STATEMENT);
1198+
}
1199+
}))
1200+
.isEqualTo(UPDATE_COUNT);
1201+
assertThat(mockSpanner.countRequestsOfType(BeginTransactionRequest.class)).isEqualTo(0L);
1202+
assertThat(mockSpanner.countRequestsOfType(ReadRequest.class)).isEqualTo(0L);
1203+
assertThat(mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)).isEqualTo(1L);
1204+
assertThat(countTransactionsStarted()).isEqualTo(1);
1205+
}
1206+
1207+
@Test
1208+
public void readAsyncWithoutCallback() {
1209+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
1210+
assertThat(
1211+
client
1212+
.readWriteTransaction()
1213+
.run(
1214+
new TransactionCallable<Long>() {
1215+
@Override
1216+
public Long run(TransactionContext transaction) throws Exception {
1217+
transaction.readAsync("FOO", KeySet.all(), Arrays.asList("ID"));
1218+
return transaction.executeUpdate(UPDATE_STATEMENT);
1219+
}
1220+
}))
1221+
.isEqualTo(UPDATE_COUNT);
1222+
assertThat(mockSpanner.countRequestsOfType(BeginTransactionRequest.class)).isEqualTo(0L);
1223+
assertThat(mockSpanner.countRequestsOfType(ReadRequest.class)).isEqualTo(0L);
1224+
assertThat(mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)).isEqualTo(1L);
1225+
assertThat(countTransactionsStarted()).isEqualTo(1);
1226+
}
1227+
1228+
@Test
1229+
public void query_ThenUpdate_ThenConsumeResultSet()
1230+
throws InterruptedException, TimeoutException {
1231+
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
1232+
assertThat(
1233+
client
1234+
.readWriteTransaction()
1235+
.run(
1236+
new TransactionCallable<Long>() {
1237+
@Override
1238+
public Long run(TransactionContext transaction) throws Exception {
1239+
ResultSet rs = transaction.executeQuery(SELECT1);
1240+
long updateCount = transaction.executeUpdate(UPDATE_STATEMENT);
1241+
// Consume the result set.
1242+
while (rs.next()) {}
1243+
return updateCount;
1244+
}
1245+
}))
1246+
.isEqualTo(UPDATE_COUNT);
1247+
// The update statement should start the transaction, and the query should use the transaction
1248+
// id returned by the update.
1249+
assertThat(mockSpanner.countRequestsOfType(BeginTransactionRequest.class)).isEqualTo(0L);
1250+
assertThat(mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)).isEqualTo(2L);
1251+
assertThat(countTransactionsStarted()).isEqualTo(1);
1252+
List<AbstractMessage> requests = mockSpanner.getRequests();
1253+
requests = requests.subList(requests.size() - 3, requests.size());
1254+
assertThat(requests.get(0)).isInstanceOf(ExecuteSqlRequest.class);
1255+
assertThat(((ExecuteSqlRequest) requests.get(0)).getSql()).isEqualTo(UPDATE_STATEMENT.getSql());
1256+
assertThat(requests.get(1)).isInstanceOf(ExecuteSqlRequest.class);
1257+
assertThat(((ExecuteSqlRequest) requests.get(1)).getSql()).isEqualTo(SELECT1.getSql());
1258+
assertThat(requests.get(2)).isInstanceOf(CommitRequest.class);
1259+
}
1260+
11421261
private int countRequests(Class<? extends AbstractMessage> requestType) {
11431262
int count = 0;
11441263
for (AbstractMessage msg : mockSpanner.getRequests()) {

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,9 +117,10 @@ public void cancel(@Nullable String message) {}
117117

118118
@Override
119119
public void request(int numMessages) {}
120-
});
120+
},
121+
false);
121122
consumer = stream.consumer();
122-
resultSet = new AbstractResultSet.GrpcResultSet(stream, new NoOpListener(), false);
123+
resultSet = new AbstractResultSet.GrpcResultSet(stream, new NoOpListener());
123124

124125
JSONArray chunks = testCase.getJSONArray("chunks");
125126
JSONObject expectedResult = testCase.getJSONObject("result");

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,11 @@ protected PartialResultSet computeNext() {
116116
public void close(@Nullable String message) {
117117
stream.close();
118118
}
119+
120+
@Override
121+
public boolean isWithBeginTransaction() {
122+
return false;
123+
}
119124
}
120125

121126
Starter starter = Mockito.mock(Starter.class);

0 commit comments

Comments
 (0)