refactor: use loop for transactions (#875)

dev/src/index.ts

@@ -789,7 +789,7 @@ export class Firestore {
     const defaultAttempts = 5;
     const tag = requestTag();
 
-    let attemptsRemaining: number;
+    let maxAttempts: number;
 
     if (transactionOptions) {
       validateObject('transactionOptions', transactionOptions);
@@ -798,84 +798,17 @@ export class Firestore {
         transactionOptions.maxAttempts,
         {optional: true, minValue: 1}
       );
-      attemptsRemaining = transactionOptions.maxAttempts || defaultAttempts;
+      maxAttempts = transactionOptions.maxAttempts || defaultAttempts;
     } else {
-      attemptsRemaining = defaultAttempts;
+      maxAttempts = defaultAttempts;
     }
 
+    const transaction = new Transaction(this, tag);
     return this.initializeIfNeeded(tag).then(() =>
-      this._runTransaction(updateFunction, {requestTag: tag, attemptsRemaining})
+      transaction.runTransaction(updateFunction, maxAttempts)
     );
   }
 
-  _runTransaction<T>(
-    updateFunction: (transaction: Transaction) => Promise<T>,
-    transactionOptions: {
-      requestTag: string;
-      attemptsRemaining: number;
-      previousTransaction?: Transaction;
-    }
-  ): Promise<T> {
-    const requestTag = transactionOptions.requestTag;
-    const attemptsRemaining = transactionOptions.attemptsRemaining;
-    const previousTransaction = transactionOptions.previousTransaction;
-
-    const transaction = new Transaction(this, requestTag, previousTransaction);
-    let result: Promise<T>;
-
-    return transaction
-      .begin()
-      .then(() => {
-        const promise = updateFunction(transaction);
-        result =
-          promise instanceof Promise
-            ? promise
-            : Promise.reject(
-                new Error(
-                  'You must return a Promise in your transaction()-callback.'
-                )
-              );
-        return result.catch(err => {
-          logger(
-            'Firestore.runTransaction',
-            requestTag,
-            'Rolling back transaction after callback error:',
-            err
-          );
-          // Rollback the transaction and return the failed result.
-          return transaction.rollback().then(() => {
-            return result;
-          });
-        });
-      })
-      .then(() => {
-        return transaction
-          .commit()
-          .then(() => result)
-          .catch(err => {
-            if (attemptsRemaining > 1) {
-              logger(
-                'Firestore.runTransaction',
-                requestTag,
-                `Retrying transaction after error: ${JSON.stringify(err)}.`
-              );
-              return this._runTransaction(updateFunction, {
-                previousTransaction: transaction,
-                requestTag,
-                attemptsRemaining: attemptsRemaining - 1,
-              });
-            }
-            logger(
-              'Firestore.runTransaction',
-              requestTag,
-              'Exhausted transaction retries, returning error: %s',
-              err
-            );
-            return Promise.reject(err);
-          });
-      });
-  }
-
   /**
    * Fetches the root collections that are associated with this Firestore
    * database.

dev/src/transaction.ts

@@ -18,6 +18,7 @@ import * as proto from '../protos/firestore_v1_proto_api';
 
 import {DocumentSnapshot, Precondition} from './document';
 import {Firestore, WriteBatch} from './index';
+import {logger} from './logger';
 import {FieldPath, validateFieldPath} from './path';
 import {
   DocumentReference,
@@ -70,17 +71,9 @@ export class Transaction {
    * @param firestore The Firestore Database client.
    * @param requestTag A unique client-assigned identifier for the scope of
    * this transaction.
-   * @param previousTransaction If available, the failed transaction that is
-   * being retried.
    */
-  constructor(
-    firestore: Firestore,
-    requestTag: string,
-    previousTransaction?: Transaction
-  ) {
+  constructor(firestore: Firestore, requestTag: string) {
     this._firestore = firestore;
-    this._transactionId =
-      previousTransaction && previousTransaction._transactionId;
     this._writeBatch = firestore.batch();
     this._requestTag = requestTag;
   }
@@ -398,12 +391,69 @@ export class Transaction {
   }
 
   /**
-   * Returns the tag to use with with all request for this Transaction.
+   * Executes `updateFunction()` and commits the transaction with retry.
+   *
    * @private
-   * @return A unique client-generated identifier for this transaction.
+   * @param updateFunction The user function to execute within the transaction
+   * context.
+   * @param requestTag A unique client-assigned identifier for the scope of
+   * this transaction.
+   * @param maxAttempts The maximum number of attempts for this transaction.
    */
-  get requestTag(): string {
-    return this._requestTag;
+  async runTransaction<T>(
+    updateFunction: (transaction: Transaction) => Promise<T>,
+    maxAttempts: number
+  ): Promise<T> {
+    let result: T;
+    let lastError: Error | undefined = undefined;
+
+    for (let attempt = 0; attempt < maxAttempts; ++attempt) {
+      if (lastError) {
+        logger(
+          'Firestore.runTransaction',
+          this._requestTag,
+          `Retrying transaction after error:`,
+          lastError
+        );
+      }
+
+      await this.begin();
+
+      try {
+        const promise = updateFunction(this);
+        if (!(promise instanceof Promise)) {
+          throw new Error(
+            'You must return a Promise in your transaction()-callback.'
+          );
+        }
+        result = await promise;
+      } catch (err) {
+        logger(
+          'Firestore.runTransaction',
+          this._requestTag,
+          'Rolling back transaction after callback error:',
+          err
+        );
+        await this.rollback();
+        return Promise.reject(err); // User callback failed
+      }
+
+      try {
+        await this.commit();
+        return result; // Success
+      } catch (err) {
+        lastError = err;
+        this._writeBatch._reset();
+      }
+    }
+
+    logger(
+      'Firestore.runTransaction',
+      this._requestTag,
+      'Exhausted transaction retries, returning error: %s',
+      lastError
+    );
+    return Promise.reject(lastError);
   }
 }
 

dev/src/write-batch.ts

@@ -652,6 +652,15 @@ export class WriteBatch {
 
     return true;
   }
+
+  /**
+   * Resets the WriteBatch and dequeues all pending operations.
+   * @private
+   */
+  _reset() {
+    this._ops.splice(0);
+    this._committed = false;
+  }
 }
 
 /**

dev/test/transaction.ts

@@ -517,10 +517,14 @@ describe('transaction operations', () => {
 
   it('enforce that gets come before writes', () => {
     return expect(
-      runTransaction((transaction, docRef) => {
-        transaction.set(docRef, {foo: 'bar'});
-        return transaction.get(docRef);
-      }, begin())
+      runTransaction(
+        (transaction, docRef) => {
+          transaction.set(docRef, {foo: 'bar'});
+          return transaction.get(docRef);
+        },
+        begin(),
+        rollback()
+      )
     ).to.eventually.be.rejectedWith(
       'Firestore transactions require all reads to be executed before all writes.'
     );
@@ -575,10 +579,14 @@ describe('transaction operations', () => {
 
   it('enforce that getAll come before writes', () => {
     return expect(
-      runTransaction((transaction, docRef) => {
-        transaction.set(docRef, {foo: 'bar'});
-        return transaction.getAll(docRef);
-      }, begin())
+      runTransaction(
+        (transaction, docRef) => {
+          transaction.set(docRef, {foo: 'bar'});
+          return transaction.getAll(docRef);
+        },
+        begin(),
+        rollback()
+      )
     ).to.eventually.be.rejectedWith(
       'Firestore transactions require all reads to be executed before all writes.'
     );

dev/test/write-batch.ts

@@ -356,6 +356,25 @@ describe('batch support', () => {
     return promise;
   });
 
+  it('can reset a committed batch', async () => {
+    const documentName = firestore.doc('col/doc');
+
+    const batch = firestore.batch();
+    batch.set(documentName, {foo: FieldValue.serverTimestamp()});
+    batch.update(documentName, {foo: 'bar'});
+    batch.create(documentName, {});
+    batch.delete(documentName);
+    await batch.commit();
+
+    batch._reset();
+
+    batch.set(documentName, {foo: FieldValue.serverTimestamp()});
+    batch.update(documentName, {foo: 'bar'});
+    batch.create(documentName, {});
+    batch.delete(documentName);
+    await batch.commit();
+  });
+
   it('can commit an unmodified batch multiple times', () => {
     const documentName = firestore.doc('col/doc');