Bulk enhancement with dynamic congestion control #32567

aditishree1 · 2025-01-15T19:17:36Z

Packages impacted by this PR

@azure/cosmos

Issues associated with this PR

Describe the problem that is addressed by this PR

This PR enhances the flow of Bulk API. It implements custom retry policy for bulk and dynamically controls degree of concurrency for each partition according to throttling and no. of items processed.

What are the possible designs available to address the problem? If there are more than one possible design, why was the one in this PR chosen?

https://excalidraw.com/#room=ac45daa258e67c6fd89d,jkOfy1GkGAKwRUYZkTGgog

Are there test cases added in this PR? (If not, why?)

yes

Provide a list of related PRs (if any)

Command used to generate this PR:***(Applicable only to SDK release request PRs)*

Checklists

Added impacted package name to the issue description
Does this PR needs any fixes in the SDK Generator?** (If so, create an Issue in the Autorest/typescript repository and link it here)
Added a changelog (if necessary)

azure-sdk · 2025-01-15T19:26:59Z

API change check

APIView has identified API level changes in this PR and created following API reviews.

@azure/cosmos

sdk/cosmosdb/cosmos/review/cosmos.api.md

sdk/cosmosdb/cosmos/src/bulk/BulkStreamerCache.ts

sdk/cosmosdb/cosmos/src/client/Item/Items.ts

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts

topshot99 · 2025-01-27T04:57:15Z

sdk/cosmosdb/cosmos/review/cosmos.api.md

