feat(datastore): Support for retrying Datastore transactions with the previous transaction id.

docs/src/main/asciidoc/datastore.adoc

@@ -82,7 +82,7 @@ The following configuration options are available:
 | `spring.cloud.gcp.datastore.emulator.firestore-in-datastore-mode` | Configures whether the emulator runs in "Cloud Firestore in Datastore Mode" | No | `false`
 | `spring.cloud.gcp.datastore.skip-null-value` | Whether skip inserting `null` values. | No | The default value is `false`.
 If configured to `true`, `null` value will not be inserted into the datastore.
-
+| `spring.cloud.gcp.datastore.previous-transaction-retry-enabled` | Whether to retry a transactions with previous transaction when combined with Retriable. | No | `true`
 |===
 
 ==== Repository settings

spring-cloud-gcp-autoconfigure/src/main/java/com/google/cloud/spring/autoconfigure/datastore/DatastoreTransactionManagerAutoConfiguration.java

@@ -25,6 +25,7 @@
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.autoconfigure.transaction.TransactionAutoConfiguration;
 import org.springframework.boot.autoconfigure.transaction.TransactionManagerCustomizers;
+import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 
 /**
@@ -44,24 +45,29 @@ private DatastoreTransactionManagerAutoConfiguration() {
 
   /** Configuration class. */
   @AutoConfiguration
