fix: BulkWriter flush logic (#1778)

tom-andersen · web-flow · commit 627619d33a05 · 2024-12-04T09:54:15.000-05:00
* Fix race condition

* Pretty

* Optimize
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
@@ -210,6 +210,10 @@ enum OperationType {
   @GuardedBy("lock")
   private ApiFuture<Void> lastOperation = ApiFutures.immediateFuture(null);
 
+  /** A pointer to the lastOperation pointer as of last flush operation. */
+  @GuardedBy("lock")
+  private ApiFuture<Void> lastFlushOperation = lastOperation;
+
   /** Whether this BulkWriter instance is closed. Once closed, it cannot be opened again. */
   @GuardedBy("lock")
   private boolean closed = false;
@@ -699,7 +703,10 @@ public ApiFuture<Void> flush() {
 
   private ApiFuture<Void> flushLocked() {
     verifyNotClosedLocked();
-    scheduleCurrentBatchLocked(/* flush= */ true);
+    if (!lastOperation.isDone()) {
+      lastFlushOperation = lastOperation;
+      scheduleCurrentBatchLocked();
+    }
     return lastOperation;
   }
 
@@ -873,14 +880,8 @@ public void addWriteErrorListener(@Nonnull Executor executor, WriteErrorCallback
     }
   }
 
-  /**
-   * Sends the current batch and resets {@link #bulkCommitBatch}.
-   *
-   * @param flush If provided, keeps re-sending operations until no more operations are enqueued.
-   *     This allows retries to resolve as part of a {@link BulkWriter#flush()} or {@link
-   *     BulkWriter#close()} call.
-   */
-  private void scheduleCurrentBatchLocked(final boolean flush) {
+  /** Sends the current batch and resets {@link #bulkCommitBatch}. */
+  private void scheduleCurrentBatchLocked() {
     if (bulkCommitBatch.getMutationsSize() == 0) return;
 
     final BulkCommitBatch pendingBatch = bulkCommitBatch;
@@ -898,15 +899,15 @@ private void scheduleCurrentBatchLocked(final boolean flush) {
     bulkWriterExecutor.schedule(
         () -> {
           synchronized (lock) {
-            sendBatchLocked(pendingBatch, flush);
+            sendBatchLocked(pendingBatch);
           }
         },
         backoffMsWithJitter,
         TimeUnit.MILLISECONDS);
   }
 
   /** Sends the provided batch once the rate limiter does not require any delay. */
-  private void sendBatchLocked(final BulkCommitBatch batch, final boolean flush) {
+  private void sendBatchLocked(final BulkCommitBatch batch) {
     // Send the batch if it does not require any delay, or schedule another attempt after the
     // appropriate timeout.
     boolean underRateLimit = rateLimiter.tryMakeRequest(batch.getMutationsSize());
@@ -926,15 +927,17 @@ private void sendBatchLocked(final BulkCommitBatch batch, final boolean flush) {
 
       try (Scope ignored = span.makeCurrent()) {
         ApiFuture<Void> result = batch.bulkCommit();
-        result.addListener(
-            () -> {
-              if (flush) {
-                synchronized (lock) {
-                  scheduleCurrentBatchLocked(/* flush= */ true);
+        if (!lastFlushOperation.isDone()) {
+          result.addListener(
+              () -> {
+                if (!lastFlushOperation.isDone()) {
+                  synchronized (lock) {
+                    scheduleCurrentBatchLocked();
+                  }
                 }
-              }
-            },
-            bulkWriterExecutor);
+              },
+              bulkWriterExecutor);
+        }
         span.endAtFuture(result);
         metricsContext.recordLatencyAtFuture(MetricType.END_TO_END_LATENCY, result);
       } catch (Exception error) {
@@ -948,7 +951,7 @@ private void sendBatchLocked(final BulkCommitBatch batch, final boolean flush) {
       bulkWriterExecutor.schedule(
           () -> {
             synchronized (lock) {
-              sendBatchLocked(batch, flush);
+              sendBatchLocked(batch);
             }
           },
           delayMs,
@@ -991,15 +994,15 @@ private void sendOperationLocked(
     // that the batch is under the 10MiB limit.
     if (op.getBackoffDuration() > 0) {
       if (bulkCommitBatch.getMutationsSize() >= RETRY_MAX_BATCH_SIZE) {
-        scheduleCurrentBatchLocked(/* flush= */ false);
+        scheduleCurrentBatchLocked();
       }
       bulkCommitBatch.setMaxBatchSize(RETRY_MAX_BATCH_SIZE);
     }
 
     if (bulkCommitBatch.has(op.getDocumentReference())) {
       // Create a new batch since the backend doesn't support batches with two writes to the same
       // document.
-      scheduleCurrentBatchLocked(/* flush= */ false);
+      scheduleCurrentBatchLocked();
     }
 
     // Run the operation on the current batch and advance the `lastOperation` pointer. This
@@ -1008,8 +1011,8 @@ private void sendOperationLocked(
     bulkCommitBatch.enqueueOperation(op);
     enqueueOperationOnBatchCallback.apply(bulkCommitBatch);
 
-    if (bulkCommitBatch.getMutationsSize() == bulkCommitBatch.getMaxBatchSize()) {
-      scheduleCurrentBatchLocked(/* flush= */ false);
+    if (bulkCommitBatch.getMutationsSize() >= bulkCommitBatch.getMaxBatchSize()) {
+      scheduleCurrentBatchLocked();
     }
   }
 
