Internal PartitionKeyRangeCache: Adds RID refresh on NotFound in GetR…

…outingMapAsync (#2096)

In GetRoutingMapAsync, we read the RID from cache & then look up the associated collection routing map.
In the case where the RID is stale & the collection routing map is not found, the stale RID should be invalidated & the operation retried.

I am also switching a method written using ContinueWith to be async instead as I found during testing that this style can cause AggregateException to be thrown which callers do not typically expect (e.g. a NotFound would end up as a 500 because of that).

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

@@ -448,12 +448,12 @@ public override async Task<string> GetCachedRIDAsync(
  return containerProperties?.ResourceId;
  }
 
- public override Task<PartitionKeyDefinition> GetPartitionKeyDefinitionAsync(CancellationToken cancellationToken = default)
+ public override async Task<PartitionKeyDefinition> GetPartitionKeyDefinitionAsync(CancellationToken cancellationToken = default)
  {
- return this.GetCachedContainerPropertiesAsync(
+ ContainerProperties cachedContainerPropertiesAsync = await this.GetCachedContainerPropertiesAsync(
  forceRefresh: false,
- cancellationToken: cancellationToken)
-  .ContinueWith(containerPropertiesTask => containerPropertiesTask.Result?.PartitionKey, cancellationToken);
+ cancellationToken: cancellationToken);
+ return cachedContainerPropertiesAsync?.PartitionKey;
  }
 
  /// <summary>
@@ -495,28 +495,34 @@ public override async Task<PartitionKeyInternal> GetNonePartitionKeyValueAsync(C
  return containerProperties.GetNoneValue();
  }
 
- public override Task<CollectionRoutingMap> GetRoutingMapAsync(CancellationToken cancellationToken)
+ public override async Task<CollectionRoutingMap> GetRoutingMapAsync(CancellationToken cancellationToken)
  {
- string collectionRID = null;
- return this.GetCachedRIDAsync(
+ string collectionRid = await this.GetCachedRIDAsync(
  forceRefresh: false,
- cancellationToken: cancellationToken)
- .ContinueWith(ridTask =>
- {
- collectionRID = ridTask.Result;
- return this.ClientContext.Client.DocumentClient.GetPartitionKeyRangeCacheAsync();
- })
- .Unwrap()
- .ContinueWith(partitionKeyRangeCachetask =>
- {
- PartitionKeyRangeCache partitionKeyRangeCache = partitionKeyRangeCachetask.Result;
- return partitionKeyRangeCache.TryLookupAsync(
- collectionRID,
- null,
- null,
- cancellationToken);
- })
- .Unwrap();
+ cancellationToken);
+
+ PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.Client.DocumentClient.GetPartitionKeyRangeCacheAsync();
+ CollectionRoutingMap collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(
+ collectionRid,
+ previousValue: null,
+ request: null,
+ cancellationToken);
+
+ // Not found.
+ if (collectionRoutingMap == null)
+ {
+ collectionRid = await this.GetCachedRIDAsync(
+ forceRefresh: true,
+ cancellationToken);
+
+ collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(
+ collectionRid,
+ previousValue: null,
+ request: null,
+ cancellationToken);
+ }
+
+ return collectionRoutingMap;
  }
 
  private async Task<ThroughputResponse> OfferRetryHelperForStaleRidCacheAsync(

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/PartitionKeyRangeCacheTests.cs

@@ -0,0 +1,53 @@
+// unset
+
+namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
+{
+ using System;
+ using System.Net;
+ using System.Threading.Tasks;
+ using Microsoft.Azure.Cosmos.Routing;
+ using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+ [TestClass]
+ public class PartitionKeyRangeCacheTests
+ {
+ [TestMethod]
+ [Owner("flnarenj")]
+ public async Task TestRidRefreshOnNotFoundAsync()
+ {
+ CosmosClient resourceClient = TestCommon.CreateCosmosClient();
+
+ string dbName = Guid.NewGuid().ToString();
+ string containerName = Guid.NewGuid().ToString();
+
+ Database db = await resourceClient.CreateDatabaseAsync(dbName);
+ Container container = await db.CreateContainerAsync(containerName, "/_id");
+
+ CosmosClient testClient = TestCommon.CreateCosmosClient();
+ ContainerInternal testContainer = (ContainerInlineCore)testClient.GetContainer(dbName, containerName);
+
+ // Populate the RID cache.
+ string cachedRidAsync = await testContainer.GetCachedRIDAsync(forceRefresh: false);
+
+ // Delete the container (using resource client).
+ await container.DeleteContainerAsync();
+
+ // Because the RID is cached, this will now try to resolve the collection routing map.
+ Assert.AreEqual(cachedRidAsync, await testContainer.GetCachedRIDAsync(forceRefresh: false));
+ CosmosException notFoundException = await Assert.ThrowsExceptionAsync<CosmosException>(() => testContainer.GetRoutingMapAsync(cancellationToken: default));
+ Assert.AreEqual(HttpStatusCode.NotFound, notFoundException.StatusCode);
+
+ await db.CreateContainerAsync(containerName, "/_id");
+
+ CollectionRoutingMap collectionRoutingMap = await testContainer.GetRoutingMapAsync(cancellationToken: default);
+ Assert.IsNotNull(collectionRoutingMap);
+ Assert.AreNotEqual(cachedRidAsync, await testContainer.GetCachedRIDAsync(forceRefresh: false));
+
+ // Delete the container (using resource client).
+ await container.DeleteContainerAsync();
+
+ CollectionRoutingMap collectionRoutingMapFromCache = await testContainer.GetRoutingMapAsync(cancellationToken: default);
+ Assert.AreEqual(collectionRoutingMap, collectionRoutingMapFromCache);
+ }
+ }
+}