From 259ece8511db71ea79cc1a080eb785a15db88756 Mon Sep 17 00:00:00 2001
From: Brian Chen <chenbrian@google.com>
Date: Thu, 1 Oct 2020 17:17:48 -0500
Subject: [PATCH] fix: bulkWriter: writing to the same doc doesn't create a new
 batch (#394)

* fix: bulkWriter: writing to the same doc doesn't create a new batch

* fix clirr rules

* resolve comments

* remove usage of BatchWriteResult.documentKey

* clirr rule update
---
 .../clirr-ignored-differences.xml             |  13 +-
 .../cloud/firestore/BatchWriteResult.java     |  13 +-
 .../cloud/firestore/BulkCommitBatch.java      |  28 +----
 .../google/cloud/firestore/BulkWriter.java    |  85 ++++----------
 .../google/cloud/firestore/UpdateBuilder.java | 111 +++++++++++-------
 .../cloud/firestore/BulkWriterTest.java       |  77 +++---------
 6 files changed, 129 insertions(+), 198 deletions(-)

diff --git a/google-cloud-firestore/clirr-ignored-differences.xml b/google-cloud-firestore/clirr-ignored-differences.xml
index d3fc69ba9..42302e973 100644
--- a/google-cloud-firestore/clirr-ignored-differences.xml
+++ b/google-cloud-firestore/clirr-ignored-differences.xml
@@ -161,7 +161,7 @@
     <method>com.google.cloud.firestore.Query collectionGroup(java.lang.String)</method>
     <to>com.google.cloud.firestore.CollectionGroup</to>
   </difference>
-  
+
   <!--
   BulkWriter
   -->
@@ -199,6 +199,17 @@
     <className>com/google/cloud/firestore/spi/v1/FirestoreRpc</className>
     <method>com.google.api.gax.rpc.UnaryCallable batchWriteCallable()</method>
   </difference>
+  <difference>
+    <differenceType>1001</differenceType>
+    <className>com/google/cloud/firestore/BatchWriteResult</className>
+  </difference>
+  <difference>
+    <differenceType>6004</differenceType>
+    <className>com/google/cloud/firestore/UpdateBuilder</className>
+    <field>pendingOperations</field>
+    <from>java.util.Map</from>
+    <to>java.util.List</to>
+  </difference>
 
   <!--
   FakeCredentials Refactor
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BatchWriteResult.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BatchWriteResult.java
index 96ce5d60b..a9847d79c 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BatchWriteResult.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BatchWriteResult.java
@@ -25,24 +25,15 @@
  * BatchWriteRequests.
  */
 @InternalApi
-public final class BatchWriteResult {
-  private final DocumentReference documentReference;
+final class BatchWriteResult {
   @Nullable private final Timestamp writeTime;
   @Nullable private final Exception exception;
 
-  BatchWriteResult(
-      DocumentReference documentReference,
-      @Nullable Timestamp timestamp,
-      @Nullable Exception exception) {
-    this.documentReference = documentReference;
+  BatchWriteResult(@Nullable Timestamp timestamp, @Nullable Exception exception) {
     this.writeTime = timestamp;
     this.exception = exception;
   }
 
-  public DocumentReference getDocumentReference() {
-    return documentReference;
-  }
-
   @Nullable
   public Timestamp getWriteTime() {
     return writeTime;
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java
index 701e6114f..8ef106fc3 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java
@@ -18,9 +18,6 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Predicate;
-import com.google.common.collect.FluentIterable;
-import java.util.Set;
 
 /** Used to represent a batch on the BatchQueue. */
 class BulkCommitBatch extends UpdateBuilder<ApiFuture<WriteResult>> {
@@ -29,21 +26,13 @@ class BulkCommitBatch extends UpdateBuilder<ApiFuture<WriteResult>> {
     super(firestore, maxBatchSize);
   }
 
-  BulkCommitBatch(
-      FirestoreImpl firestore,
-      BulkCommitBatch retryBatch,
-      final Set<DocumentReference> docsToRetry) {
+  BulkCommitBatch(FirestoreImpl firestore, BulkCommitBatch retryBatch) {
     super(firestore);
-    this.writes.addAll(
-        FluentIterable.from(retryBatch.writes)
-            .filter(
-                new Predicate<WriteOperation>() {
-                  @Override
-                  public boolean apply(WriteOperation writeOperation) {
-                    return docsToRetry.contains(writeOperation.documentReference);
-                  }
-                })
-            .toList());
+
+    // Create a new BulkCommitBatch containing only the indexes from the provided indexes to retry.
+    for (int index : retryBatch.getPendingIndexes()) {
+      this.writes.add(retryBatch.writes.get(index));
+    }
 
     Preconditions.checkState(
         retryBatch.state == BatchState.SENT,
@@ -55,9 +44,4 @@ public boolean apply(WriteOperation writeOperation) {
   ApiFuture<WriteResult> wrapResult(ApiFuture<WriteResult> result) {
     return result;
   }
-
-  @Override
-  boolean allowDuplicateDocs() {
-    return false;
-  }
 }
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java
index 5a90bd240..1fcfada8e 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java
@@ -32,7 +32,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
@@ -129,7 +128,7 @@ final class BulkWriter implements AutoCloseable {
   public ApiFuture<WriteResult> create(
       @Nonnull DocumentReference documentReference, @Nonnull Map<String, Object> fields) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.create(documentReference, fields);
     sendReadyBatches();
     return future;
@@ -147,7 +146,7 @@ public ApiFuture<WriteResult> create(
   public ApiFuture<WriteResult> create(
       @Nonnull DocumentReference documentReference, @Nonnull Object pojo) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.create(documentReference, pojo);
     sendReadyBatches();
     return future;
@@ -163,7 +162,7 @@ public ApiFuture<WriteResult> create(
   @Nonnull
   public ApiFuture<WriteResult> delete(@Nonnull DocumentReference documentReference) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.delete(documentReference);
     sendReadyBatches();
     return future;
@@ -181,7 +180,7 @@ public ApiFuture<WriteResult> delete(@Nonnull DocumentReference documentReferenc
   public ApiFuture<WriteResult> delete(
       @Nonnull DocumentReference documentReference, @Nonnull Precondition precondition) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.delete(documentReference, precondition);
     sendReadyBatches();
     return future;
@@ -200,7 +199,7 @@ public ApiFuture<WriteResult> delete(
   public ApiFuture<WriteResult> set(
       @Nonnull DocumentReference documentReference, @Nonnull Map<String, Object> fields) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.set(documentReference, fields);
     sendReadyBatches();
     return future;
@@ -222,7 +221,7 @@ public ApiFuture<WriteResult> set(
       @Nonnull Map<String, Object> fields,
       @Nonnull SetOptions options) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.set(documentReference, fields, options);
     sendReadyBatches();
     return future;
@@ -242,7 +241,7 @@ public ApiFuture<WriteResult> set(
   public ApiFuture<WriteResult> set(
       @Nonnull DocumentReference documentReference, Object pojo, @Nonnull SetOptions options) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.set(documentReference, pojo, options);
     sendReadyBatches();
     return future;
@@ -259,7 +258,7 @@ public ApiFuture<WriteResult> set(
   @Nonnull
   public ApiFuture<WriteResult> set(@Nonnull DocumentReference documentReference, Object pojo) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.set(documentReference, pojo);
     sendReadyBatches();
     return future;
@@ -282,7 +281,7 @@ public ApiFuture<WriteResult> set(@Nonnull DocumentReference documentReference,
   public ApiFuture<WriteResult> update(
       @Nonnull DocumentReference documentReference, @Nonnull Map<String, Object> fields) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.update(documentReference, fields);
     sendReadyBatches();
     return future;
@@ -308,7 +307,7 @@ public ApiFuture<WriteResult> update(
       @Nonnull Map<String, Object> fields,
       Precondition precondition) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future = bulkCommitBatch.update(documentReference, fields, precondition);
     sendReadyBatches();
     return future;
@@ -336,7 +335,7 @@ public ApiFuture<WriteResult> update(
       @Nullable Object value,
       Object... moreFieldsAndValues) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future =
         bulkCommitBatch.update(documentReference, field, value, moreFieldsAndValues);
     sendReadyBatches();
@@ -365,7 +364,7 @@ public ApiFuture<WriteResult> update(
       @Nullable Object value,
       Object... moreFieldsAndValues) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future =
         bulkCommitBatch.update(documentReference, fieldPath, value, moreFieldsAndValues);
     sendReadyBatches();
@@ -395,7 +394,7 @@ public ApiFuture<WriteResult> update(
       @Nullable Object value,
       Object... moreFieldsAndValues) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future =
         bulkCommitBatch.update(documentReference, precondition, field, value, moreFieldsAndValues);
     sendReadyBatches();
@@ -426,7 +425,7 @@ public ApiFuture<WriteResult> update(
       @Nullable Object value,
       Object... moreFieldsAndValues) {
     verifyNotClosed();
-    BulkCommitBatch bulkCommitBatch = getEligibleBatch(documentReference);
+    BulkCommitBatch bulkCommitBatch = getEligibleBatch();
     ApiFuture<WriteResult> future =
         bulkCommitBatch.update(
             documentReference, precondition, fieldPath, value, moreFieldsAndValues);
@@ -493,11 +492,10 @@ private void verifyNotClosed() {
    * Return the first eligible batch that can hold a write to the provided reference, or creates one
    * if no eligible batches are found.
    */
-  private BulkCommitBatch getEligibleBatch(DocumentReference documentReference) {
+  private BulkCommitBatch getEligibleBatch() {
     if (batchQueue.size() > 0) {
       BulkCommitBatch lastBatch = batchQueue.get(batchQueue.size() - 1);
-      if (lastBatch.getState() == UpdateBuilder.BatchState.OPEN
-          && !lastBatch.hasDocument(documentReference)) {
+      if (lastBatch.getState() == UpdateBuilder.BatchState.OPEN) {
         return lastBatch;
       }
     }
@@ -538,7 +536,8 @@ public boolean apply(BulkCommitBatch batch) {
             .toList();
 
     int index = 0;
-    while (index < unsentBatches.size() && isBatchSendable(unsentBatches.get(index))) {
+    while (index < unsentBatches.size()
+        && unsentBatches.get(index).state == BatchState.READY_TO_SEND) {
       final BulkCommitBatch batch = unsentBatches.get(index);
 
       // Send the batch if it is under the rate limit, or schedule another attempt after the
@@ -631,8 +630,8 @@ public ApiFuture<Void> apply(Void ignored) {
                 public ApiFuture<List<BatchWriteResult>> apply(Exception exception) {
                   List<BatchWriteResult> results = new ArrayList<>();
                   // If the BatchWrite RPC fails, map the exception to each individual result.
-                  for (DocumentReference documentReference : batch.getPendingDocuments()) {
-                    results.add(new BatchWriteResult(documentReference, null, exception));
+                  for (int i = 0; i < batch.getPendingDocumentPaths().size(); ++i) {
+                    results.add(new BatchWriteResult(null, exception));
                   }
                   return ApiFutures.immediateFuture(results);
                 }
@@ -655,8 +654,8 @@ public ProcessBulkCommitCallback(BulkCommitBatch batch, int attempt) {
 
     @Override
     public ApiFuture<Void> apply(List<BatchWriteResult> results) {
-      batch.processResults(results);
-      Set<DocumentReference> remainingOps = batch.getPendingDocuments();
+      batch.processResults(results, /* allowRetry= */ true);
+      List<String> remainingOps = batch.getPendingDocumentPaths();
       if (!remainingOps.isEmpty()) {
         logger.log(
             Level.WARNING,
@@ -666,52 +665,16 @@ public ApiFuture<Void> apply(List<BatchWriteResult> results) {
 
         if (attempt < MAX_RETRY_ATTEMPTS) {
           nextAttempt = backoff.createNextAttempt(nextAttempt);
-          BulkCommitBatch newBatch = new BulkCommitBatch(firestore, batch, remainingOps);
+          BulkCommitBatch newBatch = new BulkCommitBatch(firestore, batch);
           return bulkCommit(newBatch, attempt + 1);
         } else {
-          batch.failRemainingOperations(results);
+          batch.processResults(results, /* allowRetry= */ false);
         }
       }
       return ApiFutures.immediateFuture(null);
     }
   }
 
-  /**
-   * Checks that the provided batch is sendable. To be sendable, a batch must: (1) be marked as
-   * READY_TO_SEND (2) not write to references that are currently in flight.
-   */
-  private boolean isBatchSendable(BulkCommitBatch batch) {
-    if (!batch.getState().equals(UpdateBuilder.BatchState.READY_TO_SEND)) {
-      return false;
-    }
-
-    for (final DocumentReference documentReference : batch.getPendingDocuments()) {
-      boolean isRefInFlight =
-          FluentIterable.from(batchQueue)
-              .anyMatch(
-                  new Predicate<BulkCommitBatch>() {
-                    @Override
-                    public boolean apply(BulkCommitBatch batch) {
-                      return batch.getState().equals(BatchState.SENT)
-                          && batch.hasDocument(documentReference);
-                    }
-                  });
-
-      if (isRefInFlight) {
-        logger.log(
-            Level.WARNING,
-            String.format(
-                "Duplicate write to document %s detected. Writing to the same document multiple"
-                    + " times will slow down BulkWriter. Write to unique documents in order to "
-                    + "maximize throughput.",
-                documentReference.getPath()));
-        return false;
-      }
-    }
-
-    return true;
-  }
-
   @VisibleForTesting
   void setMaxBatchSize(int size) {
     maxBatchSize = size;
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java
index 056ce6840..fc9e04902 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java
@@ -25,7 +25,9 @@
 import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.UserDataConverter.EncodingOptions;
 import com.google.cloud.firestore.v1.FirestoreSettings;
+import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.firestore.v1.BatchWriteRequest;
@@ -38,7 +40,6 @@
 import io.opencensus.trace.AttributeValue;
 import io.opencensus.trace.Tracing;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -65,14 +66,31 @@ static class WriteOperation {
     }
   }
 
+  /**
+   * Used to represent a pending write operation.
+   *
+   * <p>Contains a pending write's index, document path, and the corresponding result.
+   */
+  static class PendingOperation {
+    int writeBatchIndex;
+    String documentKey;
+    SettableApiFuture<WriteResult> future;
+
+    PendingOperation(
+        int writeBatchIndex, String documentKey, SettableApiFuture<WriteResult> future) {
+      this.writeBatchIndex = writeBatchIndex;
+      this.documentKey = documentKey;
+      this.future = future;
+    }
+  }
+
   final FirestoreImpl firestore;
   protected final List<WriteOperation> writes;
   private boolean committed;
   private final int maxBatchSize;
 
   protected BatchState state = BatchState.OPEN;
-  protected Map<DocumentReference, SettableApiFuture<WriteResult>> pendingOperations =
-      new HashMap<>();
+  protected List<PendingOperation> pendingOperations = new ArrayList<>();
 
   /**
    * Used to represent the state of batch.
@@ -105,11 +123,6 @@ enum BatchState {
    */
   abstract T wrapResult(ApiFuture<WriteResult> result);
 
-  /** Whether to allow multiple writes to the same document in a batch. */
-  boolean allowDuplicateDocs() {
-    return true;
-  }
-
   /**
    * Turns a field map that contains field paths into a nested map. Turns {a.b : c} into {a : {b :
    * c}}.
@@ -712,8 +725,7 @@ public List<BatchWriteResult> apply(BatchWriteResponse batchWriteResponse) {
               } else {
                 exception = FirestoreException.serverRejected(code, status.getMessage());
               }
-              result.add(
-                  new BatchWriteResult(writes.get(i).documentReference, updateTime, exception));
+              result.add(new BatchWriteResult(updateTime, exception));
             }
 
             return result;
@@ -736,16 +748,40 @@ BatchState getState() {
     return state;
   }
 
-  boolean hasDocument(DocumentReference documentReference) {
-    return pendingOperations.containsKey(documentReference);
+  List<String> getPendingDocumentPaths() {
+    return FluentIterable.from(pendingOperations)
+        .transform(
+            new Function<PendingOperation, String>() {
+              @Override
+              public String apply(PendingOperation input) {
+                return input.documentKey;
+              }
+            })
+        .toList();
   }
 
-  Set<DocumentReference> getPendingDocuments() {
-    return pendingOperations.keySet();
+  List<Integer> getPendingIndexes() {
+    return FluentIterable.from(pendingOperations)
+        .transform(
+            new Function<PendingOperation, Integer>() {
+              @Override
+              public Integer apply(PendingOperation input) {
+                return input.writeBatchIndex;
+              }
+            })
+        .toList();
   }
 
-  Collection<SettableApiFuture<WriteResult>> getPendingFutures() {
-    return pendingOperations.values();
+  List<SettableApiFuture<WriteResult>> getPendingFutures() {
+    return FluentIterable.from(pendingOperations)
+        .transform(
+            new Function<PendingOperation, SettableApiFuture<WriteResult>>() {
+              @Override
+              public SettableApiFuture<WriteResult> apply(PendingOperation input) {
+                return input.future;
+              }
+            })
+        .toList();
   }
 
   int getPendingOperationCount() {
@@ -753,12 +789,11 @@ int getPendingOperationCount() {
   }
 
   private ApiFuture<WriteResult> processOperation(DocumentReference documentReference) {
-    Preconditions.checkState(
-        allowDuplicateDocs() || !pendingOperations.containsKey(documentReference),
-        "Batch should not contain writes to the same document");
     Preconditions.checkState(state == BatchState.OPEN, "Batch should be OPEN when adding writes");
     SettableApiFuture<WriteResult> result = SettableApiFuture.create();
-    pendingOperations.put(documentReference, result);
+    int nextBatchIndex = getPendingOperationCount();
+    pendingOperations.add(
+        new PendingOperation(nextBatchIndex, documentReference.getPath(), result));
 
     if (getPendingOperationCount() == maxBatchSize) {
       state = BatchState.READY_TO_SEND;
@@ -771,19 +806,24 @@ private ApiFuture<WriteResult> processOperation(DocumentReference documentRefere
    * Resolves the individual operations in the batch with the results and removes the entry from the
    * pendingOperations map if the result is not retryable.
    */
-  void processResults(List<BatchWriteResult> results) {
-    for (BatchWriteResult result : results) {
-      if (result.getException() == null || !shouldRetry(result.getException())) {
-        convertBatchWriteResult(result, pendingOperations.get(result.getDocumentReference()));
-        pendingOperations.remove(result.getDocumentReference());
+  void processResults(List<BatchWriteResult> results, boolean allowRetry) {
+    List<PendingOperation> newPendingOperations = new ArrayList<>();
+    for (int i = 0; i < results.size(); ++i) {
+      BatchWriteResult result = results.get(i);
+      PendingOperation operation = pendingOperations.get(i);
+
+      if (result.getException() == null) {
+        operation.future.set(new WriteResult(result.getWriteTime()));
+      } else if (!allowRetry || !shouldRetry(result.getException())) {
+        operation.future.setException(result.getException());
+      } else {
+        // Retry the operation if it has not been processed. Store the current index of
+        // pendingOperations to preserve the mapping of this operation's index in the underlying
+        // writes array.
+        newPendingOperations.add(new PendingOperation(i, operation.documentKey, operation.future));
       }
     }
-  }
-
-  void failRemainingOperations(List<BatchWriteResult> results) {
-    for (BatchWriteResult result : results) {
-      convertBatchWriteResult(result, pendingOperations.get(result.getDocumentReference()));
-    }
+    pendingOperations = newPendingOperations;
   }
 
   private boolean shouldRetry(Exception exception) {
@@ -800,15 +840,6 @@ private boolean shouldRetry(Exception exception) {
     return false;
   }
 
-  private void convertBatchWriteResult(
-      BatchWriteResult result, SettableApiFuture<WriteResult> future) {
-    if (result.getWriteTime() != null) {
-      future.set(new WriteResult(result.getWriteTime()));
-    } else {
-      future.setException(result.getException());
-    }
-  }
-
   void markReadyToSend() {
     if (state == BatchState.OPEN) {
       state = BatchState.READY_TO_SEND;
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/BulkWriterTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/BulkWriterTest.java
index e5f72af65..f51471009 100644
--- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/BulkWriterTest.java
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/BulkWriterTest.java
@@ -33,7 +33,6 @@
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.LocalFirestoreHelper.ResponseStubber;
-import com.google.cloud.firestore.LocalFirestoreHelper.SerialResponseStubber;
 import com.google.cloud.firestore.spi.v1.FirestoreRpc;
 import com.google.firestore.v1.BatchWriteRequest;
 import com.google.firestore.v1.BatchWriteResponse;
@@ -43,7 +42,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.concurrent.Callable;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
@@ -52,6 +50,7 @@
 import javax.annotation.Nonnull;
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.Timeout;
@@ -332,16 +331,15 @@ public void cannotCallMethodsAfterClose() throws Exception {
   }
 
   @Test
-  public void sendsWritesToSameDocInSeparateBatches() throws Exception {
+  public void canSendWritesToSameDocInSameBatch() throws Exception {
     ResponseStubber responseStubber =
         new ResponseStubber() {
           {
             put(
-                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
-                successResponse(1));
-            put(
-                batchWrite(update(LocalFirestoreHelper.UPDATED_FIELD_PROTO, "coll/doc1")),
-                successResponse(2));
+                batchWrite(
+                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
+                    update(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
+                mergeResponses(successResponse(1), successResponse(2)));
           }
         };
     responseStubber.initializeStub(batchWriteCapture, firestoreMock);
@@ -350,7 +348,7 @@ public void sendsWritesToSameDocInSeparateBatches() throws Exception {
     DocumentReference sameDoc = firestoreMock.document(doc1.getPath());
     ApiFuture<WriteResult> result1 = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
     ApiFuture<WriteResult> result2 =
-        bulkWriter.update(sameDoc, LocalFirestoreHelper.UPDATED_FIELD_MAP);
+        bulkWriter.update(sameDoc, LocalFirestoreHelper.SINGLE_FIELD_MAP);
     bulkWriter.close();
 
     List<BatchWriteRequest> requests = batchWriteCapture.getAllValues();
@@ -429,7 +427,7 @@ public void retriesIndividualWritesThatFailWithAbortedOrUnavailable() throws Exc
             put(
                 batchWrite(
                     set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
-                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2"),
+                    set(LocalFirestoreHelper.UPDATED_FIELD_PROTO, "coll/doc1"),
                     set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc3")),
                 mergeResponses(
                     failedResponse(),
@@ -437,7 +435,7 @@ public void retriesIndividualWritesThatFailWithAbortedOrUnavailable() throws Exc
                     failedResponse(Code.ABORTED_VALUE)));
             put(
                 batchWrite(
-                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2"),
+                    set(LocalFirestoreHelper.UPDATED_FIELD_PROTO, "coll/doc1"),
                     set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc3")),
                 mergeResponses(successResponse(2), failedResponse(Code.ABORTED_VALUE)));
             put(
@@ -447,8 +445,10 @@ public void retriesIndividualWritesThatFailWithAbortedOrUnavailable() throws Exc
         };
     responseStubber.initializeStub(batchWriteCapture, firestoreMock);
 
+    // Test writes to the same document in order to verify that retry logic unaffected by document
+    // key.
     ApiFuture<WriteResult> result1 = bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    ApiFuture<WriteResult> result2 = bulkWriter.set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    ApiFuture<WriteResult> result2 = bulkWriter.set(doc1, LocalFirestoreHelper.UPDATED_FIELD_MAP);
     ApiFuture<WriteResult> result3 =
         bulkWriter.set(firestoreMock.document("coll/doc3"), LocalFirestoreHelper.SINGLE_FIELD_MAP);
     bulkWriter.close();
@@ -549,57 +549,8 @@ public void flushCompletesWhenAllWritesComplete() throws Exception {
   }
 
   @Test
-  public void doesNotSendBatchesIfSameDocIsInFlight() throws Exception {
-    final SerialResponseStubber responseStubber =
-        new SerialResponseStubber() {
-          {
-            put(
-                batchWrite(
-                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
-                    set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2")),
-                mergeResponses(successResponse(1), successResponse(2)));
-            put(
-                batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1")),
-                successResponse(3));
-          }
-        };
-    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
-    bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    bulkWriter.set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-
-    // Schedule flush on separate thread to avoid blocking main thread while waiting for
-    // activeRequestComplete.
-    ScheduledFuture<ApiFuture<Void>> flush1 =
-        Executors.newSingleThreadScheduledExecutor()
-            .schedule(
-                new Callable<ApiFuture<Void>>() {
-                  public ApiFuture<Void> call() {
-                    return bulkWriter.flush();
-                  }
-                },
-                0,
-                TimeUnit.MILLISECONDS);
-
-    // Wait for flush() to perform logic and reach the stubbed response. This simulates a first
-    // batch that has been sent with its response still pending.
-    responseStubber.awaitRequest();
-
-    bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
-    ApiFuture<Void> flush2 = bulkWriter.flush();
-
-    // Wait for flush() to receive its response and process the batch.
-    responseStubber.markAllRequestsComplete();
-    flush1.get().get();
-    flush2.get();
-    bulkWriter.close();
-
-    List<BatchWriteRequest> requests = batchWriteCapture.getAllValues();
-    assertEquals(responseStubber.size(), requests.size());
-
-    verifyRequests(requests, responseStubber);
-  }
-
-  @Test
+  @Ignore
+  // TODO(chenbrian): Fix this test after throttling options are added.
   public void doesNotSendBatchesIfDoingSoExceedsRateLimit() {
     final boolean[] timeoutCalled = {false};
     final ScheduledExecutorService timeoutExecutor =