Repository: Azure/azure-sdk-for-js
Pull request #33869: [@azure/cosmos] fix: changefeed merge
Author: login: amanrao23
Labels: name: Cosmos
Created: 2025-04-16T09:50:51Z | Merged: 2025-05-12T08:52:51Z

### Packages impacted by this PR
@azure/cosmos

### Issues associated with this PR


### Describe the problem that is addressed by this PR
This PR addresses the bug of changefeed for entire container returning wrong results after partition merge. 

### 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
- [ ] 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](https://github.com/Azure/autorest.typescript) repository and link it here)_
- [ ] Added a changelog (if necessary)



Diff:
diff --git a/sdk/cosmosdb/cosmos/src/client/ChangeFeed/ChangeFeedForEpkRange.ts b/sdk/cosmosdb/cosmos/src/client/ChangeFeed/ChangeFeedForEpkRange.ts
index a7dc76d1d590..2952edbf54e1 100644
--- a/sdk/cosmosdb/cosmos/src/client/ChangeFeed/ChangeFeedForEpkRange.ts
+++ b/sdk/cosmosdb/cosmos/src/client/ChangeFeed/ChangeFeedForEpkRange.ts
@@ -310,7 +310,12 @@ export class ChangeFeedForEpkRange<T> implements ChangeFeedPullModelIterator<T>
       // resolvedRanges.length > 1 in case of split.
       // resolvedRanges.length === 1 in case of merge. EpkRange headers will be added in this case.
       if (resolvedRanges.length >= 1) {
-        await this.handleSplit(false, resolvedRanges, queryRange, feedRange.continuationToken);
+        await this.handleSplitOrMerge(
+          false,
+          resolvedRanges,
+          queryRange,
+          feedRange.continuationToken,
+        );
       }
       return true;
     }
@@ -319,7 +324,7 @@ export class ChangeFeedForEpkRange<T> implements ChangeFeedPullModelIterator<T>
   /*
    * Enqueues all the children feed ranges for the given feed range.
    */
