Skip to content

Commit ffe2560

Browse files
fornwallfornwall
committed
fix(spanner): cancel the in-flight Commit RPC when the commit is abandoned
Signed-off-by: Fredrik Fornwall <fredrik@fornwall.net>
1 parent 2f915ab commit ffe2560

2 files changed

Lines changed: 180 additions & 4 deletions

File tree

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

Lines changed: 54 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,14 @@ public void removeListener(Runnable listener) {
187187
@GuardedBy("precommitTokenLock")
188188
private MultiplexedSessionPrecommitToken latestPrecommitToken;
189189

190+
private final Object commitCancellationLock = new Object();
191+
192+
@GuardedBy("commitCancellationLock")
193+
private boolean commitCancelled;
194+
195+
@GuardedBy("commitCancellationLock")
196+
private ApiFuture<com.google.spanner.v1.CommitResponse> inFlightCommitFuture;
197+
190198
@GuardedBy("lock")
191199
private volatile SettableApiFuture<Void> finishedAsyncOperations = SettableApiFuture.create();
192200

@@ -365,9 +373,7 @@ void commit() {
365373
rpc.getCommitRetrySettings().getTotalTimeout().getSeconds() + 5,
366374
TimeUnit.SECONDS);
367375
} catch (InterruptedException | TimeoutException e) {
368-
if (commitFuture != null) {
369-
commitFuture.cancel(true);
370-
}
376+
cancelInFlightCommit();
371377
if (e instanceof InterruptedException) {
372378
throw SpannerExceptionFactory.propagateInterrupt((InterruptedException) e);
373379
} else {
@@ -378,7 +384,36 @@ void commit() {
378384
}
379385
}
380386

381-
volatile ApiFuture<CommitResponse> commitFuture;
387+
private void cancelInFlightCommit() {
388+
ApiFuture<com.google.spanner.v1.CommitResponse> commitFuture;
389+
synchronized (commitCancellationLock) {
390+
commitCancelled = true;
391+
commitFuture = inFlightCommitFuture;
392+
}
393+
if (commitFuture != null) {
394+
commitFuture.cancel(true);
395+
}
396+
}
397+
398+
private void publishOrCancelInFlightCommit(
399+
ApiFuture<com.google.spanner.v1.CommitResponse> commitFuture) {
400+
boolean cancelled;
401+
synchronized (commitCancellationLock) {
402+
cancelled = commitCancelled;
403+
if (!cancelled) {
404+
inFlightCommitFuture = commitFuture;
405+
}
406+
}
407+
if (cancelled) {
408+
commitFuture.cancel(true);
409+
}
410+
}
411+
412+
private boolean isCommitCancelled() {
413+
synchronized (commitCancellationLock) {
414+
return commitCancelled;
415+
}
416+
}
382417

383418
ApiFuture<CommitResponse> commitAsync() {
384419
close();
@@ -430,6 +465,13 @@ ApiFuture<CommitResponse> commitAsync() {
430465
new CommitRunnable(
431466
res, finishOps, builder, /* retryAttemptDueToCommitProtocolExtension= */ false),
432467
MoreExecutors.directExecutor());
468+
res.addListener(
469+
() -> {
470+
if (res.isCancelled()) {
471+
cancelInFlightCommit();
472+
}
473+
},
474+
MoreExecutors.directExecutor());
433475
return res;
434476
}
435477

@@ -484,13 +526,21 @@ public void run() {
484526
"Retrying commit operation with a new precommit token obtained from the previous"
485527
+ " CommitResponse");
486528
}
529+
if (isCommitCancelled()) {
530+
res.setException(
531+
newSpannerException(
532+
ErrorCode.CANCELLED,
533+
"The commit was cancelled before the Commit RPC was sent"));
534+
return;
535+
}
487536
final CommitRequest commitRequest = requestBuilder.build();
488537
span.addAnnotation("Starting Commit");
489538
final ApiFuture<com.google.spanner.v1.CommitResponse> commitFuture;
490539
final ISpan opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span);
491540
try (IScope ignore = tracer.withSpan(opSpan)) {
492541
commitFuture = rpc.commitAsync(commitRequest, getTransactionChannelHint());
493542
}
543+
publishOrCancelInFlightCommit(commitFuture);
494544
session.markUsed(clock.instant());
495545
commitFuture.addListener(
496546
() -> {

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

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import static org.junit.Assert.assertArrayEquals;
2121
import static org.junit.Assert.assertEquals;
2222
import static org.junit.Assert.assertThrows;
23+
import static org.junit.Assert.assertTrue;
2324
import static org.mockito.ArgumentMatchers.any;
2425
import static org.mockito.ArgumentMatchers.eq;
2526
import static org.mockito.Mockito.doThrow;
@@ -29,7 +30,9 @@
2930
import static org.mockito.Mockito.verify;
3031
import static org.mockito.Mockito.when;
3132

33+
import com.google.api.core.ApiFuture;
3234
import com.google.api.core.ApiFutures;
35+
import com.google.api.core.SettableApiFuture;
3336
import com.google.cloud.grpc.GrpcTransportOptions;
3437
import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
3538
import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler;
@@ -195,6 +198,129 @@ public void testCommitWithClientContext() {
195198
assertEquals(clientContext, capturedOptions.clientContext());
196199
}
197200

201+
@Test
202+
public void commitCancelsInFlightRpcWhenCallingThreadInterrupted() {
203+
when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
204+
TransactionContextImpl transaction =
205+
TransactionContextImpl.newBuilder()
206+
.setSession(session)
207+
.setTransactionId(ByteString.copyFromUtf8("test-txn"))
208+
.setOptions(Options.fromTransactionOptions())
209+
.setRpc(rpc)
210+
.setTracer(tracer)
211+
.setSpan(span)
212+
.build();
213+
SettableApiFuture<CommitResponse> inFlightCommit = SettableApiFuture.create();
214+
when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap()))
215+
.thenAnswer(
216+
invocation -> {
217+
Thread.currentThread().interrupt();
218+
return inFlightCommit;
219+
});
220+
221+
try {
222+
SpannerException e = assertThrows(SpannerException.class, transaction::commit);
223+
assertEquals(ErrorCode.CANCELLED, e.getErrorCode());
224+
assertTrue("in-flight Commit RPC was not cancelled", inFlightCommit.isCancelled());
225+
} finally {
226+
// Clear the interrupt flag so it cannot leak into other tests.
227+
Thread.interrupted();
228+
}
229+
}
230+
231+
@Test
232+
public void commitAsyncCancelsInFlightRpcWhenReturnedFutureIsCancelled() {
233+
when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
234+
TransactionContextImpl transaction =
235+
TransactionContextImpl.newBuilder()
236+
.setSession(session)
237+
.setTransactionId(ByteString.copyFromUtf8("test-txn"))
238+
.setOptions(Options.fromTransactionOptions())
239+
.setRpc(rpc)
240+
.setTracer(tracer)
241+
.setSpan(span)
242+
.build();
243+
SettableApiFuture<CommitResponse> inFlightCommit = SettableApiFuture.create();
244+
when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap()))
245+
.thenReturn(inFlightCommit);
246+
247+
ApiFuture<com.google.cloud.spanner.CommitResponse> commitFuture = transaction.commitAsync();
248+
assertTrue(commitFuture.cancel(true));
249+
250+
assertTrue("in-flight Commit RPC was not cancelled", inFlightCommit.isCancelled());
251+
}
252+
253+
@Test
254+
public void commitAsyncSkipsCommitRpcWhenCancelledBeforeItIsSent() {
255+
SettableApiFuture<Transaction> beginTransaction = SettableApiFuture.create();
256+
setUpTransactionThatCommitsAfterBeginTransaction(beginTransaction);
257+
TransactionContextImpl transaction = newTransactionWithoutTransactionId();
258+
259+
ApiFuture<com.google.cloud.spanner.CommitResponse> commitFuture = transaction.commitAsync();
260+
// The Commit RPC is only sent once BeginTransaction has finished, so cancelling here means
261+
// that the commit is abandoned before there is any RPC to cancel.
262+
assertTrue(commitFuture.cancel(true));
263+
verify(rpc, never()).commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap());
264+
265+
beginTransaction.set(
266+
Transaction.newBuilder().setId(ByteString.copyFromUtf8("test-txn")).build());
267+
268+
// The commit was already abandoned, so the Commit RPC should not be sent at all.
269+
verify(rpc, never()).commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap());
270+
}
271+
272+
@Test
273+
public void commitSkipsCommitRpcWhenCallingThreadInterruptedBeforeItIsSent() {
274+
SettableApiFuture<Transaction> beginTransaction = SettableApiFuture.create();
275+
setUpTransactionThatCommitsAfterBeginTransaction(beginTransaction);
276+
TransactionContextImpl transaction = newTransactionWithoutTransactionId();
277+
278+
try {
279+
// Interrupting before commit() waits for the result means that it gives up while
280+
// BeginTransaction is still pending, and therefore before the Commit RPC has been sent.
281+
Thread.currentThread().interrupt();
282+
SpannerException e = assertThrows(SpannerException.class, transaction::commit);
283+
assertEquals(ErrorCode.CANCELLED, e.getErrorCode());
284+
verify(rpc, never()).commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap());
285+
286+
// commit() restores the interrupt flag, and BeginTransaction finishes on a gax thread that
287+
// does not have that flag. Clear it so the listeners below are not interrupted either.
288+
Thread.interrupted();
289+
beginTransaction.set(
290+
Transaction.newBuilder().setId(ByteString.copyFromUtf8("test-txn")).build());
291+
292+
// The commit was already abandoned, so the Commit RPC should not be sent at all.
293+
verify(rpc, never()).commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap());
294+
} finally {
295+
// Clear the interrupt flag so it cannot leak into other tests.
296+
Thread.interrupted();
297+
}
298+
}
299+
300+
private void setUpTransactionThatCommitsAfterBeginTransaction(
301+
SettableApiFuture<Transaction> beginTransaction) {
302+
when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
303+
when(session.beginTransactionAsync(
304+
Mockito.any(Options.class),
305+
Mockito.anyBoolean(),
306+
Mockito.anyMap(),
307+
Mockito.any(),
308+
Mockito.any()))
309+
.thenReturn(beginTransaction);
310+
when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap()))
311+
.thenReturn(SettableApiFuture.create());
312+
}
313+
314+
private TransactionContextImpl newTransactionWithoutTransactionId() {
315+
return TransactionContextImpl.newBuilder()
316+
.setSession(session)
317+
.setOptions(Options.fromTransactionOptions())
318+
.setRpc(rpc)
319+
.setTracer(tracer)
320+
.setSpan(span)
321+
.build();
322+
}
323+
198324
@SuppressWarnings("unchecked")
199325
@Test
200326
public void usesPreparedTransaction() {

0 commit comments

Comments
 (0)