diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index be31594e7..25e434227 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -51,6 +51,8 @@ final class DatastoreImpl extends BaseService implements Datas private final RetrySettings retrySettings; private static final ExceptionHandler TRANSACTION_EXCEPTION_HANDLER = TransactionExceptionHandler.build(); + private static final ExceptionHandler TRANSACTION_OPERATION_EXCEPTION_HANDLER = + TransactionOperationExceptionHandler.build(); DatastoreImpl(DatastoreOptions options) { super(options); @@ -182,7 +184,9 @@ public com.google.datastore.v1.RunQueryResponse call() throws DatastoreException } }, retrySettings, - EXCEPTION_HANDLER, + requestPb.getReadOptions().getTransaction().isEmpty() + ? EXCEPTION_HANDLER + : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); @@ -394,7 +398,9 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException { } }, retrySettings, - EXCEPTION_HANDLER, + requestPb.getReadOptions().getTransaction().isEmpty() + ? EXCEPTION_HANDLER + : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); @@ -532,7 +538,9 @@ public com.google.datastore.v1.CommitResponse call() throws DatastoreException { } }, retrySettings, - EXCEPTION_HANDLER, + requestPb.getTransaction().isEmpty() + ? EXCEPTION_HANDLER + : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { throw DatastoreException.translateAndThrow(e); diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionOperationExceptionHandler.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionOperationExceptionHandler.java new file mode 100644 index 000000000..0556555d0 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionOperationExceptionHandler.java @@ -0,0 +1,68 @@ +/* + * Copyright 2017 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.datastore; + +import com.google.api.core.BetaApi; +import com.google.cloud.BaseService; +import com.google.cloud.ExceptionHandler; +import com.google.cloud.ExceptionHandler.Interceptor; + +@BetaApi +public class TransactionOperationExceptionHandler { + + public static final Interceptor TRANSACTION_OPERATION_EXCEPTION_HANDLER_INTERCEPTOR = + new Interceptor() { + + private static final long serialVersionUID = -1240723093072535978L; + + private static final int ABORTED_CODE = 10; + + @Override + public RetryResult beforeEval(Exception exception) { + if (exception instanceof DatastoreException) { + DatastoreException e = getInnerException((DatastoreException) exception); + if (e.getCode() == ABORTED_CODE + || e.getReason() != null && e.getReason().equals("ABORTED")) { + return RetryResult.NO_RETRY; + } + } + return BaseService.EXCEPTION_HANDLER_INTERCEPTOR.beforeEval(exception); + } + + @Override + public RetryResult afterEval(Exception exception, RetryResult retryResult) { + return BaseService.EXCEPTION_HANDLER_INTERCEPTOR.afterEval(exception, retryResult); + } + + private DatastoreException getInnerException(DatastoreException exception) { + while (exception.getCause() instanceof DatastoreException) { + exception = (DatastoreException) exception.getCause(); + } + return exception; + } + }; + + public static ExceptionHandler build() { + return ExceptionHandler.newBuilder() + .abortOn(RuntimeException.class) + .addInterceptors(TRANSACTION_OPERATION_EXCEPTION_HANDLER_INTERCEPTOR) + .build(); + } + + /** Intentionally private empty constructor to disable instantiation of this class. */ + private TransactionOperationExceptionHandler() {} +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java index 8049e0957..442a6eb72 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java @@ -1135,6 +1135,56 @@ public void testRetryableException() { EasyMock.verify(rpcFactoryMock, rpcMock); } + @Test + public void testRetryableExceptionForOperationWithTxn() { + ByteString txnBytes = ByteString.copyFromUtf8("txn1"); + LookupRequest requestPb = + LookupRequest.newBuilder() + .addKeys(KEY1.toPb()) + .setReadOptions(ReadOptions.newBuilder().setTransaction(txnBytes).build()) + .build(); + LookupResponse responsePb = + LookupResponse.newBuilder() + .addFound(EntityResult.newBuilder().setEntity(ENTITY1.toPb())) + .build(); + EasyMock.expect(rpcMock.beginTransaction(EasyMock.anyObject(BeginTransactionRequest.class))) + .andReturn(BeginTransactionResponse.newBuilder().setTransaction(txnBytes).build()); + EasyMock.expect(rpcMock.lookup(requestPb)) + .andThrow(new DatastoreException(14, "UNAVAILABLE", "UNAVAILABLE", null)) + .andReturn(responsePb); + EasyMock.replay(rpcFactoryMock, rpcMock); + Datastore datastore = rpcMockOptions.getService(); + Transaction transaction = datastore.newTransaction(); + Entity entity = transaction.get(KEY1); + assertEquals(ENTITY1, entity); + EasyMock.verify(rpcFactoryMock, rpcMock); + } + + @Test + public void testNonRetryableExceptionForOperationWithTxn() { + ByteString txnBytes = ByteString.copyFromUtf8("txn1"); + LookupRequest requestPb = + LookupRequest.newBuilder() + .addKeys(KEY1.toPb()) + .setReadOptions(ReadOptions.newBuilder().setTransaction(txnBytes).build()) + .build(); + EasyMock.expect(rpcMock.beginTransaction(EasyMock.anyObject(BeginTransactionRequest.class))) + .andReturn(BeginTransactionResponse.newBuilder().setTransaction(txnBytes).build()); + EasyMock.expect(rpcMock.lookup(requestPb)) + .andThrow(new DatastoreException(10, "ABORTED", "ABORTED", null)) + .times(1); + EasyMock.replay(rpcFactoryMock, rpcMock); + try { + Datastore datastore = rpcMockOptions.getService(); + Transaction transaction = datastore.newTransaction(); + transaction.get(KEY1); + Assert.fail(); + EasyMock.verify(rpcFactoryMock, rpcMock); + } catch (DatastoreException ex) { + assertEquals("ABORTED", ex.getMessage()); + } + } + @Test public void testNonRetryableException() { LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build();