-  private async handleSplit(
+  private async handleSplitOrMerge(
     shiftLeft: boolean,
     resolvedRanges: any,
     oldFeedRange: QueryRange,
@@ -327,7 +332,7 @@ export class ChangeFeedForEpkRange<T> implements ChangeFeedPullModelIterator<T>
   ): Promise<void> {
     let flag = 0;
     if (shiftLeft) {
-      // This section is only applicable when handleSplit is called by getPartitionRangeId().
+      // This section is only applicable when handleSplitOrMerge is called by getPartitionRangeId().
       // used only when existing partition key range cache is used to check for any overlapping ranges.
       // Modifies the first element with the first overlapping range.
       const [epkMinHeader, epkMaxHeader] = await extractOverlappingRanges(
@@ -384,8 +389,13 @@ export class ChangeFeedForEpkRange<T> implements ChangeFeedPullModelIterator<T>
       throw new ErrorResponse("No overlapping ranges found.");
     }
     const firstResolvedRange = resolvedRanges[0];
-    if (resolvedRanges.length > 1) {
-      await this.handleSplit(true, resolvedRanges, queryRange, feedRange.continuationToken);
+    const isPartitionRangeChanged =
+      feedRange.minInclusive !== firstResolvedRange.minInclusive ||
+      feedRange.maxExclusive !== firstResolvedRange.maxExclusive ||
+      resolvedRanges.length > 1;
+    // If the partition range is changed, we need to handle split/merge
+    if (isPartitionRangeChanged) {
+      await this.handleSplitOrMerge(true, resolvedRanges, queryRange, feedRange.continuationToken);
     }
     return firstResolvedRange.id;
   }
@@ -437,6 +447,7 @@ export class ChangeFeedForEpkRange<T> implements ChangeFeedPullModelIterator<T>
     }
     try {
       // startEpk and endEpk are only valid in case we want to fetch result for a part of partition and not the entire partition.
+      const finalFeedRange = this.fetchFinalFeedRange();
       const response: Response<Array<T & Resource>> = await (this.clientContext.queryFeed<T>({
         path: this.resourceLink,
         resourceType: ResourceType.item,
@@ -447,8 +458,8 @@ export class ChangeFeedForEpkRange<T> implements ChangeFeedPullModelIterator<T>
         diagnosticNode,
         partitionKey: undefined,
         partitionKeyRangeId: rangeId,
-        startEpk: feedRange.epkMinHeader,
-        endEpk: feedRange.epkMaxHeader,
+        startEpk: finalFeedRange.epkMinHeader,
+        endEpk: finalFeedRange.epkMaxHeader,
       }) as Promise<any>);
 
       return new ChangeFeedIteratorResponse(
@@ -477,4 +488,14 @@ export class ChangeFeedForEpkRange<T> implements ChangeFeedPullModelIterator<T>
       throw errorResponse;
     }
   }
+  private fetchFinalFeedRange(): ChangeFeedRange {
+    // this is used to fetch the final feed range before making a call to fetch the results.
+    // In case of merge, the final updated feed range is present in the queue and needs to be returned.
+    const feedRange = this.queue.peek();
+    if (feedRange) {
+      return feedRange;
+    } else {
+      throw new ErrorResponse("No feed range found.");
+    }
+  }
 }
diff --git a/sdk/cosmosdb/cosmos/test/internal/unit/changeFeed/ChangeFeedForEpkRange.spec.ts b/sdk/cosmosdb/cosmos/test/internal/unit/changeFeed/ChangeFeedForEpkRange.spec.ts
index c191e05b9bdb..9a149a3549e2 100644
--- a/sdk/cosmosdb/cosmos/test/internal/unit/changeFeed/ChangeFeedForEpkRange.spec.ts
+++ b/sdk/cosmosdb/cosmos/test/internal/unit/changeFeed/ChangeFeedForEpkRange.spec.ts
@@ -1,13 +1,26 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { ChangeFeedForEpkRange, ChangeFeedMode } from "../../../../src/client/ChangeFeed/index.js";
+import {
+  ChangeFeedForEpkRange,
+  ChangeFeedIteratorResponse,
+  ChangeFeedMode,
+} from "../../../../src/client/ChangeFeed/index.js";
 import type { ClientContext } from "../../../../src/index.js";
-import { Container, ErrorResponse, TimeoutError } from "../../../../src/index.js";
+import {
+  Constants,
+  Container,
+  ErrorResponse,
+  StatusCodes,
+  TimeoutError,
+} from "../../../../src/index.js";
 import { PartitionKeyRangeCache, QueryRange } from "../../../../src/routing/index.js";
 import { ChangeFeedRange } from "../../../../src/client/ChangeFeed/ChangeFeedRange.js";
 import { MockedClientContext } from "../../../public/common/MockClientContext.js";
 import { describe, it, assert, expect, beforeEach, afterEach, vi } from "vitest";
+import { getEmptyCosmosDiagnostics } from "../../../../src/utils/diagnostics.js";
+import { SubStatusCodes } from "../../../../src/common/statusCodes.js";
+import { FeedRangeQueue } from "../../../../src/client/ChangeFeed/FeedRangeQueue.js";
 
 interface Resource {
   id: string;
@@ -94,4 +107,213 @@ describe("ChangeFeedForEpkRange Unit Tests", () => {
       assert.strictEqual(err.code, timeOutError.code);
     }
   });
+
+  it("should handle partition split into three ranges and retry accordingly", async () => {
+    // first call: return 410/Gone to trigger shouldRetryOnFailure
+    const queue: FeedRangeQueue<ChangeFeedRange> = new FeedRangeQueue();
+    queue.enqueue(new ChangeFeedRange("", "DD", "cont1"));
+    queue.enqueue(new ChangeFeedRange("DD", "FF", "cont2"));
+    (changeFeedForEpkRange as any).queue = queue;
+
+    const goneResp = new ChangeFeedIteratorResponse(
+      [],
+      0,
+      StatusCodes.Gone,
+      { [Constants.HttpHeaders.ETag]: "" },
+      getEmptyCosmosDiagnostics(),
+      SubStatusCodes.PartitionKeyRangeGone,
+    );
+    // second call: return OK
+    const okResp = new ChangeFeedIteratorResponse(
+      [{ id: "ok" }],
+      1,
+      StatusCodes.Ok,
+      {},
+      getEmptyCosmosDiagnostics(),
+    );
+    const feedSpy = vi
+      .spyOn(changeFeedForEpkRange as any, "getFeedResponse")
+      .mockResolvedValueOnce(goneResp)
+      .mockResolvedValueOnce(okResp);
+
+    // simulate three child ranges when shouldRetryOnFailure calls getOverlappingRanges
+    vi.spyOn(partitionKeyRangeCache, "getOverlappingRanges").mockResolvedValue([
+      {
+        id: "1",
+        minInclusive: "",
+        maxExclusive: "AA",
+        ridPrefix: 0,
+        throughputFraction: 0,
+        status: "",
+        parents: [],
+      },
+      {
+        id: "2",
+        minInclusive: "AA",
+        maxExclusive: "BB",
+        ridPrefix: 0,
+        throughputFraction: 0,
+        status: "",
+        parents: [],
+      },
+      {
+        id: "3",
+        minInclusive: "BB",
+        maxExclusive: "DD",
+        ridPrefix: 0,
+        throughputFraction: 0,
+        status: "",
+        parents: [],
+      },
+    ]);
+
+    const result = await changeFeedForEpkRange.readNext();
+
+    // ensure getFeedResponse was called twice (first Gone, then Ok)
+    expect(feedSpy).toHaveBeenCalledTimes(2);
+
+    // verify the queue snapshot now has all three new feed‐ranges
+    const snapshot = (changeFeedForEpkRange as any).queue.returnSnapshot();
+    const ranges = snapshot.map((f: ChangeFeedRange) => [f.minInclusive, f.maxExclusive]);
+    expect(ranges).toEqual([
+      ["", "AA"],
+      ["AA", "BB"],
+      ["BB", "DD"],
+      ["DD", "FF"],
+    ]);
+    // final result should be the OK response
+    expect(result.statusCode).toBe(StatusCodes.Ok);
+  });
+
+  it("should modify first element when shiftLeft=true in handleSplitOrMerge", async () => {
+    // replace queue with a single original feed range
+    (changeFeedForEpkRange as any).queue = new FeedRangeQueue<ChangeFeedRange>();
+    const queue: FeedRangeQueue<ChangeFeedRange> = new FeedRangeQueue();
+    queue.enqueue(new ChangeFeedRange("", "DD", "cont1"));
+    queue.enqueue(new ChangeFeedRange("DD", "FF", "cont2"));
+    (changeFeedForEpkRange as any).queue = queue;
+    // prepare three resolved ranges
+    const resolvedRanges = [
+      { minInclusive: "", maxExclusive: "AA" },
+      { minInclusive: "AA", maxExclusive: "BB" },
+      { minInclusive: "BB", maxExclusive: "DD" },
+    ];
+    const oldQueryRange = new QueryRange("", "DD", true, false);
+
+    // invoke handleSplitOrMerge with shiftLeft=true
+    await (changeFeedForEpkRange as any).handleSplitOrMerge(
+      true,
+      resolvedRanges,
+      oldQueryRange,
+      "newCont",
+    );
+
+    const queueSnapshot = (changeFeedForEpkRange as any).queue.returnSnapshot();
+    const ranges = queueSnapshot.map((f: ChangeFeedRange) => [f.minInclusive, f.maxExclusive]);
+    expect(ranges).toEqual([
+      ["", "AA"],
+      ["DD", "FF"],
+      ["AA", "BB"],
+      ["BB", "DD"],
+    ]);
+  });
+
+  it("should handle partition merge retry accordingly", async () => {
+    // first call: return 410/Gone to trigger shouldRetryOnFailure
+    const queue: FeedRangeQueue<ChangeFeedRange> = new FeedRangeQueue();
+    queue.enqueue(new ChangeFeedRange("", "DD", "cont1"));
+    queue.enqueue(new ChangeFeedRange("DD", "FF", "cont2"));
+    (changeFeedForEpkRange as any).queue = queue;
+
+    const goneResp = new ChangeFeedIteratorResponse(
+      [],
+      0,
+      StatusCodes.Gone,
+      { [Constants.HttpHeaders.ETag]: "" },
+      getEmptyCosmosDiagnostics(),
+      SubStatusCodes.PartitionKeyRangeGone,
+    );
+    // second call: return OK
+    const okResp = new ChangeFeedIteratorResponse(
+      [{ id: "ok" }],
+      1,
+      StatusCodes.Ok,
+      {},
+      getEmptyCosmosDiagnostics(),
+    );
+    const feedSpy = vi
+      .spyOn(changeFeedForEpkRange as any, "getFeedResponse")
+      .mockResolvedValueOnce(goneResp)
+      .mockResolvedValueOnce(goneResp)
+      .mockResolvedValueOnce(okResp);
+
+    vi.spyOn(partitionKeyRangeCache, "getOverlappingRanges").mockResolvedValue([
+      {
+        id: "1",
+        minInclusive: "",
+        maxExclusive: "FF",
+        ridPrefix: 0,
+        throughputFraction: 0,
+        status: "",
+        parents: [],
+      },
+    ]);
+
+    await changeFeedForEpkRange.readNext();
+
+    // ensure getFeedResponse was called thrice
+    expect(feedSpy).toHaveBeenCalledTimes(3);
+    // verify the queue snapshot now has updated feed‐ranges
+    const queueFirstResult = (changeFeedForEpkRange as any).queue.dequeue();
+    const queueSecondResult = (changeFeedForEpkRange as any).queue.dequeue();
+    expect(queueFirstResult).toEqual({
+      minInclusive: "",
+      maxExclusive: "FF",
+      continuationToken: "cont2",
+      epkMinHeader: "DD",
+      epkMaxHeader: "FF",
+    });
+    expect(queueSecondResult).toEqual({
+      minInclusive: "",
+      maxExclusive: "FF",
+      continuationToken: undefined,
+      epkMinHeader: "",
+      epkMaxHeader: "DD",
+    });
+  });
+
+  it("should modify first element with merge when shiftLeft=true in handleSplitOrMerge", async () => {
+    (changeFeedForEpkRange as any).queue = new FeedRangeQueue<ChangeFeedRange>();
+    const queue: FeedRangeQueue<ChangeFeedRange> = new FeedRangeQueue();
+    queue.enqueue(new ChangeFeedRange("", "DD", "cont1"));
+    queue.enqueue(new ChangeFeedRange("DD", "FF", "cont2"));
+    (changeFeedForEpkRange as any).queue = queue;
+    // prepare three resolved ranges
+    const resolvedRanges = [{ minInclusive: "", maxExclusive: "FF" }];
+    const oldQueryRange = new QueryRange("", "DD", true, false);
+
+    // invoke handleSplitOrMerge with shiftLeft=true
+    await (changeFeedForEpkRange as any).handleSplitOrMerge(
+      true,
+      resolvedRanges,
+      oldQueryRange,
+      "newCont",
+    );
+    const queueFirstResult = (changeFeedForEpkRange as any).queue.dequeue();
+    expect(queueFirstResult).toEqual({
+      minInclusive: "",
+      maxExclusive: "FF",
+      continuationToken: "newCont",
+      epkMinHeader: "",
+      epkMaxHeader: "DD",
+    });
+    const queueSecondResult = (changeFeedForEpkRange as any).queue.dequeue();
+    expect(queueSecondResult).toEqual({
+      minInclusive: "DD",
+      maxExclusive: "FF",
+      continuationToken: "cont2",
+      epkMinHeader: undefined,
+      epkMaxHeader: undefined,
+    });
+  });
 });


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR fixes the bug where the changefeed for an entire container returns incorrect results after a partition merge by adjusting the logic for handling split and merge scenarios. Key changes include renaming the existing handleSplit method to handleSplitOrMerge to better represent its dual functionality, modifying logic in getPartitionRangeId to decide when to reprocess partition boundaries, and updating the extraction of effective feed range information using a queue peek.


