Skip to content

Commit 20bc33f

Browse files
olavloitesduskis
olavloite
authored and
sduskis
committed
Spanner: Aborted batch dml should be retried (#4955)
* batchDml can return aborted for one statement When executing batchDml statements, it is possible that the transaction is aborted by Spanner for one of the statements. This used to throw a SpannerBatchUpdateException, which would cause the transaction to fail, while the appropriate exception is AbortedException to trigger a retry of the transaction. * added missing copyright header * added test for batch dml exception in TransactionRunnerImplTest * moved asserts to the test methods * fix flaky test case * ran code formatter
1 parent cd4d635 commit 20bc33f

4 files changed

Lines changed: 162 additions & 2 deletions

File tree

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.google.common.annotations.VisibleForTesting;
2929
import com.google.common.collect.ImmutableMap;
3030
import com.google.protobuf.ByteString;
31+
import com.google.rpc.Code;
3132
import com.google.spanner.v1.CommitRequest;
3233
import com.google.spanner.v1.CommitResponse;
3334
import com.google.spanner.v1.ExecuteBatchDmlRequest;
@@ -270,7 +271,12 @@ public com.google.spanner.v1.ExecuteBatchDmlResponse call() throws Exception {
270271
results[i] = response.getResultSets(i).getStats().getRowCountExact();
271272
}
272273

273-
if (response.getStatus().getCode() != 0) {
274+
// If one of the DML statements was aborted, we should throw an aborted exception.
275+
// In all other cases, we should throw a BatchUpdateException.
276+
if (response.getStatus().getCode() == Code.ABORTED_VALUE) {
277+
throw newSpannerException(
278+
ErrorCode.fromRpcStatus(response.getStatus()), response.getStatus().getMessage());
279+
} else if (response.getStatus().getCode() != 0) {
274280
throw newSpannerBatchUpdateException(
275281
ErrorCode.fromRpcStatus(response.getStatus()),
276282
response.getStatus().getMessage(),

‎google-cloud-clients/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright 2019 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner;
18+
19+
import static org.mockito.Mockito.mock;
20+
import static org.mockito.Mockito.when;
21+
22+
import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl;
23+
import com.google.cloud.spanner.spi.v1.SpannerRpc;
24+
import com.google.protobuf.ByteString;
25+
import com.google.rpc.Code;
26+
import com.google.rpc.Status;
27+
import com.google.spanner.v1.ExecuteBatchDmlRequest;
28+
import com.google.spanner.v1.ExecuteBatchDmlResponse;
29+
import java.util.Arrays;
30+
import org.junit.Test;
31+
import org.junit.runner.RunWith;
32+
import org.junit.runners.JUnit4;
33+
import org.mockito.Mockito;
34+
35+
@RunWith(JUnit4.class)
36+
public class TransactionContextImplTest {
37+
38+
@Test(expected = AbortedException.class)
39+
public void batchDmlAborted() {
40+
batchDml(Code.ABORTED_VALUE);
41+
}
42+
43+
@Test(expected = SpannerBatchUpdateException.class)
44+
public void batchDmlException() {
45+
batchDml(Code.FAILED_PRECONDITION_VALUE);
46+
}
47+
48+
@SuppressWarnings("unchecked")
49+
private void batchDml(int status) {
50+
SessionImpl session = mock(SessionImpl.class);
51+
when(session.getName()).thenReturn("test");
52+
SpannerRpc rpc = mock(SpannerRpc.class);
53+
ExecuteBatchDmlResponse response =
54+
ExecuteBatchDmlResponse.newBuilder()
55+
.setStatus(Status.newBuilder().setCode(status).build())
56+
.build();
57+
Statement statement = Statement.of("UPDATE FOO SET BAR=1");
58+
59+
when(rpc.executeBatchDml(Mockito.any(ExecuteBatchDmlRequest.class), Mockito.anyMap()))
60+
.thenReturn(response);
61+
try (TransactionContextImpl impl =
62+
new TransactionContextImpl(session, ByteString.copyFromUtf8("test"), rpc, 10)) {
63+
impl.batchUpdate(Arrays.asList(statement));
64+
}
65+
}
66+
}

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

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,16 +27,29 @@
2727

2828
import com.google.api.client.util.BackOff;
2929
import com.google.cloud.spanner.TransactionRunner.TransactionCallable;
30+
import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl;
3031
import com.google.cloud.spanner.spi.v1.SpannerRpc;
32+
import com.google.common.base.Preconditions;
33+
import com.google.protobuf.ByteString;
34+
import com.google.protobuf.Timestamp;
35+
import com.google.rpc.Code;
36+
import com.google.spanner.v1.CommitRequest;
37+
import com.google.spanner.v1.CommitResponse;
38+
import com.google.spanner.v1.ExecuteBatchDmlRequest;
39+
import com.google.spanner.v1.ExecuteBatchDmlResponse;
40+
import com.google.spanner.v1.ResultSet;
41+
import com.google.spanner.v1.ResultSetStats;
3142
import io.grpc.Context;
3243
import io.grpc.Status;
3344
import io.grpc.StatusRuntimeException;
45+
import java.util.Arrays;
3446
import java.util.concurrent.atomic.AtomicInteger;
3547
import org.junit.Before;
3648
import org.junit.Test;
3749
import org.junit.runner.RunWith;
3850
import org.junit.runners.JUnit4;
3951
import org.mockito.Mock;
52+
import org.mockito.Mockito;
4053
import org.mockito.MockitoAnnotations;
4154

4255
/** Unit test for {@link com.google.cloud.spanner.SpannerImpl.TransactionRunnerImpl} */
@@ -141,6 +154,77 @@ public void runResourceExhaustedNoRetry() throws Exception {
141154
verify(txn).rollback();
142155
}
143156

157+
@Test
158+
public void batchDmlAborted() {
159+
long updateCount[] = batchDmlException(Code.ABORTED_VALUE);
160+
assertThat(updateCount.length).isEqualTo(2);
161+
assertThat(updateCount[0]).isEqualTo(1L);
162+
assertThat(updateCount[1]).isEqualTo(1L);
163+
}
164+
165+
@Test
166+
public void batchDmlFailedPrecondition() {
167+
try {
168+
batchDmlException(Code.FAILED_PRECONDITION_VALUE);
169+
fail("Expected exception");
170+
} catch (SpannerBatchUpdateException e) {
171+
assertThat(e.getUpdateCounts().length).isEqualTo(1);
172+
assertThat(e.getUpdateCounts()[0]).isEqualTo(1L);
173+
assertThat(e.getCode() == Code.FAILED_PRECONDITION_VALUE);
174+
}
175+
}
176+
177+
@SuppressWarnings("unchecked")
178+
private long[] batchDmlException(int status) {
179+
Preconditions.checkArgument(status != Code.OK_VALUE);
180+
TransactionContextImpl transaction =
181+
new TransactionContextImpl(session, ByteString.copyFromUtf8("test"), rpc, 10);
182+
when(session.newTransaction()).thenReturn(transaction);
183+
when(session.getName()).thenReturn("test");
184+
TransactionRunnerImpl runner = new TransactionRunnerImpl(session, rpc, 10);
185+
ExecuteBatchDmlResponse response1 =
186+
ExecuteBatchDmlResponse.newBuilder()
187+
.addResultSets(
188+
ResultSet.newBuilder()
189+
.setStats(ResultSetStats.newBuilder().setRowCountExact(1L))
190+
.build())
191+
.setStatus(com.google.rpc.Status.newBuilder().setCode(status).build())
192+
.build();
193+
ExecuteBatchDmlResponse response2 =
194+
ExecuteBatchDmlResponse.newBuilder()
195+
.addResultSets(
196+
ResultSet.newBuilder()
197+
.setStats(ResultSetStats.newBuilder().setRowCountExact(1L))
198+
.build())
199+
.addResultSets(
200+
ResultSet.newBuilder()
201+
.setStats(ResultSetStats.newBuilder().setRowCountExact(1L))
202+
.build())
203+
.setStatus(com.google.rpc.Status.newBuilder().setCode(Code.OK_VALUE).build())
204+
.build();
205+
when(rpc.executeBatchDml(Mockito.any(ExecuteBatchDmlRequest.class), Mockito.anyMap()))
206+
.thenReturn(response1, response2);
207+
CommitResponse commitResponse =
208+
CommitResponse.newBuilder().setCommitTimestamp(Timestamp.getDefaultInstance()).build();
209+
when(rpc.commit(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenReturn(commitResponse);
210+
final Statement statement = Statement.of("UPDATE FOO SET BAR=1");
211+
final AtomicInteger numCalls = new AtomicInteger(0);
212+
long updateCount[] =
213+
runner.run(
214+
new TransactionCallable<long[]>() {
215+
@Override
216+
public long[] run(TransactionContext transaction) throws Exception {
217+
numCalls.incrementAndGet();
218+
return transaction.batchUpdate(Arrays.asList(statement, statement));
219+
}
220+
});
221+
if (status == Code.ABORTED_VALUE) {
222+
// Assert that the method ran twice because the first response aborted.
223+
assertThat(numCalls.get()).isEqualTo(2);
224+
}
225+
return updateCount;
226+
}
227+
144228
private void runTransaction(final Exception exception) {
145229
transactionRunner.run(
146230
new TransactionCallable<Void>() {

‎google-cloud-clients/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionManagerTest.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,11 @@ public void abortAndRetry() throws InterruptedException {
170170
break;
171171
} catch (AbortedException e) {
172172
Thread.sleep(e.getRetryDelayInMillis() / 1000);
173-
txn1 = manager1.resetForRetry();
173+
// It is possible that it was txn2 that aborted.
174+
// In that case we should just retry without resetting anything.
175+
if (manager1.getState() == TransactionState.ABORTED) {
176+
txn1 = manager1.resetForRetry();
177+
}
174178
}
175179
}
176180

0 commit comments

Comments
 (0)