Skip to content
Permalink
Browse files
fix: fix aborted handling of batchUpdateAsync (#421)
  • Loading branch information
olavloite committed Sep 11, 2020
1 parent b35304e commit 61540085c971d7885e4938b486e051a1ed9cf35f
@@ -540,7 +540,7 @@ public ApiFuture<long[]> batchUpdateAsync(Iterable<Statement> statements) {
decreaseAsyncOperations();
throw t;
}
final ApiFuture<long[]> updateCounts =
ApiFuture<long[]> updateCounts =
ApiFutures.transform(
response,
new ApiFunction<ExecuteBatchDmlResponse, long[]>() {
@@ -565,19 +565,24 @@ public long[] apply(ExecuteBatchDmlResponse input) {
}
},
MoreExecutors.directExecutor());
updateCounts =
ApiFutures.catching(
updateCounts,
Throwable.class,
new ApiFunction<Throwable, long[]>() {
@Override
public long[] apply(Throwable input) {
SpannerException e = SpannerExceptionFactory.newSpannerException(input);
onError(e);
throw e;
}
},
MoreExecutors.directExecutor());
updateCounts.addListener(
new Runnable() {
@Override
public void run() {
try {
updateCounts.get();
} catch (ExecutionException e) {
onError(SpannerExceptionFactory.newSpannerException(e.getCause()));
} catch (InterruptedException e) {
onError(SpannerExceptionFactory.propagateInterrupt(e));
} finally {
decreaseAsyncOperations();
}
decreaseAsyncOperations();
}
},
MoreExecutors.directExecutor());
@@ -734,6 +734,47 @@ public ApiFuture<long[]> apply(TransactionContext txn, Void input)
CommitRequest.class);
}

@Test
public void asyncTransactionManagerBatchUpdateAbortedBeforeFirstStatement() throws Exception {
final AtomicInteger attempt = new AtomicInteger();
try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) {
TransactionContextFuture txn = mgr.beginAsync();
while (true) {
try {
txn.then(
new AsyncTransactionFunction<Void, long[]>() {
@Override
public ApiFuture<long[]> apply(TransactionContext txn, Void input)
throws Exception {
if (attempt.incrementAndGet() == 1) {
mockSpanner.abortTransaction(txn);
}
return txn.batchUpdateAsync(
ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT));
}
},
executor)
.commitAsync()
.get();
break;
} catch (AbortedException e) {
txn = mgr.resetForRetryAsync();
}
}
}
assertThat(attempt.get()).isEqualTo(2);
// There should only be 1 CommitRequest, as the first attempt should abort already after the
// ExecuteBatchDmlRequest.
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
BatchCreateSessionsRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
}

@Test
public void asyncTransactionManagerWithBatchUpdateCommitAborted() throws Exception {
try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) {

0 comments on commit 6154008

Please sign in to comment.