Skip to content

Commit 8c4b537

Browse files
neozwuneozwu
committed
retry datastore transactions (#1242)
1 parent a27d8c7 commit 8c4b537

5 files changed

Lines changed: 155 additions & 5 deletions

File tree

‎google-cloud-core/src/main/java/com/google/cloud/BaseService.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ public RetryResult afterEval(Exception exception, RetryResult retryResult) {
3939
public RetryResult beforeEval(Exception exception) {
4040
if (exception instanceof BaseServiceException) {
4141
boolean retriable = ((BaseServiceException) exception).isRetryable();
42-
return retriable ? Interceptor.RetryResult.RETRY : Interceptor.RetryResult.NO_RETRY;
42+
return retriable ? Interceptor.RetryResult.RETRY : RetryResult.CONTINUE_EVALUATION;
4343
}
4444
return Interceptor.RetryResult.CONTINUE_EVALUATION;
4545
}

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,11 @@
1616

1717
package com.google.cloud.datastore;
1818

19+
import com.google.api.gax.core.RetrySettings;
1920
import com.google.cloud.BaseService;
21+
import com.google.cloud.ExceptionHandler;
2022
import com.google.cloud.RetryHelper;
2123
import com.google.cloud.RetryHelper.RetryHelperException;
22-
import com.google.api.gax.core.RetrySettings;
2324
import com.google.cloud.ServiceOptions;
2425
import com.google.cloud.datastore.ReadOption.EventualConsistency;
2526
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
@@ -31,7 +32,6 @@
3132
import com.google.common.collect.Sets;
3233
import com.google.datastore.v1.ReadOptions.ReadConsistency;
3334
import com.google.protobuf.ByteString;
34-
3535
import java.util.ArrayList;
3636
import java.util.Arrays;
3737
import java.util.Collections;
@@ -66,8 +66,23 @@ public Transaction newTransaction() {
6666
}
6767

6868
@Override
69-
public <T> T runInTransaction(TransactionCallable<T> callable) {
70-
return DatastoreHelper.runInTransaction(this, callable);
69+
public <T> T runInTransaction(final TransactionCallable<T> callable) {
70+
ExceptionHandler TRANSACTION_EXCEPTION_HANDLER = TransactionExceptionHandler.build();
71+
final DatastoreImpl self = this;
72+
try {
73+
return RetryHelper.runWithRetries(
74+
new Callable<T>() {
75+
@Override
76+
public T call() throws DatastoreException {
77+
return DatastoreHelper.runInTransaction(self, callable);
78+
}
79+
},
80+
retrySettings,
81+
TRANSACTION_EXCEPTION_HANDLER,
82+
getOptions().getClock());
83+
} catch (RetryHelperException e) {
84+
throw DatastoreException.translateAndThrow(e);
85+
}
7186
}
7287

7388
@Override

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionExceptionHandler.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*
2+
* Copyright 2017 Google Inc. All Rights Reserved.
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.datastore;
18+
19+
import com.google.cloud.ExceptionHandler;
20+
import com.google.cloud.ExceptionHandler.Interceptor;
21+
22+
public class TransactionExceptionHandler {
23+
public static final Interceptor TRANSACTION_EXCEPTION_HANDLER_INTERCEPTOR =
24+
new Interceptor() {
25+
@Override
26+
public RetryResult beforeEval(Exception exception) {
27+
if (exception instanceof DatastoreException) {
28+
DatastoreException e = getInnerException((DatastoreException) exception);
29+
if (e.getCode() == 10 || e.getReason() == "ABORTED") {
30+
return RetryResult.RETRY;
31+
}
32+
}
33+
return Interceptor.RetryResult.CONTINUE_EVALUATION;
34+
}
35+
36+
@Override
37+
public RetryResult afterEval(Exception exception, RetryResult retryResult) {
38+
return Interceptor.RetryResult.CONTINUE_EVALUATION;
39+
}
40+
41+
private DatastoreException getInnerException(DatastoreException exception) {
42+
while (exception.getCause() instanceof DatastoreException) {
43+
exception = (DatastoreException) exception.getCause();
44+
}
45+
return exception;
46+
}
47+
};
48+
49+
public static ExceptionHandler build() {
50+
return ExceptionHandler.newBuilder()
51+
.abortOn(RuntimeException.class)
52+
.addInterceptors(
53+
DatastoreImpl.EXCEPTION_HANDLER_INTERCEPTOR, TRANSACTION_EXCEPTION_HANDLER_INTERCEPTOR)
54+
.build();
55+
}
56+
}

‎google-cloud-datastore/src/test/java/com/google/cloud/datastore/TransactionExceptionHandlerTest.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright 2017 Google Inc. All Rights Reserved.
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.datastore;
18+
19+
import static junit.framework.TestCase.assertFalse;
20+
import static junit.framework.TestCase.assertTrue;
21+
22+
import com.google.cloud.BaseServiceException;
23+
import com.google.cloud.ExceptionHandler;
24+
import org.junit.Rule;
25+
import org.junit.Test;
26+
import org.junit.rules.ExpectedException;
27+
28+
/** Tests for {@link TransactionExceptionHandler}. */
29+
public class TransactionExceptionHandlerTest {
30+
31+
@Rule public ExpectedException thrown = ExpectedException.none();
32+
33+
@Test
34+
public void testShouldTry() {
35+
ExceptionHandler handler =
36+
ExceptionHandler.newBuilder()
37+
.abortOn(RuntimeException.class)
38+
.addInterceptors(DatastoreImpl.EXCEPTION_HANDLER_INTERCEPTOR)
39+
.build();
40+
ExceptionHandler transactionHandler = TransactionExceptionHandler.build();
41+
42+
assertFalse(handler.accept(new DatastoreException(10, "", "ABORTED", false, null)));
43+
assertFalse(handler.accept(new DatastoreException(10, "", "", false, null)));
44+
assertFalse(handler.accept(new DatastoreException(0, "", "", false, null)));
45+
46+
assertTrue(transactionHandler.accept(new DatastoreException(10, "", "ABORTED", false, null)));
47+
assertTrue(transactionHandler.accept(new DatastoreException(10, "", "", false, null)));
48+
assertFalse(transactionHandler.accept(new DatastoreException(0, "", "", false, null)));
49+
50+
DatastoreException nestedDatastoreException =
51+
new DatastoreException(
52+
BaseServiceException.UNKNOWN_CODE,
53+
"",
54+
null,
55+
new DatastoreException(10, "", "ABORTED", false, null));
56+
57+
assertTrue(transactionHandler.accept(nestedDatastoreException));
58+
assertFalse(handler.accept(nestedDatastoreException));
59+
}
60+
}

‎google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import com.google.cloud.datastore.Datastore;
3131
import com.google.cloud.datastore.DatastoreException;
3232
import com.google.cloud.datastore.DatastoreOptions;
33+
import com.google.cloud.datastore.DatastoreReaderWriter;
3334
import com.google.cloud.datastore.DateTime;
3435
import com.google.cloud.datastore.DateTimeValue;
3536
import com.google.cloud.datastore.Entity;
@@ -729,4 +730,22 @@ public void testDelete() {
729730
assertNull(keys.next());
730731
assertFalse(keys.hasNext());
731732
}
733+
734+
@Test
735+
public void testRunInTransaction() {
736+
Datastore.TransactionCallable<Integer> callable = new Datastore.TransactionCallable<Integer>() {
737+
private Integer attempts = 0;
738+
public Integer run(DatastoreReaderWriter transaction) {
739+
transaction.get(KEY1);
740+
if (attempts < 1) {
741+
++attempts;
742+
throw new DatastoreException(10,"","ABORTED",false,null);
743+
}
744+
745+
return attempts;
746+
}
747+
};
748+
int result = DATASTORE.runInTransaction(callable);
749+
assertTrue(result == 1);
750+
}
732751
}

0 commit comments

Comments
 (0)