fix: update retry logic for operations that can be associated with a transaction

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

@@ -51,6 +51,8 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> 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);

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() {}
+}

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();