Skip to content
This repository was archived by the owner on May 8, 2026. It is now read-only.

Commit a7caff2

Browse files
author
Brian Chen
authored
feat: add buffering layer to BulkWriter (#611)
1 parent 0a2bc53 commit a7caff2

2 files changed

Lines changed: 144 additions & 14 deletions

File tree

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java‎

Lines changed: 107 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,13 @@
2525
import com.google.api.core.ApiFutures;
2626
import com.google.api.core.BetaApi;
2727
import com.google.api.core.SettableApiFuture;
28+
import com.google.api.gax.rpc.ApiException;
2829
import com.google.api.gax.rpc.StatusCode.Code;
2930
import com.google.cloud.firestore.v1.FirestoreSettings;
3031
import com.google.common.annotations.VisibleForTesting;
3132
import com.google.common.util.concurrent.MoreExecutors;
33+
import java.util.ArrayList;
34+
import java.util.List;
3235
import java.util.Map;
3336
import java.util.Set;
3437
import java.util.concurrent.ExecutionException;
@@ -110,6 +113,14 @@ enum OperationType {
110113
*/
111114
private static final int RATE_LIMITER_MULTIPLIER_MILLIS = 5 * 60 * 1000;
112115

116+
/**
117+
* The default maximum number of pending operations that can be enqueued onto a BulkWriter
118+
* instance. An operation is considered pending if BulkWriter has sent it via RPC and is awaiting
119+
* the result. BulkWriter buffers additional writes after this many pending operations in order to
120+
* avoiding going OOM.
121+
*/
122+
private static final int DEFAULT_MAXIMUM_PENDING_OPERATIONS_COUNT = 500;
123+
113124
/**
114125
* The default jitter to apply to the exponential backoff used in retries. For example, a factor
115126
* of 0.3 means a 30% jitter is applied.
@@ -158,6 +169,26 @@ public boolean onError(BulkWriterException error) {
158169
@GuardedBy("lock")
159170
private final RateLimiter rateLimiter;
160171

172+
/**
173+
* The number of pending operations enqueued on this BulkWriter instance. An operation is
174+
* considered pending if BulkWriter has sent it via RPC and is awaiting the result.
175+
*/
176+
@GuardedBy("lock")
177+
private int pendingOpsCount = 0;
178+
179+
/**
180+
* An array containing buffered BulkWriter operations after the maximum number of pending
181+
* operations has been enqueued.
182+
*/
183+
@GuardedBy("lock")
184+
private final List<Runnable> bufferedOperations = new ArrayList<>();
185+
186+
/**
187+
* The maximum number of pending operations that can be enqueued onto this BulkWriter instance.
188+
* Once the this number of writes have been enqueued, subsequent writes are buffered.
189+
*/
190+
private int maxPendingOpCount = DEFAULT_MAXIMUM_PENDING_OPERATIONS_COUNT;
191+
161192
/**
162193
* The batch that is currently used to schedule operations. Once this batch reaches maximum
163194
* capacity, a new batch is created.
@@ -627,7 +658,7 @@ private ApiFuture<WriteResult> executeWrite(
627658
final DocumentReference documentReference,
628659
final OperationType operationType,
629660
final ApiFunction<BulkCommitBatch, ApiFuture<WriteResult>> enqueueOperationOnBatchCallback) {
630-
BulkWriterOperation operation =
661+
final BulkWriterOperation operation =
631662
new BulkWriterOperation(
632663
documentReference,
633664
operationType,
@@ -660,10 +691,73 @@ public ApiFuture<Boolean> apply(BulkWriterException e) {
660691
synchronized (lock) {
661692
verifyNotClosedLocked();
662693
writesEnqueued = true;
663-
sendOperationLocked(enqueueOperationOnBatchCallback, operation);
694+
695+
// Advance the lastOperation pointer. This ensures that lastOperation only completes when
696+
// both the previous and the current write complete.
697+
lastOperation =
698+
ApiFutures.transformAsync(
699+
lastOperation,
700+
new ApiAsyncFunction<Void, Void>() {
701+
@Override
702+
public ApiFuture<Void> apply(Void aVoid) {
703+
return silenceFuture(operation.getFuture());
704+
}
705+
},
706+
MoreExecutors.directExecutor());
707+
708+
// Schedule the operation if the BulkWriter has fewer than the maximum number of allowed
709+
// pending operations, or add the operation to the buffer.
710+
if (pendingOpsCount < maxPendingOpCount) {
711+
pendingOpsCount++;
712+
sendOperationLocked(enqueueOperationOnBatchCallback, operation);
713+
} else {
714+
bufferedOperations.add(
715+
new Runnable() {
716+
@Override
717+
public void run() {
718+
synchronized (lock) {
719+
pendingOpsCount++;
720+
sendOperationLocked(enqueueOperationOnBatchCallback, operation);
721+
}
722+
}
723+
});
724+
}
664725
}
665726

666-
return operation.getFuture();
727+
ApiFuture<WriteResult> processedOperationFuture =
728+
ApiFutures.transformAsync(
729+
operation.getFuture(),
730+
new ApiAsyncFunction<WriteResult, WriteResult>() {
731+
public ApiFuture<WriteResult> apply(WriteResult result) throws Exception {
732+
pendingOpsCount--;
733+
processBufferedOperations();
734+
return ApiFutures.immediateFuture(result);
735+
}
736+
},
737+
MoreExecutors.directExecutor());
738+
739+
return ApiFutures.catchingAsync(
740+
processedOperationFuture,
741+
ApiException.class,
742+
new ApiAsyncFunction<ApiException, WriteResult>() {
743+
public ApiFuture<WriteResult> apply(ApiException e) throws Exception {
744+
pendingOpsCount--;
745+
processBufferedOperations();
746+
throw e;
747+
}
748+
},
749+
MoreExecutors.directExecutor());
750+
}
751+
752+
/**
753+
* Manages the pending operation counter and schedules the next BulkWriter operation if we're
754+
* under the maximum limit.
755+
*/
756+
private void processBufferedOperations() {
757+
if (pendingOpsCount < maxPendingOpCount && bufferedOperations.size() > 0) {
758+
Runnable nextOp = bufferedOperations.remove(0);
759+
nextOp.run();
760+
}
667761
}
668762

669763
/**
@@ -927,6 +1021,16 @@ RateLimiter getRateLimiter() {
9271021
return rateLimiter;
9281022
}
9291023

1024+
@VisibleForTesting
1025+
int getBufferedOperationsCount() {
1026+
return bufferedOperations.size();
1027+
}
1028+
1029+
@VisibleForTesting
1030+
void setMaxPendingOpCount(int newMax) {
1031+
maxPendingOpCount = newMax;
1032+
}
1033+
9301034
/**
9311035
* Schedules the provided operations on the current BulkCommitBatch. Sends the BulkCommitBatch if
9321036
* it reaches maximum capacity.
@@ -946,17 +1050,6 @@ private void sendOperationLocked(
9461050
bulkCommitBatch.enqueueOperation(op);
9471051
enqueueOperationOnBatchCallback.apply(bulkCommitBatch);
9481052

949-
lastOperation =
950-
ApiFutures.transformAsync(
951-
lastOperation,
952-
new ApiAsyncFunction<Void, Void>() {
953-
@Override
954-
public ApiFuture<Void> apply(Void aVoid) {
955-
return silenceFuture(op.getFuture());
956-
}
957-
},
958-
MoreExecutors.directExecutor());
959-
9601053
if (bulkCommitBatch.getMutationsSize() == maxBatchSize) {
9611054
scheduleCurrentBatchLocked(/* flush= */ false);
9621055
}

