Initial interface for transaction manager for manually managing retries of a transaction

google-cloud-examples/src/main/java/com/google/cloud/examples/spanner/snippets/DatabaseClientSnippets.java

@@ -23,13 +23,15 @@
 package com.google.cloud.examples.spanner.snippets;
 
 import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.AbortedException;
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.Key;
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.ReadOnlyTransaction;
 import com.google.cloud.spanner.Struct;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.TransactionContext;
+import com.google.cloud.spanner.TransactionManager;
 import com.google.cloud.spanner.TransactionRunner;
 import com.google.cloud.spanner.TransactionRunner.TransactionCallable;
 import java.util.Collections;
@@ -222,4 +224,31 @@ public Void run(TransactionContext transaction) throws Exception {
         });
     // [END readWriteTransaction]
   }
+
+  /**
+   * Example of using {@link TransactionManager}.
+   */
+  // [TARGET transactionManager()]
+  // [VARIABLE my_singer_id]
+  public void transactionManager(final long singerId) throws InterruptedException {
+    // [START transactionManager]
+    try (TransactionManager manager = dbClient.transactionManager()) {
+      TransactionContext txn = manager.begin();
+      while (true) {
+        String column = "FirstName";
+        Struct row = txn.readRow("Singers", Key.of(singerId), Collections.singleton(column));
+        String name = row.getString(column);
+        txn.buffer(
+            Mutation.newUpdateBuilder("Singers").set(column).to(name.toUpperCase()).build());
+        try {
+          manager.commit();
+          break;
+        } catch (AbortedException e) {
+          Thread.sleep(e.getRetryDelayInMillis() / 1000);
+          txn = manager.resetForRetry();
+        }
+      }
+    }
+    // [END transactionManager]
+  }
 }

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

@@ -237,6 +237,56 @@ public interface DatabaseClient {
    *     });
    * </code></pre>
    *
+   * <p>Example of a read write transaction.
+   * <pre> {@code
+   * long singerId = my_singer_id;
+   * TransactionRunner runner = dbClient.readWriteTransaction();
+   * runner.run(
+   *     new TransactionCallable<Void>() {
+   * 
+   *       @Override
+   *       public Void run(TransactionContext transaction) throws Exception {
+   *         String column = "FirstName";
+   *         Struct row =
+   *             transaction.readRow("Singers", Key.of(singerId), Collections.singleton(column));
+   *         String name = row.getString(column);
+   *         transaction.buffer(
+   *             Mutation.newUpdateBuilder("Singers").set(column).to(name.toUpperCase()).build());
+   *         return null;
+   *       }
+   *     });
+   * }</pre>
+   *
    */
   TransactionRunner readWriteTransaction();
+
+  /**
+   * Returns a transaction manager which allows manual management of transaction lifecycle. This
+   * API is meant for advanced users. Most users should instead use the
+   * {@link #readWriteTransaction()} API instead.
+   *
+   * <p>Example of using {@link TransactionManager}.
+   * <pre> {@code
+   * long singerId = my_singer_id;
+   * try (TransactionManager manager = dbClient.transactionManager()) {
+   *   TransactionContext txn = manager.begin();
+   *   while (true) {
+   *     String column = "FirstName";
+   *     Struct row = txn.readRow("Singers", Key.of(singerId), Collections.singleton(column));
+   *     String name = row.getString(column);
+   *     txn.buffer(
+   *         Mutation.newUpdateBuilder("Singers").set(column).to(name.toUpperCase()).build());
+   *     try {
+   *       manager.commit();
+   *       break;
+   *     } catch (AbortedException e) {
+   *       Thread.sleep(e.getRetryDelayInMillis() / 1000);
+   *       txn = manager.resetForRetry();
+   *     }
+   *   }
+   * }
+   * }</pre>
+   *
+   */
+  TransactionManager transactionManager();
 }

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

@@ -144,6 +144,17 @@ public TransactionRunner readWriteTransaction() {
     }
   }
 
+  @Override
+  public TransactionManager transactionManager() {
+    Span span = tracer.spanBuilder(READ_WRITE_TRANSACTION).startSpan();
+    try (Scope s = tracer.withSpan(span)) {
+      return pool.getReadWriteSession().transactionManager();
+    } catch (RuntimeException e) {
+      TraceUtil.endSpanWithFailure(span, e);
+      throw e;
+    }
+  }
+
   ListenableFuture<Void> closeAsync() {
     return pool.closeAsync();
   }

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

@@ -195,6 +195,65 @@ public Timestamp getReadTimestamp() {
       return txn.getReadTimestamp();
     }
   }
