FeedRanges: Fixes GetFeedRangesAsync throwing DocumentClientException

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs

@@ -270,50 +270,57 @@ public async Task<IReadOnlyList<FeedRange>> GetFeedRangesAsync(
                 trace,
                 cancellationToken);
 
-            IReadOnlyList<PartitionKeyRange> partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync(
-                containerRId,
-                ContainerCore.allRanges,
-                trace,
-                forceRefresh: true);
-
-            if (partitionKeyRanges == null)
-            {
-                string refreshedContainerRId;
-                refreshedContainerRId = await this.GetCachedRIDAsync(
-                    forceRefresh: true,
-                    trace,
-                    cancellationToken);
-
-                if (string.Equals(containerRId, refreshedContainerRId))
-                {
-                    throw CosmosExceptionFactory.CreateInternalServerErrorException(
-                        $"Container rid {containerRId} did not have a partition key range after refresh",
-                        headers: new Headers(),
-                        trace: trace);
-                }
-
-                partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync(
-                    refreshedContainerRId,
-                    ContainerCore.allRanges,
-                    trace,
-                    forceRefresh: true);
-
-                if (partitionKeyRanges == null)
-                {
-                    throw CosmosExceptionFactory.CreateInternalServerErrorException(
-                        $"Container rid {containerRId} returned partitionKeyRanges null after Container RID refresh",
-                        headers: new Headers(),
-                        trace: trace);
-                }
-            }
-
-            List<FeedRange> feedTokens = new List<FeedRange>(partitionKeyRanges.Count);
-            foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges)
-            {
-                feedTokens.Add(new FeedRangeEpk(partitionKeyRange.ToRange()));
+            try
+            {
+                IReadOnlyList<PartitionKeyRange> partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync(
+                        containerRId,
+                        ContainerCore.allRanges,
+                        trace,
+                        forceRefresh: true);
+
+                if (partitionKeyRanges == null)
+                {
+                    string refreshedContainerRId;
+                    refreshedContainerRId = await this.GetCachedRIDAsync(
+                        forceRefresh: true,
+                        trace,
+                        cancellationToken);
+
+                    if (string.Equals(containerRId, refreshedContainerRId))
+                    {
+                        throw CosmosExceptionFactory.CreateInternalServerErrorException(
+                            $"Container rid {containerRId} did not have a partition key range after refresh",
+                            headers: new Headers(),
+                            trace: trace);
+                    }
+
+                    partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync(
+                        refreshedContainerRId,
+                        ContainerCore.allRanges,
+                        trace,
+                        forceRefresh: true);
+
+                    if (partitionKeyRanges == null)
+                    {
+                        throw CosmosExceptionFactory.CreateInternalServerErrorException(
+                            $"Container rid {containerRId} returned partitionKeyRanges null after Container RID refresh",
+                            headers: new Headers(),
+                            trace: trace);
+                    }
+                }
+
+                List<FeedRange> feedTokens = new List<FeedRange>(partitionKeyRanges.Count);
+                foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges)
+                {
+                    feedTokens.Add(new FeedRangeEpk(partitionKeyRange.ToRange()));
+                }
+
+                return feedTokens;
+            }
+            catch (DocumentClientException dce)
+            {
+                throw CosmosExceptionFactory.Create(dce, trace);
             }
-
-            return feedTokens;
         }
 
         public override FeedIterator GetChangeFeedStreamIterator(

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs

@@ -14,6 +14,8 @@ namespace Microsoft.Azure.Cosmos.Tests.FeedRange
     using Microsoft.Azure.Cosmos.Routing;
     using Moq;
     using Microsoft.Azure.Cosmos.Tracing;
+    using System.Net.Http;
+    using System.Text;
 
     [TestClass]
     public class FeedRangeTests
@@ -220,5 +222,58 @@ public void FeedRangePKRangeId_ToJsonFromJson()
             Assert.IsNotNull(feedRangePartitionKeyRangeDeserialized);
             Assert.AreEqual(feedRangePartitionKeyRange.PartitionKeyRangeId, feedRangePartitionKeyRangeDeserialized.PartitionKeyRangeId);
         }
+
+        /// <summary>
+        /// Upon failures in PartitionKeyRanges calls, the failure should be a CosmosException
+        /// </summary>
+        [TestMethod]
+        public async Task GetFeedRangesThrowsCosmosException()
+        {
+            Mock<IHttpHandler> mockHttpHandler = new Mock<IHttpHandler>();
+            Uri endpoint = MockSetupsHelper.SetupSingleRegionAccount(
+                "mockAccountInfo",
+                consistencyLevel: ConsistencyLevel.Session,
+                mockHttpHandler,
+                out string primaryRegionEndpoint);
+
+            string databaseName = "mockDbName";
+            string containerName = "mockContainerName";
+            string containerRid = "ccZ1ANCszwk=";
+            Documents.ResourceId cRid = Documents.ResourceId.Parse(containerRid);
+            MockSetupsHelper.SetupContainerProperties(
+                mockHttpHandler: mockHttpHandler,
+                regionEndpoint: primaryRegionEndpoint,
+                databaseName: databaseName,
+                containerName: containerName,
+                containerRid: containerRid);
+
+            // Return a 503 on PKRange call
+            bool invokedPkRanges = false;
+            Uri partitionKeyUri = new Uri($"{primaryRegionEndpoint}/dbs/{cRid.DatabaseId}/colls/{cRid.DocumentCollectionId}/pkranges");
+            mockHttpHandler.Setup(x => x.SendAsync(It.Is<HttpRequestMessage>(x => x.RequestUri == partitionKeyUri), It.IsAny<CancellationToken>()))
+              .Returns(() => Task.FromResult(new HttpResponseMessage()
+              {
+                  StatusCode = HttpStatusCode.ServiceUnavailable,
+                  Content = new StringContent("ServiceUnavailable")
+              }))
+              .Callback(() => invokedPkRanges = true);
+
+            CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
+            {
+                ConsistencyLevel = Cosmos.ConsistencyLevel.Session,
+                HttpClientFactory = () => new HttpClient(new HttpHandlerHelper(mockHttpHandler.Object)),
+            };
+
+            using (CosmosClient customClient = new CosmosClient(
+                   endpoint.ToString(),
+                   Convert.ToBase64String(Encoding.UTF8.GetBytes(Guid.NewGuid().ToString())),
+                   cosmosClientOptions))
+            {
+                Container container = customClient.GetContainer(databaseName, containerName);
+                CosmosException ex = await Assert.ThrowsExceptionAsync<CosmosException>(() => container.GetFeedRangesAsync(CancellationToken.None));
+                Assert.AreEqual(HttpStatusCode.ServiceUnavailable, ex.StatusCode);
+                Assert.IsTrue(invokedPkRanges);
+            }
+        }
     }
 }