‎google-cloud-firestore/src/test/java/com/google/cloud/firestore/BulkWriterTest.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,41 @@ public void sendWritesToDifferentDocsInSameBatch() throws Exception {
422422
assertEquals(Timestamp.ofTimeSecondsAndNanos(2, 0), result2.get().getUpdateTime());
423423
}
424424

425+
@Test
426+
public void buffersSubsequentOpsAfterReachingMaxPendingOpCount() throws Exception {
427+
ResponseStubber responseStubber =
428+
new ResponseStubber() {
429+
{
430+
put(
431+
batchWrite(
432+
set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc1"),
433+
set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc2"),
434+
set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc3")),
435+
mergeResponses(
436+
successResponse(1),
437+
successResponse(2),
438+
failedResponse(Code.FAILED_PRECONDITION_VALUE)));
439+
put(
440+
batchWrite(
441+
set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc4"),
442+
set(LocalFirestoreHelper.SINGLE_FIELD_PROTO, "coll/doc5")),
443+
mergeResponses(successResponse(4), successResponse(5)));
444+
}
445+
};
446+
responseStubber.initializeStub(batchWriteCapture, firestoreMock);
447+
448+
bulkWriter.setMaxPendingOpCount(3);
449+
bulkWriter.set(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
450+
bulkWriter.set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP);
451+
bulkWriter.set(firestoreMock.document("coll/doc3"), LocalFirestoreHelper.SINGLE_FIELD_MAP);
452+
bulkWriter.set(firestoreMock.document("coll/doc4"), LocalFirestoreHelper.SINGLE_FIELD_MAP);
453+
assertEquals(1, bulkWriter.getBufferedOperationsCount());
454+
bulkWriter.set(firestoreMock.document("coll/doc5"), LocalFirestoreHelper.SINGLE_FIELD_MAP);
455+
assertEquals(2, bulkWriter.getBufferedOperationsCount());
456+
bulkWriter.close();
457+
responseStubber.verifyAllRequestsSent();
458+
}
459+
425460
@Test
426461
public void runsSuccessHandler() throws Exception {
427462
ResponseStubber responseStubber =
@@ -1260,6 +1295,7 @@ public boolean onError(BulkWriterException error) {
12601295
bulkWriter.create(doc1, LocalFirestoreHelper.SINGLE_FIELD_MAP);
12611296
bulkWriter.set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP);
12621297
bulkWriter.close();
1298+
responseStubber.verifyAllRequestsSent();
12631299
assertEquals(2, retryAttempts[0]);
12641300
}
12651301

@@ -1291,6 +1327,7 @@ public boolean onError(BulkWriterException error) {
12911327
bulkWriter.flush();
12921328
bulkWriter.set(doc2, LocalFirestoreHelper.SINGLE_FIELD_MAP);
12931329
bulkWriter.close();
1330+
responseStubber.verifyAllRequestsSent();
12941331
}
12951332

12961333
@Test

0 commit comments

Comments
 (0)