@@ -558,6 +582,9 @@ export const Constants: {
    SDKVersion: string;
    CosmosDbDiagnosticLevelEnvVarName: string;
    DefaultMaxBulkRequestBodySizeInBytes: number;
+    MaxBulkOperationsCount: number;


any reason for adding bulk in the names? Can't it be just Operations count?

not any specific reason. we are using these constants for bulk, that's why. Should we change these? In .Net, these are named as MaxOperationsInDirectModeBatchRequest and DefaultMaxBulkRequestBodySizeInBytes

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts

topshot99 · 2025-01-27T12:01:22Z

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts

+    } finally {
+      for (const [key, streamer] of this.streamersByPartitionKeyRangeId.entries()) {
+        streamer.disposeTimers();
+        this.limitersByPartitionKeyRangeId.delete(key);


can't we just delete the entire map in one go?

topshot99 · 2025-01-28T10:47:09Z

sdk/cosmosdb/cosmos/src/bulk/Limiter.ts

+// Licensed under the MIT License.
+import semaphore from "semaphore";
+/**
+ * Semaphores and locks for execution of Bulk


can you add a better description? Can it be used for other flows too like changefeed, query?

The methods (isStopped and stopDispatch) are specific for bulk only. They are used to block requests in case of 410. If there are similar requirements then this could be used.
I'll add more description.

topshot99 · 2025-01-28T12:28:50Z

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts

+
+  private async reBatchOperation(operation: ItemBulkOperation): Promise<void> {
+    const partitionKeyRangeId = await this.resolvePartitionKeyRangeId(operation.operationInput);
+    operation.operationContext.reRouteOperation(partitionKeyRangeId);


instead of reRouteOperation can it be just update partitionKeyRangeID?

topshot99 · 2025-01-28T12:30:09Z

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts

+    return limiter;
+  }
+
+  private getOrCreateStreamerForPKRange(pkRangeId: string): BulkStreamerPerPartition {


rename to getStreamerForPKRange

topshot99 · 2025-01-29T04:00:11Z

sdk/cosmosdb/cosmos/src/bulk/BulkStreamerPerPartition.ts

+   * @param operation - operation to add
+   */
+  add(operation: ItemBulkOperation): void {
+    let toDispatch: BulkBatcher | null = null;


why is this initialized with null

we could keep it as undefined also. I'll update.

topshot99 · 2025-01-29T04:26:35Z

sdk/cosmosdb/cosmos/src/bulk/BulkBatcher.ts

+        this.options,
+        this.diagnosticNode,
+      );
+      if (response.statusCode === 0) {


add a log about significance of 0 status code

topshot99 · 2025-01-29T04:40:47Z

sdk/cosmosdb/cosmos/src/bulk/BulkStreamerPerPartition.ts

+    this.lock.take(() => {
+      try {
+        // attempt to add operation until it fits in the current batch for the streamer
+        while (!this.currentBatcher.tryAdd(operation)) {


If there's a backend issue where the backend keeps timing out and, after some retries, returns an error, the dispatcher might get stuck. In such cases, does it really make sense to keep a thread busy while adding the batcher? Wouldn't it be better to make this method async and wait for tryAdd to complete? Maybe the batcher can handle the issue internally.

This lock is only for adding of operation in the batch. The tryAdd method will try to add operation in current batch. This will only return false when the batch is full. In that case, the full batch will be assigned to toDispatch and a new batch will get created. So, when this loop will run again, that operation will get added to the new batch.

topshot99 · 2025-01-29T04:59:57Z

sdk/cosmosdb/cosmos/src/bulk/BulkStreamerPerPartition.ts

+
+    if (toDispatch) {
+      // dispatch with fire and forget. No need to wait for the dispatch to complete.
+      toDispatch.dispatch(this.partitionMetric);


We are calling dispatch with every add operation. I thought it's controlled by timer.

We have both the mechanisms for dispatching the batch. The add operation checks the operation count and batch size condition. Based on this, if the batch is full, it will dispatch it. Timer mechanism ensures timely dispatch to prevent the batch from waiting for too long when there aren't enough operations to fill it.

topshot99 · 2025-01-29T05:30:29Z

sdk/cosmosdb/cosmos/src/bulk/BulkCongestionAlgorithm.ts

+      Constants.BulkMaxDegreeOfConcurrency
+    ) {
+      if (this.limiter.current() > 0) {
+        this.limiter.leave(this.congestionIncreaseFactor);


help me understand this? how updating limiter helps in increasing concurrency.

the semaphore is initialized with max 50 slots. We start with degree of concurrency as 1 and block 49 slots. According to the algo, we take or leave slots. Suppose we have current degree of concurrency as 1 and we didn't get any throttles, the algo will run and we will release one more slot to increase concurrency. So, now max of 2 requests can be sent parallely.

topshot99 · 2025-01-29T05:33:49Z

sdk/cosmosdb/cosmos/src/bulk/ItemBulkOperation.ts

+ * @hidden
+ */
+
+export class ItemBulkOperation {


can it be an interface?

topshot99 · 2025-01-29T05:39:23Z

sdk/cosmosdb/cosmos/src/bulk/BulkOperationResult.ts

+/**
+ * Represents a result for a specific operation that was part of a batch request
+ */
+export class BulkOperationResult {


convert it to an iterface

topshot99 · 2025-01-29T05:53:18Z

sdk/cosmosdb/cosmos/src/retry/bulkExecutionRetryPolicy.ts

+    // API can return 413 which means the response is bigger than 4Mb.
+    // Operations that exceed the 4Mb limit are returned as 413/3402, while the operations within the 4Mb limit will be 200
+    if (
+      err.code === StatusCodes.RequestEntityTooLarge &&


we are already validating the batch size, do we really need this retry?

we get this status code and substatus code when the response size from backend exceeds 4 Mb. I was not able to replicate this error to see the exact behaviour from backend. But I included it since it was present in the .NET SDK's bulk retry policy.

There is a possibility that addition of headers might exceed the size limit.

topshot99 · 2025-01-29T05:58:52Z

sdk/cosmosdb/cosmos/src/retry/bulkExecutionRetryPolicy.ts

+      shouldRetryForThrottle = Array.isArray(retryResult) ? retryResult[0] : retryResult;
+    }
+    if (shouldRetryForThrottle) {
+      await sleep(this.nextRetryPolicy.retryAfterInMs);


isn't it already covered:

azure-sdk-for-js/sdk/cosmosdb/cosmos/src/retry/retryUtility.ts

Line 171 in 0c3aaff

await sleep(retryPolicy.retryAfterInMs);

Bulk retry policy will apply on top of current retry policies since we are applying retry on operation level. So, it will not take this path.

topshot99

check comments

aditishree1 · 2025-01-29T06:59:33Z

check comments

addressed comments

topshot99 · 2025-01-29T10:26:56Z

sdk/cosmosdb/cosmos/test/public/functional/item/bulkStreamer.item.spec.ts

+      operations.forEach((operation) => bulkStreamer.addOperations(operation));
+      const response = await bulkStreamer.endStream();
+      // Create
+      response.forEach((res, index) =>


add a check for number of responses

topshot99 · 2025-01-29T10:28:53Z

sdk/cosmosdb/cosmos/test/public/functional/item/bulkStreamer.item.spec.ts

+        assert.strictEqual(res.statusCode, 201, `Status should be 201 for operation ${index}`),
+      );
+    });
+    it("Check case when cumulative size of all operations is greater than threshold - payload size is 25x threshold", async function () {


how about one operation having size greater than threshold?

In that case, it will send the request anyway and throw the 413 error. I'll add a test.

topshot99 · 2025-01-29T10:32:52Z

sdk/cosmosdb/cosmos/test/public/functional/item/bulkStreamer.item.spec.ts

+        assert.equal(response[4].resourceBody.name, "nice");
+        assert.equal(response[4].statusCode, 200);
+      });
+      it("Check case when cumulative size of all operations is less than threshold", async function () {


Only operation input is changing in the test, can we combine these tests together?

topshot99 · 2025-01-29T10:47:40Z

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts

+   *
+   * @internal
+   */
+  // TODO: Add a public reset() method which will utilize this to reset the state of the streamer to use it again.


is it future TODO?

topshot99 · 2025-01-29T10:59:10Z

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts

+   */
+  async endStream(): Promise<BulkStreamerResponse> {
+    if (this.streamEnded) {
+      throw new ErrorResponse("Bulk streamer has already ended.");


can we attach some status code too? some 4xx code conveying bad request

that will be http like status code. should we set that for throwing sdk error?

topshot99 · 2025-01-29T11:01:12Z

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts

+    if (this.streamEnded) {
+      throw new ErrorResponse("Bulk streamer has already ended.");
+    }
+    this.streamEnded = true;


add it after line 136

Wouldn't this get stuck if user will keep on adding operation after calling endStream while promise are settling? We were trying to avoid this case with streamEnded flag.

github-actions bot added the Cosmos label Jan 15, 2025

aditishree1 mentioned this pull request Jan 16, 2025

Bulk enhancement with congestion control #32434

Closed

3 tasks

aditishree1 marked this pull request as ready for review January 21, 2025 03:08

aditishree1 requested review from amanrao23, sajeetharan and topshot99 as code owners January 21, 2025 03:08

amanrao23 reviewed Jan 22, 2025

View reviewed changes

sdk/cosmosdb/cosmos/review/cosmos.api.md Outdated Show resolved Hide resolved

sdk/cosmosdb/cosmos/src/bulk/BulkStreamerCache.ts Outdated Show resolved Hide resolved

topshot99 reviewed Jan 24, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/client/Item/Items.ts Outdated Show resolved Hide resolved

amanrao23 reviewed Jan 27, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts Outdated Show resolved Hide resolved

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts Outdated Show resolved Hide resolved

topshot99 reviewed Jan 27, 2025

View reviewed changes

amanrao23 reviewed Jan 27, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/bulk/BulkStreamer.ts Outdated Show resolved Hide resolved

Aditishree . and others added 18 commits January 27, 2025 10:49

modified bulk implementation

bfbfa4b

remove logs

5b0432c

add bulkExecutionRetryPolicy and update response

19a940a

remove errors

f3cfd89

rush format

a8d50fb

add congestion control logic

cc88b23

remove internal tag

9984c3c

format

3cfb4da

remove executor cache usage

12f231c

delete resources from executor and add checks for semaphore

c4aca61

modify public methods

eb9b281

format

2b1a23d

add option to pass a list of operations

9a31b46

format

ffc2a07

fix container leaking in tests

831336d

add bulk request handler and fix operation Index

68bc101

fix retry error and degree of concurrency

0206c2e

Add limiter to stop batches on 410

42a9d5d

topshot99 reviewed Jan 27, 2025

View reviewed changes

remove separate operation index, add checks and test

7f82e21

topshot99 reviewed Jan 28, 2025

View reviewed changes

refactor names and add comments

0c3aaff

topshot99 reviewed Jan 29, 2025

View reviewed changes

address comments

6e68420

topshot99 reviewed Jan 29, 2025

View reviewed changes

Aditishree . added 6 commits January 29, 2025 18:29

add lock on operationIndex

63ebbbd

add comment for TODO

8c8eb80

remove redundant tests

a25b5e8

Merge branch 'main' into bulk-congestion-control

9d89d0c

format

f415870

extend operation response and add lock to degree of concurrency

028d437

Bulk enhancement with dynamic congestion control #32567

Are you sure you want to change the base?

Bulk enhancement with dynamic congestion control #32567

Conversation

aditishree1 commented Jan 15, 2025

Packages impacted by this PR

Issues associated with this PR

Describe the problem that is addressed by this PR

What are the possible designs available to address the problem? If there are more than one possible design, why was the one in this PR chosen?

Are there test cases added in this PR? (If not, why?)

Provide a list of related PRs (if any)

Command used to generate this PR:**(Applicable only to SDK release request PRs)

Checklists

azure-sdk commented Jan 15, 2025

topshot99 Jan 27, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

topshot99 Jan 28, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

topshot99 left a comment

Choose a reason for hiding this comment

aditishree1 commented Jan 29, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

aditishree1 Jan 29, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Command used to generate this PR:***(Applicable only to SDK release request PRs)*

topshot99 Jan 27, 2025 •

edited

Loading

topshot99 Jan 28, 2025 •

edited

Loading

aditishree1 Jan 29, 2025 •

edited

Loading