Add TransactionOptions to allow control over number of attempts to commit, before transaction fails.

Author: tom-andersen

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/TransactionTest.java

@@ -21,6 +21,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -29,7 +30,6 @@
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.firestore.FirebaseFirestoreException.Code;
-import com.google.firebase.firestore.core.TransactionRunner;
 import com.google.firebase.firestore.testutil.IntegrationTestUtil;
 import com.google.firebase.firestore.util.AsyncQueue.TimerId;
 import java.util.ArrayList;
@@ -651,7 +651,38 @@ public void testMakesDefaultMaxAttempts() {
 
     Exception e = waitForException(transactionTask);
     assertEquals(Code.FAILED_PRECONDITION, ((FirebaseFirestoreException) e).getCode());
-    assertEquals(TransactionRunner.DEFAULT_MAX_ATTEMPTS_COUNT, count.get());
+    assertEquals(TransactionOptions.DEFAULT_MAX_ATTEMPTS_COUNT, count.get());
+  }
+
+  @Test
+  public void testMakesOptionSpecifiedMaxAttempts() {
+    TransactionOptions options = new TransactionOptions.Builder().setMaxAttempts(1).build();
+
+    FirebaseFirestore firestore = testFirestore();
+    DocumentReference doc1 = firestore.collection("counters").document();
+    AtomicInteger count = new AtomicInteger(0);
+    waitFor(doc1.set(map("count", 15)));
+    Task<Void> transactionTask =
+        firestore.runTransaction(
+            options,
+            transaction -> {
+              // Get the first doc.
+              transaction.get(doc1);
+              // Do a write outside of the transaction to cause the transaction to fail.
+              waitFor(doc1.set(map("count", 1234 + count.incrementAndGet())));
+              return null;
+            });
+
+    Exception e = waitForException(transactionTask);
+    assertEquals(Code.FAILED_PRECONDITION, ((FirebaseFirestoreException) e).getCode());
+    assertEquals(options.getMaxAttempts(), count.get());
+  }
+
+  @Test
+  public void testTransactionOptionsZeroMaxAttempts_shouldThrowIllegalArgumentException() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> new TransactionOptions.Builder().setMaxAttempts(0).build());
   }
 
   @Test

firebase-firestore/src/main/java/com/google/firebase/firestore/FirebaseFirestore.java

