feat: add transaction runner for connections

Adds a `runTransaction` method to `Connection` to allow applications to
execute read/write transactions that are automatically retried in the same
way as in the standard client library. This feature will be extended to
the JDBC driver, so transaction retries can be defined using a runner
there as well.

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