fix: Resolve RejectedExecutionException that occurred on BulkWriter close or executor shutdown (#2223)

MarkDuckworth · web-flow · commit e919ebbef229 · 2025-08-21T12:18:11.000-06:00
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
@@ -661,8 +661,10 @@ private ApiFuture<WriteResult> executeWrite(
         ApiFutures.transformAsync(
             operation.getFuture(),
             result -> {
-              pendingOpsCount--;
-              processBufferedOperations();
+              synchronized (lock) {
+                pendingOpsCount--;
+                processBufferedOperations();
+              }
               return ApiFutures.immediateFuture(result);
             },
             MoreExecutors.directExecutor());
@@ -671,8 +673,10 @@ private ApiFuture<WriteResult> executeWrite(
         processedOperationFuture,
         ApiException.class,
         e -> {
-          pendingOpsCount--;
-          processBufferedOperations();
+          synchronized (lock) {
+            pendingOpsCount--;
+            processBufferedOperations();
+          }
           throw e;
         },
         MoreExecutors.directExecutor());
@@ -951,7 +955,7 @@ private void sendBatchLocked(final BulkCommitBatch batch) {
                   }
                 }
               },
-              bulkWriterExecutor);
+              MoreExecutors.directExecutor());
         }
         span.endAtFuture(result);
         metricsContext.recordLatencyAtFuture(MetricType.END_TO_END_LATENCY, result);
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
@@ -158,7 +158,7 @@ public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit)
 
     bulkWriter =
         firestoreMock.bulkWriter(BulkWriterOptions.builder().setExecutor(timeoutExecutor).build());
-    bulkWriter.autoShutdownBulkWriterExecutor = true;
+
     doc1 = firestoreMock.document("coll/doc1");
     doc2 = firestoreMock.document("coll/doc2");
   }
@@ -1401,4 +1401,30 @@ public void optionsInitialAndMaxRatesAreProperlySet() throws Exception {
     assertEquals(bulkWriter.getRateLimiter().getInitialCapacity(), Integer.MAX_VALUE);
     assertEquals(bulkWriter.getRateLimiter().getMaximumRate(), Integer.MAX_VALUE);
   }
+
+  @Test
+  public void closeHandlesLargeNumberOfBufferedOps() throws Exception {
+    final int numOps = 100;
+
+    bulkWriter.setMaxPendingOpCount(5);
+    bulkWriter.setMaxBatchSize(1);
+    bulkWriter.autoShutdownBulkWriterExecutor = true;
+
+    ResponseStubber responseStubber = new ResponseStubber();
+
+    for (int i = 0; i < numOps; i += 1) {
+      responseStubber.put(
+          batchWrite(set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc" + i)),
+          successResponse(1));
+    }
+
+    responseStubber.initializeStub(batchWriteCapture, firestoreMock);
+
+    for (int i = 0; i < numOps; ++i) {
+      bulkWriter.set(firestoreMock.document("coll/doc" + i), LocalFirestoreHelper.SINGLE_FIELD_MAP);
+    }
+    bulkWriter.close();
+    responseStubber.verifyAllRequestsSent();
+    assertEquals(numOps, responseStubber.actualRequestList.size());
+  }
 }
