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

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,18 @@
     <className>com/google/cloud/firestore/spi/v1/FirestoreRpc</className>
     <method>com.google.api.gax.rpc.UnaryCallable batchWriteCallable()</method>
   </difference>
+  <difference>
+    <differenceType>7002</differenceType>
+    <className>com/google/cloud/firestore/BatchWriteResult</className>
+    <method>com.google.cloud.firestore.DocumentReference getDocumentReference()</method>
+  </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

google-cloud-firestore/src/main/java/com/google/cloud/firestore/BatchWriteResult.java

@@ -26,21 +26,18 @@
  */
 @InternalApi
 public final class BatchWriteResult {
-  private final DocumentReference documentReference;
+  private final String key;
   @Nullable private final Timestamp writeTime;
   @Nullable private final Exception exception;
 
-  BatchWriteResult(
-      DocumentReference documentReference,
-      @Nullable Timestamp timestamp,
-      @Nullable Exception exception) {
-    this.documentReference = documentReference;
+  BatchWriteResult(String key, @Nullable Timestamp timestamp, @Nullable Exception exception) {
+    this.key = key;
     this.writeTime = timestamp;
     this.exception = exception;
   }
 
-  public DocumentReference getDocumentReference() {
-    return documentReference;
+  public String getKey() {
+    return key;
   }
 
   @Nullable

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

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 (String documentPath : batch.getPendingDocumentPaths()) {
+                    results.add(new BatchWriteResult(documentPath, 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;