<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**sdk/cosmosdb/cosmos/src/client/ChangeFeed/ChangeFeedForEpkRange.ts:313**
* Since the method has been renamed to handleSplitOrMerge to cover both split and merge scenarios, please update the accompanying TSDoc and inline comments to clearly reflect its new combined functionality.
```
await this.handleSplitOrMerge(false, resolvedRanges, queryRange, feedRange.continuationToken);
```
</details>


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T11:01:44Z
user: login: topshot99
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T07:14:38Z
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR addresses a bug in the changefeed logic for an entire container by updating the method handling partition splits and merges.  
- Renames and updates the method for handling split/merge scenarios.  
- Adds a new condition to correctly process partition range changes.  
- Updates usage of the feed range headers to fetch data from the queue.


<details>
<summary>Comments suppressed due to low confidence (2)</summary>

**sdk/cosmosdb/cosmos/src/client/ChangeFeed/ChangeFeedForEpkRange.ts:312**
* [nitpick] The method has been renamed to 'handleSplitOrMerge', but the corresponding inline comments and TSDoc should be updated to reflect its dual handling of both split and merge scenarios.
```
await this.handleSplitOrMerge(false, resolvedRanges, queryRange, feedRange.continuationToken);
```
**sdk/cosmosdb/cosmos/src/client/ChangeFeed/ChangeFeedForEpkRange.ts:392**
* [nitpick] Consider extracting the boundary comparisons into a well-named helper or constant to improve readability and maintainability.
```
const isPartitionRangeChanged = feedRange.minInclusive !== firstResolvedRange.minInclusive || feedRange.maxExclusive !== firstResolvedRange.maxExclusive || resolvedRanges.length > 1;
```
</details>


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T07:16:02Z
user: login: topshot99
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T15:49:00Z
user: login: amanrao23
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-09T16:21:12Z
user: login: topshot99
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-09T16:40:34Z