+  @EnableConfigurationProperties(GcpDatastoreProperties.class)
   static class DatastoreTransactionManagerConfiguration {
 
     private final DatastoreProvider datastore;
 
     private final TransactionManagerCustomizers transactionManagerCustomizers;
 
+    private final boolean previousTransactionRetryEnabled;
+
     DatastoreTransactionManagerConfiguration(
+        GcpDatastoreProperties gcpDatastoreProperties,
         DatastoreProvider datastore,
         ObjectProvider<TransactionManagerCustomizers> transactionManagerCustomizers) {
       this.datastore = datastore;
       this.transactionManagerCustomizers = transactionManagerCustomizers.getIfAvailable();
+      this.previousTransactionRetryEnabled = gcpDatastoreProperties.isPreviousTransactionRetryEnabled();
     }
 
     @Bean
     @ConditionalOnMissingBean
     public DatastoreTransactionManager datastoreTransactionManager() {
       DatastoreTransactionManager transactionManager =
-          new DatastoreTransactionManager(this.datastore);
+          new DatastoreTransactionManager(this.datastore, this.previousTransactionRetryEnabled);
       if (this.transactionManagerCustomizers != null) {
         this.transactionManagerCustomizers.customize(transactionManager);
       }

spring-cloud-gcp-autoconfigure/src/main/java/com/google/cloud/spring/autoconfigure/datastore/GcpDatastoreProperties.java

@@ -49,6 +49,9 @@ public class GcpDatastoreProperties implements CredentialsSupplier {
   /** Whether skip the insertion if the value is null */
   private boolean skipNullValue;
 
+  /** Whether to retry a transactions with previous transaction when combined with Retriable. */
+  private boolean previousTransactionRetryEnabled = true;
+
   @Override
   public Credentials getCredentials() {
     return this.credentials;
@@ -97,4 +100,12 @@ public boolean isSkipNullValue() {
   public void setSkipNullValue(boolean skipNullValue) {
     this.skipNullValue = skipNullValue;
   }
+
+  public boolean isPreviousTransactionRetryEnabled() {
+    return previousTransactionRetryEnabled;
+  }
+
+  public void setPreviousTransactionRetryEnabled(boolean previousTransactionRetryEnabled) {
+    this.previousTransactionRetryEnabled = previousTransactionRetryEnabled;
+  }
 }

spring-cloud-gcp-data-datastore/pom.xml

@@ -26,6 +26,10 @@
       <artifactId>spring-data-commons</artifactId>
     </dependency>
     <dependency>
+      <groupId>org.springframework</groupId>
+      <artifactId>spring-aspects</artifactId>
+    </dependency>
+      <dependency>
       <groupId>org.springframework</groupId>
       <artifactId>spring-context</artifactId>
     </dependency>
@@ -41,5 +45,9 @@
       <groupId>org.springframework</groupId>
       <artifactId>spring-tx</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.springframework.retry</groupId>
+      <artifactId>spring-retry</artifactId>
+    </dependency>
   </dependencies>
 </project>

spring-cloud-gcp-data-datastore/src/main/java/com/google/cloud/spring/data/datastore/core/DatastoreTransactionManager.java

@@ -20,7 +20,11 @@
 import com.google.cloud.datastore.DatastoreException;
 import com.google.cloud.datastore.Transaction;
 import com.google.datastore.v1.TransactionOptions;
+import com.google.datastore.v1.TransactionOptions.ReadWrite;
+import com.google.protobuf.ByteString;
 import java.util.function.Supplier;
+import org.springframework.retry.RetryContext;
+import org.springframework.retry.support.RetrySynchronizationManager;
 import org.springframework.transaction.TransactionDefinition;
 import org.springframework.transaction.TransactionException;
 import org.springframework.transaction.TransactionSystemException;
@@ -34,15 +38,25 @@
  * @since 1.1
  */
 public class DatastoreTransactionManager extends AbstractPlatformTransactionManager {
+
+  static final String PREVIOUS_TRANSACTION_ID_ATTRIBUTE = "datastore-transaction-id";
+
   private static final TransactionOptions READ_ONLY_OPTIONS =
       TransactionOptions.newBuilder()
           .setReadOnly(TransactionOptions.ReadOnly.newBuilder().build())
           .build();
 
   private final Supplier<Datastore> datastore;
+
+  private final boolean previousTransactionRetryEnabled;
 
   public DatastoreTransactionManager(final Supplier<Datastore> datastore) {
+    this(datastore, true);
+  }
+
+  public DatastoreTransactionManager(final Supplier<Datastore> datastore, boolean previousTransactionRetryEnabled) {
     this.datastore = datastore;
+    this.previousTransactionRetryEnabled = previousTransactionRetryEnabled;
   }
 
   @Override
@@ -71,9 +85,16 @@ protected void doBegin(Object transactionObject, TransactionDefinition transacti
           "DatastoreTransactionManager supports only propagation behavior "
               + "TransactionDefinition.PROPAGATION_REQUIRED");
     }
+    RetryContext retryContext = RetrySynchronizationManager.getContext();
     Tx tx = (Tx) transactionObject;
     if (transactionDefinition.isReadOnly()) {
       tx.transaction = tx.datastore.newTransaction(READ_ONLY_OPTIONS);
+    } else if (this.previousTransactionRetryEnabled && retryContext != null && retryContext.hasAttribute(PREVIOUS_TRANSACTION_ID_ATTRIBUTE)) {
+      tx.transaction = tx.datastore.newTransaction(
+          TransactionOptions.newBuilder()
+              .setReadWrite(ReadWrite.newBuilder().setPreviousTransaction(
+                  (ByteString) retryContext.getAttribute(PREVIOUS_TRANSACTION_ID_ATTRIBUTE)))
+              .build());
     } else {
       tx.transaction = tx.datastore.newTransaction();
     }
@@ -85,13 +106,20 @@ protected void doBegin(Object transactionObject, TransactionDefinition transacti
   protected void doCommit(DefaultTransactionStatus defaultTransactionStatus)
       throws TransactionException {
     Tx tx = (Tx) defaultTransactionStatus.getTransaction();
+    RetryContext retryContext = RetrySynchronizationManager.getContext();
     try {
       if (tx.transaction.isActive()) {
         tx.transaction.commit();
+        if (retryContext != null) {
+          retryContext.removeAttribute(PREVIOUS_TRANSACTION_ID_ATTRIBUTE);
+        }
       } else {
         this.logger.debug("Transaction was not committed because it is no longer active.");
       }
     } catch (DatastoreException ex) {
+      if (this.previousTransactionRetryEnabled && retryContext != null) {
+        retryContext.setAttribute(PREVIOUS_TRANSACTION_ID_ATTRIBUTE, tx.transaction.getTransactionId());
+      }
       throw new TransactionSystemException("Cloud Datastore transaction failed to commit.", ex);
     }
   }
@@ -100,6 +128,10 @@ protected void doCommit(DefaultTransactionStatus defaultTransactionStatus)
   protected void doRollback(DefaultTransactionStatus defaultTransactionStatus)
       throws TransactionException {
     Tx tx = (Tx) defaultTransactionStatus.getTransaction();
+    RetryContext retryContext = RetrySynchronizationManager.getContext();
+    if (retryContext != null) {
+      retryContext.removeAttribute(PREVIOUS_TRANSACTION_ID_ATTRIBUTE);
+    }
     try {
       if (tx.transaction.isActive()) {
         tx.transaction.rollback();

spring-cloud-gcp-data-datastore/src/test/java/com/google/cloud/spring/data/datastore/core/DatastoreTransactionManagerTests.java

@@ -18,6 +18,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -34,6 +35,8 @@
 import org.junit.rules.ExpectedException;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.springframework.retry.RetryContext;
+import org.springframework.retry.support.RetrySynchronizationManager;
 import org.springframework.transaction.TransactionDefinition;
 import org.springframework.transaction.TransactionSystemException;
 import org.springframework.transaction.support.DefaultTransactionDefinition;
@@ -46,6 +49,8 @@ class DatastoreTransactionManagerTests {
 
   @Mock Transaction transaction;
 
+  @Mock RetryContext retryContext;
+
   private Tx tx;
 
   private DatastoreTransactionManager manager;
@@ -61,6 +66,7 @@ void initMocks() {
 
     when(this.datastore.newTransaction()).thenReturn(this.transaction);
     when(this.status.getTransaction()).thenReturn(this.tx);
+    RetrySynchronizationManager.register(retryContext);
   }
 
   @Test
@@ -148,4 +154,36 @@ void testDoRollbackNotActive() {
     this.manager.doRollback(this.status);
     verify(this.transaction, never()).rollback();
   }
+
+  @Test
+  void testFailCommitSetsPreviousTransactionId() {
+    DatastoreException exception = new DatastoreException(0, "", "");
+    when(this.transaction.isActive()).thenReturn(true);
+    when(this.transaction.commit()).thenThrow(exception);
+    this.tx.setTransaction(this.transaction);
+
+    assertThatThrownBy(() -> this.manager.doCommit(this.status))
+            .isInstanceOf(TransactionSystemException.class)
+            .hasMessage("Cloud Datastore transaction failed to commit.")
+            .hasCause(exception);
+    verify(retryContext, times(1)).setAttribute(any(), any());
+  }
+
+  @Test
+  void testFailCommitRespectsIgnoringRetry() {
+    this.manager = new DatastoreTransactionManager(() -> datastore, false);
+    this.tx = (Tx) manager.doGetTransaction();
+    when(this.status.getTransaction()).thenReturn(this.tx);
+
+    DatastoreException exception = new DatastoreException(0, "", "");
+    when(this.transaction.isActive()).thenReturn(true);
+    when(this.transaction.commit()).thenThrow(exception);
+    this.tx.setTransaction(this.transaction);
+
+    assertThatThrownBy(() -> this.manager.doCommit(this.status))
+            .isInstanceOf(TransactionSystemException.class)
+            .hasMessage("Cloud Datastore transaction failed to commit.")
+            .hasCause(exception);
+    verify(retryContext, never()).setAttribute(any(), any());
+  }
 }