@@ -419,7 +419,7 @@ public Query collectionGroup(@NonNull String collectionId) {
    * @return The task returned from the updateFunction.
    */
   private <ResultT> Task<ResultT> runTransaction(
-      Transaction.Function<ResultT> updateFunction, Executor executor) {
+      TransactionOptions options, Transaction.Function<ResultT> updateFunction, Executor executor) {
     ensureClientConfigured();
 
     // We wrap the function they provide in order to
@@ -434,7 +434,7 @@ private <ResultT> Task<ResultT> runTransaction(
                     updateFunction.apply(
                         new Transaction(internalTransaction, FirebaseFirestore.this)));
 
-    return client.transaction(wrappedUpdateFunction);
+    return client.transaction(options, wrappedUpdateFunction);
   }
 
   /**
@@ -448,9 +448,26 @@ private <ResultT> Task<ResultT> runTransaction(
   @NonNull
   public <TResult> Task<TResult> runTransaction(
       @NonNull Transaction.Function<TResult> updateFunction) {
+    return runTransaction(TransactionOptions.DEFAULT, updateFunction);
+  }
+
+  /**
+   * Executes the given updateFunction and then attempts to commit the changes applied within the
+   * transaction. If any document read within the transaction has changed, the updateFunction will
+   * be retried. If it fails to commit after the maxmimum number of attempts specified in
+   * transactionOptions, the transaction will fail.
+   *
+   * @param options The function to execute within the transaction context.
+   * @param updateFunction The function to execute within the transaction context.
+   * @return The task returned from the updateFunction.
+   */
+  public <TResult> Task<TResult> runTransaction(
+      @NonNull TransactionOptions options, @NonNull Transaction.Function<TResult> updateFunction) {
     checkNotNull(updateFunction, "Provided transaction update function must not be null.");
     return runTransaction(
-        updateFunction, com.google.firebase.firestore.core.Transaction.getDefaultExecutor());
+        options,
+        updateFunction,
+        com.google.firebase.firestore.core.Transaction.getDefaultExecutor());
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/TransactionOptions.java

@@ -0,0 +1,96 @@
+// Copyright 2022 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.firebase.firestore;
+
+/**
+ * Parameter for {@link FirebaseFirestore#runTransaction(TransactionOptions, Transaction.Function)}.
+ */
+public final class TransactionOptions {
+
+  static final TransactionOptions DEFAULT = new TransactionOptions.Builder().build();
+  static final int DEFAULT_MAX_ATTEMPTS_COUNT = 5;
+
+  private final int maxAttempts;
+
+  private TransactionOptions(int maxAttempts) {
+    this.maxAttempts = maxAttempts;
+  }
+
+  /** A Builder for creating {@code TransactionOptions}. */
+  public static final class Builder {
+    private int maxAttempts = DEFAULT_MAX_ATTEMPTS_COUNT;
+
+    /** Constructs a new {@code TransactionOptions} Builder object. */
+    public Builder() {}
+
+    /**
+     * Constructs a new {@code TransactionOptions} Builder based on an existing {@code
+     * TransactionOptions} object.
+     */
+    public Builder(TransactionOptions options) {
+      maxAttempts = options.maxAttempts;
+    }
+
+    /**
+     * Set maximum number of attempts to commit, after which transaction fails. Default is 5.
+     *
+     * @return this builder
+     */
+    public Builder setMaxAttempts(int maxAttempts) {
+      if (maxAttempts < 1) throw new IllegalArgumentException("Max attempts must be at least 1");
+      this.maxAttempts = maxAttempts;
+      return this;
+    }
+
+
+    /**
+     * Build the {@code TransactionOptions} object.
+     *
+     * @return the built {@code TransactionOptions} object
+     */
+    public TransactionOptions build() {
+      return new TransactionOptions(maxAttempts);
+    }
+  }
+
+  /**
+   * Get maximum number of attempts to commit, after which transaction fails. Default is 5.
+   *
+   * @return maximum number of attempts
+   */
+  public int getMaxAttempts() {
+    return maxAttempts;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+
+    TransactionOptions that = (TransactionOptions) o;
+
+    return maxAttempts == that.maxAttempts;
+  }
+
+  @Override
+  public int hashCode() {
+    return maxAttempts;
+  }
+
+  @Override
+  public String toString() {
+    return "TransactionOptions{" + "maxAttempts=" + maxAttempts + '}';
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/core/FirestoreClient.java

@@ -26,6 +26,7 @@
 import com.google.firebase.firestore.FirebaseFirestoreException.Code;
 import com.google.firebase.firestore.FirebaseFirestoreSettings;
 import com.google.firebase.firestore.LoadBundleTask;
+import com.google.firebase.firestore.TransactionOptions;
 import com.google.firebase.firestore.auth.CredentialsProvider;
 import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.bundle.BundleReader;
@@ -228,10 +229,12 @@ public Task<Void> write(final List<Mutation> mutations) {
   }
 
   /** Tries to execute the transaction in updateFunction. */
-  public <TResult> Task<TResult> transaction(Function<Transaction, Task<TResult>> updateFunction) {
+  public <TResult> Task<TResult> transaction(
+      TransactionOptions options, Function<Transaction, Task<TResult>> updateFunction) {
     this.verifyNotTerminated();
     return AsyncQueue.callTask(
-        asyncQueue.getExecutor(), () -> syncEngine.transaction(asyncQueue, updateFunction));
+        asyncQueue.getExecutor(),
+        () -> syncEngine.transaction(asyncQueue, options, updateFunction));
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/core/SyncEngine.java

@@ -26,6 +26,7 @@
 import com.google.firebase.firestore.FirebaseFirestoreException;
 import com.google.firebase.firestore.LoadBundleTask;
 import com.google.firebase.firestore.LoadBundleTaskProgress;
+import com.google.firebase.firestore.TransactionOptions;
 import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.bundle.BundleElement;
 import com.google.firebase.firestore.bundle.BundleLoader;
@@ -307,8 +308,10 @@ private void addUserCallback(int batchId, TaskCompletionSource<Void> userTask) {
    * <p>The Task returned is resolved when the transaction is fully committed.
    */
   public <TResult> Task<TResult> transaction(
-      AsyncQueue asyncQueue, Function<Transaction, Task<TResult>> updateFunction) {
-    return new TransactionRunner<TResult>(asyncQueue, remoteStore, updateFunction).run();
+      AsyncQueue asyncQueue,
+      TransactionOptions options,
+      Function<Transaction, Task<TResult>> updateFunction) {
+    return new TransactionRunner<TResult>(asyncQueue, remoteStore, options, updateFunction).run();
   }
 
   /** Called by FirestoreClient to notify us of a new remote event. */

firebase-firestore/src/main/java/com/google/firebase/firestore/core/TransactionRunner.java

@@ -18,6 +18,7 @@
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.firebase.firestore.FirebaseFirestoreException;
+import com.google.firebase.firestore.TransactionOptions;
 import com.google.firebase.firestore.remote.Datastore;
 import com.google.firebase.firestore.remote.RemoteStore;
 import com.google.firebase.firestore.util.AsyncQueue;
@@ -27,7 +28,6 @@
 
 /** TransactionRunner encapsulates the logic needed to run and retry transactions with backoff. */
 public class TransactionRunner<TResult> {
-  public static final int DEFAULT_MAX_ATTEMPTS_COUNT = 5;
   private AsyncQueue asyncQueue;
   private RemoteStore remoteStore;
   private Function<Transaction, Task<TResult>> updateFunction;
@@ -39,12 +39,13 @@ public class TransactionRunner<TResult> {
   public TransactionRunner(
       AsyncQueue asyncQueue,
       RemoteStore remoteStore,
+      TransactionOptions options,
       Function<Transaction, Task<TResult>> updateFunction) {
 
     this.asyncQueue = asyncQueue;
     this.remoteStore = remoteStore;
     this.updateFunction = updateFunction;
-    this.attemptsRemaining = DEFAULT_MAX_ATTEMPTS_COUNT;
+    this.attemptsRemaining = options.getMaxAttempts();
 
     backoff = new ExponentialBackoff(asyncQueue, TimerId.RETRY_TRANSACTION);
   }