QueryPlan : Fixes 410 Gone Exception on non-x64 platforms (#5257)

Windows 365 Defender team faced an issue with their HPK container set
up. Problem was when they set the root level partition key in the
QueryRequestOptions and run a query they would get the data back when
running on windows but for other environments(like docker) they were
getting an exception
For Docker they were getting 410(Gone) exception 
Microsoft.Azure.Cosmos.CosmosException : Response status code does not
indicate success: Gone (410); Substatus: 1002; ActivityId:
623305e2-9a69-4b65-aa9b-911ee8ac3533; Reason: (Epk Range:
[0DCEB8CE51C6BFE84F4BD9409F69B9BB,0DCEB8CE51C6BFE84F4BD9409F69B9BBFF) is
gone.);

The root of the problem was that for non-window(x64) environment we were
getting the query plan from gateway(unlike windows where we use the
ServiceInterop dll) and after we got the plan we funnel it through
RequestInvokerHandler pipeline. The pipeline has a check where it sees
that if a feedRange is provided and if it returns more than 1
overlapping ranges it throws a gone exception but in this particular
case this is a valid scenario(edge case).
e.g
datatable (MinEPK: string, MaxEPK: string, PKrangeId) 
[

"0D4DC2CD8F49C65A8E0C5306B61B4343","0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED",0

"0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED","1080F600C27CF98DC13F8639E94E7676"
1
]
If our provided partition key results in a hash e.g
0DCEB8CE51C6BFE84F4BD9409F69B9BB that lies between the above 2 records,
we will get overlapping ranges back.

Fix:
Fix it to explicitly make a check and see if its a Query Plan , if it is
then allow it to progress.

Testing:
I was able to get access to the Defender Team DB where they were facing
this issue. I manually tested the following E2E use cases and got the
desired results.
On Windows -> with PartitionKey specified
Diagnostics is : {"Summary":{"DirectCalls":{"(449, 5350)":6,"(200,
0)":2},"GatewayCalls":{"(200, 0)":4,"(304, 0)":1}}
Count: 24025; Severity: Medium; 
Count: 24114; Severity: Informational; 
Count: 23; Severity: High; 

On Non_windows - with PartitionKey specified(in the QueryRequestOptions)
Diagnostics is : {"Summary":{"DirectCalls":{"(449, 5350)":2,"(200,
0)":2},"GatewayCalls":{"(200, 0)":5,"(304, 0)":1,"(0, 0)":1}}
Count: 24025; Severity: Medium; 
Count: 24114; Severity: Informational; 
Count: 23; Severity: High;

On Windows - no QueryRequestOption specified
Diagnostics is : {"Summary":{"DirectCalls":{"(200,
0)":2},"GatewayCalls":{"(200, 0)":4,"(304, 0)":1}}
Count: 24025; Severity: Medium; 
Count: 24114; Severity: Informational; 
Count: 23; Severity: High; 

On Non_windows - no QueryRequestOption specified
Diagnostics is : {"Summary":{"DirectCalls":{"(200,
0)":2},"GatewayCalls":{"(200, 0)":5,"(304, 0)":1}}
Count: 24025; Severity: Medium; 
Count: 24114; Severity: Informational; 
Count: 23; Severity: High; 

Problem writing a new reliable E2E test in the pipeline is the set up
needed to return overlapping ranges and we don't have control over when
partition split happens.

Talked to Ananth on how he tested it and he worked with Elasticity team
to create the database in the desired state, we can do the same in the
pipeline but need to check if our database accounts/containers get
cleaned up by some background job in the future(will check with Nalu on
this one).
I also explored the PartitionKeyHashRangeSplitterAndMerger class but
that works only with in-memory container and doesn't follow the
RequestInvokerHandler pipeline flow.

for the time being I added a unit test which will avoid any regression
in the future.


closes #5220
#5220

---------

Co-authored-by: Kiran Kumar Kolli <kirankk@microsoft.com>

Author: dibahlfi

Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs

@@ -305,17 +305,22 @@ public virtual async Task<ResponseMessage> SendAsync(
                             // For epk range filtering we can end up in one of 3 cases:
                             if (overlappingRanges.Count > 1)
                             {
-                                // 1) The EpkRange spans more than one physical partition
-                                // In this case it means we have encountered a split and 
-                                // we need to bubble that up to the higher layers to update their datastructures
-                                CosmosException goneException = new CosmosException(
-                                    message: $"Epk Range: {feedRangeEpk.Range} is gone.",
-                                    statusCode: System.Net.HttpStatusCode.Gone,
-                                    subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone,
-                                    activityId: Guid.NewGuid().ToString(),
-                                    requestCharge: default);
-
-                                return goneException.ToCosmosResponseMessage(request);
+                                //If we are running a query plan and our provided partition key results in a hash that resolves to more than one EPKRanges then its a valid use case
+                                bool isQueryPlanOperation = request.ResourceType == ResourceType.Document && request.OperationType == OperationType.QueryPlan;
+                                if (!isQueryPlanOperation)
+                                {
+                                    // 1) The EpkRange spans more than one physical partition
+                                    // In this case it means we have encountered a split and 
+                                    // we need to bubble that up to the higher layers to update their datastructures
+                                    CosmosException goneException = new CosmosException(
+                                        message: $"Epk Range: {feedRangeEpk.Range} is gone.",
+                                        statusCode: System.Net.HttpStatusCode.Gone,
+                                        subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone,
+                                        activityId: Guid.NewGuid().ToString(),
+                                        requestCharge: default);
+
+                                    return goneException.ToCosmosResponseMessage(request);
+                                }
                             }
                             // overlappingRanges.Count == 1
                             else

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/RetryHandlerTests.cs

@@ -5,11 +5,16 @@
 namespace Microsoft.Azure.Cosmos.Tests
 {
     using System;
+    using System.Collections.Generic;
+    using System.IO;
     using System.Net;
     using System.Net.Http;
     using System.Threading;
     using System.Threading.Tasks;
+    using global::Azure;
     using Microsoft.Azure.Cosmos.Handlers;
+    using Microsoft.Azure.Cosmos.Routing;
+    using Microsoft.Azure.Cosmos.Tracing;
     using Microsoft.Azure.Documents;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
     using Moq;
@@ -18,6 +23,139 @@ namespace Microsoft.Azure.Cosmos.Tests
     public class RetryHandlerTests
     {
         private static readonly Uri TestUri = new Uri("https://dummy.documents.azure.com:443/dbs");
+        [TestMethod]
+        public async Task ValidateQueryPlanDoesNotThrowExceptionForOverlappingRanges()
+        {
+            await this.ValidateOverlappingRangesBehaviorAsync(
+                operationType: OperationType.QueryPlan,
+                shouldThrowGoneException: false);
+        }
+
+        [TestMethod]
+        public async Task ValidateQueryThrowsGoneExceptionForOverlappingRanges()
+        {
+            await this.ValidateOverlappingRangesBehaviorAsync(
+                operationType: OperationType.Query,
+                shouldThrowGoneException: true);
+        }
+
+        private async Task ValidateOverlappingRangesBehaviorAsync(
+            OperationType operationType,
+            bool shouldThrowGoneException)
+        {
+            // Create overlapping ranges for the test
+            List<PartitionKeyRange> overlappingRanges = new List<PartitionKeyRange>
+            {
+                new PartitionKeyRange { Id = "0", MinInclusive = "0D4DC2CD8F49C65A8E0C5306B61B4343", MaxExclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED" },
+                new PartitionKeyRange { Id = "1", MinInclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED", MaxExclusive = "1080F600C27CF98DC13F8639E94E7676" }
+            };
+
+            // Create a custom document client with our TestPartitionKeyRangeCache
+            var testPartitionKeyRangeCache = new TestPartitionKeyRangeCache(overlappingRanges);
+            var customDocClient = new CustomMockDocumentClient(testPartitionKeyRangeCache);
+
+            // Create CosmosClient with our custom document client
+            using CosmosClient client = new CosmosClient(
+                "https://localhost:8081",
+                MockCosmosUtil.RandomInvalidCorrectlyFormatedAuthKey,
+                new CosmosClientOptions(),
+                customDocClient);
+
+            // Create mock container 
+            Mock<ContainerInternal> containerMock = MockCosmosUtil.CreateMockContainer("testDb", "testColl");
+
+            // Setup container properties
+            ContainerProperties containerProps = new ContainerProperties("testColl", "/pk");
+            var resourceIdProperty = typeof(ContainerProperties).GetProperty(
+                "ResourceId",
+                System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance);
+            resourceIdProperty.SetValue(containerProps, "testCollRid");
+
+            // Set up additional mocks as needed
+            containerMock.Setup(c => c.GetCachedContainerPropertiesAsync(
+                It.IsAny<bool>(), It.IsAny<ITrace>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(containerProps);
+
+            Mock<Cosmos.Database> databaseMock = new Mock<Cosmos.Database>();
+            databaseMock.Setup(d => d.Id).Returns("testDb");
+            containerMock.Setup(c => c.Database).Returns(databaseMock.Object);
+
+            // FeedRangeEpk for the test - use a range that overlaps both partition key ranges
+            FeedRangeEpk feedRange = new FeedRangeEpk(new Documents.Routing.Range<string>(
+                "0DCEB8CE51C6BFE84F4BD9409F69B9BB",
+                "0DCEB8CE51C6BFE84F4BD9409F69B9BBFF", 
+                true, false));
+
+            RequestInvokerHandler invoker = new RequestInvokerHandler(client, null, null, null)
+            {
+                InnerHandler = new TestHandler((request, token) => TestHandler.ReturnSuccess())
+            };
+
+            // Act
+            ResponseMessage response = await invoker.SendAsync(
+                "dbs/testDb/colls/testColl",
+                ResourceType.Document,
+                operationType,
+                null,
+                containerMock.Object,
+                feedRange,
+                null,
+                null,
+                NoOpTrace.Singleton,
+                CancellationToken.None);
+
+            // Assert
+            Assert.IsNotNull(response, "Response should not be null.");
+
+            if (shouldThrowGoneException)
+            {
+                Assert.IsFalse(response.IsSuccessStatusCode, "Expected a failure status code for Query operation.");
+                Assert.AreEqual(HttpStatusCode.Gone, response.StatusCode, "Expected a 410 Gone status code.");
+                Assert.AreEqual((int)SubStatusCodes.PartitionKeyRangeGone, (int)response.Headers.SubStatusCode, "Expected PartitionKeyRangeGone sub-status code.");
+            }
+            else
+            {
+                Assert.IsTrue(response.IsSuccessStatusCode, $"Expected a successful status code, but got {response.StatusCode}.");
+            }
+        }
+
+        // Custom MockDocumentClient that allows injecting our TestPartitionKeyRangeCache
+        private class CustomMockDocumentClient : MockDocumentClient
+        {
+            private readonly TestPartitionKeyRangeCache testPartitionKeyRangeCache;
+
+            public CustomMockDocumentClient(TestPartitionKeyRangeCache testPartitionKeyRangeCache)
+                : base(new ConnectionPolicy())
+            {
+                this.testPartitionKeyRangeCache = testPartitionKeyRangeCache;
+            }
+
+            internal override Task<PartitionKeyRangeCache> GetPartitionKeyRangeCacheAsync(ITrace trace)
+            {
+                return Task.FromResult<PartitionKeyRangeCache>(this.testPartitionKeyRangeCache);
+            }
+        }
+
+        private class TestPartitionKeyRangeCache : PartitionKeyRangeCache
+        {
+            private readonly IReadOnlyList<PartitionKeyRange> overlappingRanges;
+
+            public TestPartitionKeyRangeCache(IReadOnlyList<PartitionKeyRange> overlappingRanges)
+                : base(null, null, null, null) // Pass nulls or mocks as needed for base constructor
+            {
+                this.overlappingRanges = overlappingRanges;
+            }
+
+            public override Task<IReadOnlyList<PartitionKeyRange>> TryGetOverlappingRangesAsync(
+                string collectionRid,
+                Documents.Routing.Range<string> range,
+                ITrace trace,
+                bool forceRefresh)
+            {
+                return Task.FromResult(this.overlappingRanges);
+            }
+        }
+
 
         [TestMethod]
         public async Task RetryHandlerDoesNotRetryOnSuccess()