+
+  private static class AutoClosingTransactionManager implements TransactionManager {
+    final TransactionManager delegate;
+    final PooledSession session;
+
+    AutoClosingTransactionManager(TransactionManager delegate, PooledSession session) {
+      this.delegate = delegate;
+      this.session = session;
+    }
+
+    @Override
+    public TransactionContext begin() {
+      return delegate.begin();
+    }
+
+    @Override
+    public void commit() {
+      try {
+        delegate.commit();
+      } finally {
+        if (getState() != TransactionState.ABORTED) {
+          close();
+        }
+      }
+    }
+
+    @Override
+    public void rollback() {
+      try {
+        delegate.rollback();
+      } finally {
+        close();
+      }
+    }
+
+    @Override
+    public TransactionContext resetForRetry() {
+      return delegate.resetForRetry();
+    }
+
+    @Override
+    public Timestamp getCommitTimestamp() {
+      return delegate.getCommitTimestamp();
+    }
+
+    @Override
+    public void close() {
+      try {
+        delegate.close();
+      } finally {
+        session.close();
+      }
+    }
+
+    @Override
+    public TransactionState getState() {
+      return delegate.getState();
+    }
+  }
 
   // Exception class used just to track the stack trace at the point when a session was handed out
   // from the pool.
@@ -386,6 +445,12 @@ private void keepAlive() {
     private void markUsed() {
       lastUseTime = clock.instant();
     }
+
+    @Override
+    public TransactionManager transactionManager() {
+      markUsed();
+      return new AutoClosingTransactionManager(delegate.transactionManager(), this);
+    }
   }
 
   private static final class SessionOrError {

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

@@ -898,14 +898,25 @@ public Transaction call() throws Exception {
       }
     }
 
-    private <T extends SessionTransaction> T setActive(@Nullable T ctx) {
+    TransactionContextImpl newTransaction() {
+      TransactionContextImpl txn = new TransactionContextImpl(this, readyTransactionId, rpc,
+          defaultPrefetchChunks);
+      return txn;
+    }
+
+    <T extends SessionTransaction> T setActive(@Nullable T ctx) {
       if (activeTransaction != null) {
         activeTransaction.invalidate();
       }
       activeTransaction = ctx;
       readyTransactionId = null;
       return ctx;
     }
+
+    @Override
+    public TransactionManager transactionManager() {
+      return new TransactionManagerImpl(this);
+    }
   }
 
   /**
@@ -914,7 +925,7 @@ private <T extends SessionTransaction> T setActive(@Nullable T ctx) {
    * transactions, and read-write transactions. The defining characteristic is that a session may
    * only have one such transaction active at a time.
    */
-  private interface SessionTransaction {
+  static interface SessionTransaction {
     /** Invalidates the transaction, generally because a new one has been started on the session. */
     void invalidate();
   }
@@ -1018,7 +1029,7 @@ ResultSet executeQueryInternalWithOptions(
           ExecuteSqlRequest.newBuilder()
               .setSql(statement.getSql())
               .setQueryMode(queryMode)
-              .setSession(session.name);
+              .setSession(session.getName());
       Map<String, Value> stmtParameters = statement.getParameters();
       if (!stmtParameters.isEmpty()) {
         com.google.protobuf.Struct.Builder paramsBuilder = builder.getParamsBuilder();
@@ -1217,10 +1228,7 @@ void backoffSleep(Context context, long backoffMillis) {
       this.session = session;
       this.sleeper = sleeper;
       this.span = Tracing.getTracer().getCurrentSpan();
-      ByteString transactionId = session.readyTransactionId;
-      session.readyTransactionId = null;
-      this.txn = new TransactionContextImpl(session, transactionId, rpc, defaultPrefetchChunks,
-          span);
+      this.txn = session.newTransaction();
     }
 
     TransactionRunnerImpl(SessionImpl session, SpannerRpc rpc, int defaultPrefetchChunks) {
@@ -1230,7 +1238,7 @@ void backoffSleep(Context context, long backoffMillis) {
     @Nullable
     @Override
     public <T> T run(TransactionCallable<T> callable) {
-      try {
+      try (Scope s = tracer.withSpan(span)) {
         return runInternal(callable);
       } catch (RuntimeException e) {
         TraceUtil.endSpanWithFailure(span, e);
@@ -1244,6 +1252,7 @@ private <T> T runInternal(TransactionCallable<T> callable) {
       BackOff backoff = newBackOff();
       final Context context = Context.current();
       int attempt = 0;
+      // TODO: Change this to use TransactionManager.
       while (true) {
         checkState(
             isValid, "TransactionRunner has been invalidated by a new operation on the session");
@@ -1318,7 +1327,7 @@ public void invalidate() {
 
     private void backoff(Context context, BackOff backoff) {
       long delay = txn.getRetryDelayInMillis(backoff);
-      txn = new TransactionContextImpl(session, null, txn.rpc, txn.defaultPrefetchChunks, span);
+      txn = session.newTransaction();
       span.addAnnotation("Backing off",
           ImmutableMap.of("Delay", AttributeValue.longAttributeValue(delay)));
       sleeper.backoffSleep(context, delay);
@@ -1344,9 +1353,8 @@ static class TransactionContextImpl extends AbstractReadContext implements Trans
         SessionImpl session,
         @Nullable ByteString transactionId,
         SpannerRpc rpc,
-        int defaultPrefetchChunks,
-        Span span) {
-      super(session, rpc, defaultPrefetchChunks, span);
+        int defaultPrefetchChunks) {
+      super(session, rpc, defaultPrefetchChunks);
       this.transactionId = transactionId;
     }