Skip to content

Commit

Permalink
FeedRanges: Fixes GetFeedRangesAsync throwing DocumentClientException (…
Browse files Browse the repository at this point in the history 
…#4640)

* Fix

* tests

* Remove unused
  • Loading branch information
@ealsur
ealsur authored Aug 13, 2024
1 parent 47596bd commit 7bc8e3c
Show file tree
Hide file tree
Showing 2 changed files with 105 additions and 43 deletions.
93 changes: 50 additions & 43 deletions Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
55 changes: 55 additions & 0 deletions Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/FeedRangeTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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);
}
}
}
}

0 comments on commit 7bc8e3c

Please sign in to comment.