Review comments:
diff_hunk: @@ -437,6 +447,7 @@ export class ChangeFeedForEpkRange<T> implements ChangeFeedPullModelIterator<T>
     }
     try {
       // startEpk and endEpk are only valid in case we want to fetch result for a part of partition and not the entire partition.
+      const finalFeedRange = this.queue.peek();
path: sdk/cosmosdb/cosmos/src/client/ChangeFeed/ChangeFeedForEpkRange.ts
user: login: topshot99
body: wrap this.queue.peek(); into a method
created_at: 2025-05-09T07:14:38Z
author_association: MEMBER
diff_hunk: @@ -437,6 +447,7 @@ export class ChangeFeedForEpkRange<T> implements ChangeFeedPullModelIterator<T>
     }
     try {
       // startEpk and endEpk are only valid in case we want to fetch result for a part of partition and not the entire partition.
+      const finalFeedRange = this.queue.peek();
path: sdk/cosmosdb/cosmos/src/client/ChangeFeed/ChangeFeedForEpkRange.ts
user: login: Copilot
body: There is a potential risk that 'this.queue.peek()' may return an undefined value; ensure that finalFeedRange is valid before accessing its properties.
```suggestion
      const finalFeedRange = this.queue.peek();
      if (!finalFeedRange) {
        throw new Error("Queue is empty. Unable to retrieve a valid feed range.");
      }
```
created_at: 2025-05-09T07:16:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,4 +107,213 @@ describe("ChangeFeedForEpkRange Unit Tests", () => {
       assert.strictEqual(err.code, timeOutError.code);
     }
   });
