feat: add transaction runner for connections

google-cloud-spanner/clirr-ignored-differences.xml

@@ -814,6 +814,12 @@
     <className>com/google/cloud/spanner/connection/TransactionRetryListener</className>
     <method>void retryDmlAsPartitionedDmlFailed(java.util.UUID, com.google.cloud.spanner.Statement, java.lang.Throwable)</method>
   </difference>
-
+
+  <!-- Added transaction runner. -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/connection/Connection</className>
+    <method>java.lang.Object runTransaction(com.google.cloud.spanner.connection.Connection$TransactionCallable)</method>
+  </difference>
 
 </differences>

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java

@@ -99,7 +99,7 @@ public void rollback() {
   public TransactionContext resetForRetry() {
     if (txn == null || !txn.isAborted() && txnState != TransactionState.ABORTED) {
       throw new IllegalStateException(
-          "resetForRetry can only be called if the previous attempt" + " aborted");
+          "resetForRetry can only be called if the previous attempt aborted");
     }
     try (IScope s = tracer.withSpan(span)) {
       boolean useInlinedBegin = txn.transactionId != null;

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java

@@ -835,6 +835,21 @@ default boolean isKeepTransactionAlive() {
    */
   ApiFuture<Void> rollbackAsync();
 
+  /** Functional interface for the {@link #runTransaction(TransactionCallable)} method. */
+  interface TransactionCallable<T> {
+    /** This method is invoked with a fresh transaction on the connection. */
+    T run(Connection transaction);
+  }
+
+  /**
+   * Runs the given callable in a transaction. The transaction type is determined by the current
+   * state of the connection. That is; if the connection is in read/write mode, the transaction type
+   * will be a read/write transaction. If the connection is in read-only mode, it will be a
+   * read-only transaction. The transaction will automatically be retried if it is aborted by
+   * Spanner.
+   */
+  <T> T runTransaction(TransactionCallable<T> callable);
+
   /** Returns the current savepoint support for this connection. */
   SavepointSupport getSavepointSupport();
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

@@ -194,6 +194,11 @@ private LeakedConnectionException() {
    */
   private final ConnectionOptions options;
 
+  enum Caller {
+    APPLICATION,
+    TRANSACTION_RUNNER,
+  }
+
   /** The supported batch modes. */
   enum BatchMode {
     NONE,
@@ -267,6 +272,9 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
    */
   private boolean transactionBeginMarked = false;
 
+  /** This field is set to true when a transaction runner is active for this connection. */
+  private boolean transactionRunnerActive = false;
+
   private BatchMode batchMode;
   private UnitOfWorkType unitOfWorkType;
   private final Stack<UnitOfWork> transactionStack = new Stack<>();
@@ -1164,16 +1172,19 @@ public void onFailure() {
 
   @Override
   public void commit() {
-    get(commitAsync(CallType.SYNC));
+    get(commitAsync(CallType.SYNC, Caller.APPLICATION));
   }
 
   @Override
   public ApiFuture<Void> commitAsync() {
-    return commitAsync(CallType.ASYNC);
+    return commitAsync(CallType.ASYNC, Caller.APPLICATION);
   }
 
-  private ApiFuture<Void> commitAsync(CallType callType) {
+  ApiFuture<Void> commitAsync(CallType callType, Caller caller) {
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ConnectionPreconditions.checkState(
+        !transactionRunnerActive || caller == Caller.TRANSACTION_RUNNER,
+        "Cannot call commit when a transaction runner is active");
     maybeAutoCommitOrFlushCurrentUnitOfWork(COMMIT_STATEMENT.getType(), COMMIT_STATEMENT);
     return endCurrentTransactionAsync(callType, commit, COMMIT_STATEMENT);
   }
@@ -1201,16 +1212,19 @@ public void onFailure() {
 
   @Override
   public void rollback() {
-    get(rollbackAsync(CallType.SYNC));
+    get(rollbackAsync(CallType.SYNC, Caller.APPLICATION));
   }
 
   @Override
   public ApiFuture<Void> rollbackAsync() {
-    return rollbackAsync(CallType.ASYNC);
+    return rollbackAsync(CallType.ASYNC, Caller.APPLICATION);
   }
 
-  private ApiFuture<Void> rollbackAsync(CallType callType) {
+  ApiFuture<Void> rollbackAsync(CallType callType, Caller caller) {
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ConnectionPreconditions.checkState(
+        !transactionRunnerActive || caller == Caller.TRANSACTION_RUNNER,
+        "Cannot call rollback when a transaction runner is active");
     maybeAutoCommitOrFlushCurrentUnitOfWork(ROLLBACK_STATEMENT.getType(), ROLLBACK_STATEMENT);
     return endCurrentTransactionAsync(callType, rollback, ROLLBACK_STATEMENT);
   }
@@ -1243,6 +1257,27 @@ private ApiFuture<Void> endCurrentTransactionAsync(
     return res;
   }
 
+  @Override
+  public <T> T runTransaction(TransactionCallable<T> callable) {
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ConnectionPreconditions.checkState(!isBatchActive(), "Cannot run transaction while in a batch");
+    ConnectionPreconditions.checkState(
+        !isTransactionStarted(), "Cannot run transaction when a transaction is already active");
+    ConnectionPreconditions.checkState(
+        !transactionRunnerActive, "A transaction runner is already active for this connection");
+    this.transactionRunnerActive = true;
+    try {
+      return new TransactionRunnerImpl(this).run(callable);
+    } finally {
+      this.transactionRunnerActive = false;
+    }
+  }
+
+  void resetForRetry(UnitOfWork retryUnitOfWork) {
+    retryUnitOfWork.resetForRetry();
+    this.currentUnitOfWork = retryUnitOfWork;
+  }
+
   @Override
   public SavepointSupport getSavepointSupport() {
     return getConnectionPropertyValue(SAVEPOINT_SUPPORT);
@@ -2000,7 +2035,7 @@ private UnitOfWork maybeStartAutoDmlBatch(UnitOfWork transaction) {
     return transaction;
   }
 
-  private UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() {
+  UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() {
     return getCurrentUnitOfWorkOrStartNewUnitOfWork(
         StatementType.UNKNOWN, /* parsedStatement = */ null, /* internalMetadataQuery = */ false);
   }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java

@@ -1261,6 +1261,11 @@ private ApiFuture<Void> rollbackAsync(CallType callType, boolean updateStatusAnd
     }
   }
 
+  @Override
+  public void resetForRetry() {
+    txContextFuture = ApiFutures.immediateFuture(txManager.resetForRetry());
+  }
+
   @Override
   String getUnitOfWorkName() {
     return "read/write transaction";

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/TransactionRunnerImpl.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2024 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.spanner.connection;
+
+import static com.google.cloud.spanner.SpannerApiFutures.get;
+
+import com.google.cloud.spanner.AbortedException;
+import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.cloud.spanner.connection.Connection.TransactionCallable;
+import com.google.cloud.spanner.connection.ConnectionImpl.Caller;
+import com.google.cloud.spanner.connection.UnitOfWork.CallType;
+
+class TransactionRunnerImpl {
+  private final ConnectionImpl connection;
+
+  TransactionRunnerImpl(ConnectionImpl connection) {
+    this.connection = connection;
+  }
+
+  <T> T run(TransactionCallable<T> callable) {
+    connection.beginTransaction();
+    // Disable internal retries during this transaction.
+    connection.setRetryAbortsInternally(/* retryAbortsInternally = */ false, /* local = */ true);
+    UnitOfWork transaction = connection.getCurrentUnitOfWorkOrStartNewUnitOfWork();
+    while (true) {
+      try {
+        T result = callable.run(connection);
+        get(connection.commitAsync(CallType.SYNC, Caller.TRANSACTION_RUNNER));
+        return result;
+      } catch (AbortedException abortedException) {
+        try {
+          //noinspection BusyWait
+          Thread.sleep(abortedException.getRetryDelayInMillis());
+          connection.resetForRetry(transaction);
+        } catch (InterruptedException interruptedException) {
+          connection.rollbackAsync(CallType.SYNC, Caller.TRANSACTION_RUNNER);
+          throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
+        } catch (Throwable t) {
+          connection.rollbackAsync(CallType.SYNC, Caller.TRANSACTION_RUNNER);
+          throw t;
+        }
+      } catch (Throwable t) {
+        connection.rollbackAsync(CallType.SYNC, Caller.TRANSACTION_RUNNER);
+        throw t;
+      }
+    }
+  }
+}

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java

@@ -125,6 +125,10 @@ interface EndTransactionCallback {
   ApiFuture<Void> rollbackAsync(
       @Nonnull CallType callType, @Nonnull EndTransactionCallback callback);
 
+  default void resetForRetry() {
+    throw new UnsupportedOperationException();
+  }
+
   /** @see Connection#savepoint(String) */
   void savepoint(@Nonnull String name, @Nonnull Dialect dialect);