+
+  it.skip("should handle partition split into three ranges and retry accordingly", async () => {
path: sdk/cosmosdb/cosmos/test/internal/unit/changeFeed/ChangeFeedForEpkRange.spec.ts
user: login: topshot99
body: skipped?
created_at: 2025-05-09T15:49:00Z
author_association: MEMBER
diff_hunk: @@ -94,4 +107,213 @@ describe("ChangeFeedForEpkRange Unit Tests", () => {
       assert.strictEqual(err.code, timeOutError.code);
     }
   });
+
+  it.skip("should handle partition split into three ranges and retry accordingly", async () => {
path: sdk/cosmosdb/cosmos/test/internal/unit/changeFeed/ChangeFeedForEpkRange.spec.ts
user: login: amanrao23
body: updated. 
created_at: 2025-05-09T16:21:12Z
author_association: MEMBER

PR comments:
user: login: azure-sdk
created_at: 2025-04-16T10:01:12Z
body: **API change check**

API changes are not detected in this pull request.
author_association: COLLABORATOR
user: login: amanrao23
created_at: 2025-05-07T11:04:35Z
body: /azp run js - cosmosdb -ci
author_association: MEMBER
user: login: azure-pipelines[bot]
created_at: 2025-05-07T11:04:42Z
body: <samp>
No pipelines are associated with this pull request.<br>

</samp>
author_association: NONE
user: login: amanrao23
created_at: 2025-05-07T11:05:52Z
body: /azp run js-cosmosdb-ci

author_association: MEMBER
user: login: azure-pipelines[bot]
created_at: 2025-05-07T11:05:59Z
body: <samp>
No pipelines are associated with this pull request.<br>

</samp>
author_association: NONE
user: login: topshot99
created_at: 2025-05-09T07:15:07Z
body: Add a unit test for merge scenario.
author_association: MEMBER