public override async Task <CollectionRoutingMap> GetRoutingMapAsync(CancellationToken cancellationToken)
{
string collectionRid = await this.GetCachedRIDAsync(
forceRefresh : false,
trace : NoOpTrace.Singleton,
cancellationToken);
PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.Client.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
CollectionRoutingMap collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(
collectionRid,
previousValue : null,
request : null,
cancellationToken,
NoOpTrace.Singleton);
// Not found.
if (collectionRoutingMap == null)
{
collectionRid = await this.GetCachedRIDAsync(
forceRefresh : true,
trace : NoOpTrace.Singleton,
cancellationToken);
collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(
collectionRid,
previousValue : null,
request : null,
cancellationToken,
NoOpTrace.Singleton);
}
return(collectionRoutingMap);
}
private async Task <ShouldRetryResult> ShouldRetryInternalAsync(
HttpStatusCode?statusCode,
SubStatusCodes?subStatusCode,
CancellationToken cancellationToken)
{
if (statusCode == HttpStatusCode.Gone)
{
if (subStatusCode == SubStatusCodes.PartitionKeyRangeGone ||
subStatusCode == SubStatusCodes.CompletingSplit ||
subStatusCode == SubStatusCodes.CompletingPartitionMigration)
{
PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
string containerRid = await this.container.GetCachedRIDAsync(forceRefresh : false, cancellationToken : cancellationToken);
await partitionKeyRangeCache.TryGetOverlappingRangesAsync(containerRid, FeedRangeEpk.FullRange.Range, forceRefresh : true);
return(ShouldRetryResult.RetryAfter(TimeSpan.Zero));
}
if (subStatusCode == SubStatusCodes.NameCacheIsStale)
{
return(ShouldRetryResult.RetryAfter(TimeSpan.Zero));
}
}
return(null);
}
internal async Task <Tuple <string, ResponseMessage> > ReadNextInternalAsync(CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
if (this.compositeContinuationToken == null)
{
PartitionKeyRangeCache pkRangeCache = await this.clientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
this.containerRid = await this.container.GetRIDAsync(cancellationToken);
this.compositeContinuationToken = await StandByFeedContinuationToken.CreateAsync(this.containerRid, this.continuationToken, pkRangeCache.TryGetOverlappingRangesAsync);
}
(CompositeContinuationToken currentRangeToken, string rangeId) = await this.compositeContinuationToken.GetCurrentTokenAsync();
string partitionKeyRangeId = rangeId;
this.continuationToken = currentRangeToken.Token;
ResponseMessage response = await this.NextResultSetDelegateAsync(this.continuationToken, partitionKeyRangeId, this.maxItemCount, this.changeFeedOptions, cancellationToken);
if (await this.ShouldRetryFailureAsync(response, cancellationToken))
{
return(await this.ReadNextInternalAsync(cancellationToken));
}
if (response.IsSuccessStatusCode ||
response.StatusCode == HttpStatusCode.NotModified)
{
// Change Feed read uses Etag for continuation
currentRangeToken.Token = response.Headers.ETag;
}
return(new Tuple <string, ResponseMessage>(partitionKeyRangeId, response));
}
public async Task <IReadOnlyList <FeedRange> > GetFeedRangesAsync(
CosmosDiagnosticsContext diagnosticsContext,
CancellationToken cancellationToken = default)
{
PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
string containerRId = await this.GetRIDAsync(cancellationToken);
IReadOnlyList <PartitionKeyRange> partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync(
containerRId,
new Range <string>(
PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey,
PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey,
isMinInclusive : true,
isMaxInclusive : false),
forceRefresh : true);
List <FeedRange> feedTokens = new List <FeedRange>(partitionKeyRanges.Count);
foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges)
{
feedTokens.Add(new FeedRangeEpk(partitionKeyRange.ToRange()));
}
return(feedTokens);
}
public async Task <IReadOnlyList <FeedRange> > GetFeedRangesAsync(
ITrace trace,
CancellationToken cancellationToken = default)
{
PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(trace);
string containerRId;
containerRId = await this.GetCachedRIDAsync(
forceRefresh : false,
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(
containerRId,
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);
}
internal static async Task ApplySessionTokenAsync(
DocumentServiceRequest request,
ConsistencyLevel defaultConsistencyLevel,
ISessionContainer sessionContainer,
PartitionKeyRangeCache partitionKeyRangeCache,
CollectionCache clientCollectionCache,
IGlobalEndpointManager globalEndpointManager)
{
if (request.Headers == null)
{
Debug.Fail("DocumentServiceRequest does not have headers.");
return;
}
// Master resource operations don't require session token.
if (GatewayStoreModel.IsMasterOperation(request.ResourceType, request.OperationType))
{
if (!string.IsNullOrEmpty(request.Headers[HttpConstants.HttpHeaders.SessionToken]))
{
request.Headers.Remove(HttpConstants.HttpHeaders.SessionToken);
}
return;
}
if (!string.IsNullOrEmpty(request.Headers[HttpConstants.HttpHeaders.SessionToken]))
{
return; // User is explicitly controlling the session.
}
string requestConsistencyLevel = request.Headers[HttpConstants.HttpHeaders.ConsistencyLevel];
bool isReadOrBatchRequest = request.IsReadOnlyRequest || request.OperationType == OperationType.Batch;
bool requestHasConsistencySet = !string.IsNullOrEmpty(requestConsistencyLevel) && isReadOrBatchRequest; // Only read requests can have their consistency modified
bool sessionConsistencyApplies =
(!requestHasConsistencySet && defaultConsistencyLevel == ConsistencyLevel.Session) ||
(requestHasConsistencySet &&
string.Equals(requestConsistencyLevel, GatewayStoreModel.sessionConsistencyAsString, StringComparison.OrdinalIgnoreCase));
bool isMultiMasterEnabledForRequest = globalEndpointManager.CanUseMultipleWriteLocations(request);
if (!sessionConsistencyApplies ||
(!isReadOrBatchRequest &&
!isMultiMasterEnabledForRequest))
{
return; // Only apply the session token in case of session consistency and the request is read only or read/write on multimaster
}
(bool isSuccess, string sessionToken) = await GatewayStoreModel.TryResolveSessionTokenAsync(
request,
sessionContainer,
partitionKeyRangeCache,
clientCollectionCache);
if (isSuccess && !string.IsNullOrEmpty(sessionToken))
{
request.Headers[HttpConstants.HttpHeaders.SessionToken] = sessionToken;
}
}
public PartitionKeyRangeGoneRetryPolicy(
CollectionCache collectionCache,
PartitionKeyRangeCache partitionKeyRangeCache,
string collectionLink,
IDocumentClientRetryPolicy nextRetryPolicy)
: this(collectionCache, partitionKeyRangeCache, collectionLink, nextRetryPolicy, NoOpTrace.Singleton)
{
}
public override async Task <IReadOnlyList <PartitionKeyRange> > TryGetOverlappingRangesAsync(
string collectionResourceId,
Range <string> range,
bool forceRefresh = false)
{
PartitionKeyRangeCache partitionKeyRangeCache = await this.GetRoutingMapProviderAsync();
return(await partitionKeyRangeCache.TryGetOverlappingRangesAsync(collectionResourceId, range, forceRefresh));
}
internal static async Task ApplySessionTokenAsync(
DocumentServiceRequest request,
ConsistencyLevel defaultConsistencyLevel,
ISessionContainer sessionContainer,
PartitionKeyRangeCache partitionKeyRangeCache,
CollectionCache clientCollectionCache)
{
if (request.Headers == null)
{
Debug.Fail("DocumentServiceRequest does not have headers.");
return;
}
// Master resource operations don't require session token.
if (GatewayStoreModel.IsMasterOperation(request.ResourceType, request.OperationType))
{
if (!string.IsNullOrEmpty(request.Headers[HttpConstants.HttpHeaders.SessionToken]))
{
request.Headers.Remove(HttpConstants.HttpHeaders.SessionToken);
}
return;
}
if (!string.IsNullOrEmpty(request.Headers[HttpConstants.HttpHeaders.SessionToken]))
{
return; // User is explicitly controlling the session.
}
string requestConsistencyLevel = request.Headers[HttpConstants.HttpHeaders.ConsistencyLevel];
bool sessionConsistency =
defaultConsistencyLevel == ConsistencyLevel.Session ||
(!string.IsNullOrEmpty(requestConsistencyLevel) &&
string.Equals(requestConsistencyLevel, ConsistencyLevel.Session.ToString(), StringComparison.OrdinalIgnoreCase));
if (!sessionConsistency || (!request.IsReadOnlyRequest && request.OperationType != OperationType.Batch))
{
return; // Only apply the session token in case of session consistency and the request is read only
}
(bool isSuccess, string sessionToken) = await GatewayStoreModel.TryResolveSessionTokenAsync(
request,
sessionContainer,
partitionKeyRangeCache,
clientCollectionCache);
if (!isSuccess)
{
sessionToken = sessionContainer.ResolveGlobalSessionToken(request);
}
if (!string.IsNullOrEmpty(sessionToken))
{
request.Headers[HttpConstants.HttpHeaders.SessionToken] = sessionToken;
}
}
public PartitionKeyRangeGoneRetryPolicy(
CollectionCache collectionCache,
PartitionKeyRangeCache partitionKeyRangeCache,
string collectionLink,
IDocumentClientRetryPolicy nextRetryPolicy)
{
this.collectionCache = collectionCache;
this.partitionKeyRangeCache = partitionKeyRangeCache;
this.collectionLink = collectionLink;
this.nextRetryPolicy = nextRetryPolicy;
}
public override async Task <DocumentServiceLease> AcquireAsync(DocumentServiceLease lease)
{
if (lease == null)
{
throw new ArgumentNullException(nameof(lease));
}
string oldOwner = lease.Owner;
// We need to add the range information to any older leases
// This would not happen with new created leases but we need to be back compat
if (lease.FeedRange == null)
{
if (!this.lazyContainerRid.ValueInitialized)
{
TryCatch <string> tryInitializeContainerRId = await this.lazyContainerRid.GetValueAsync(NoOpTrace.Singleton, default);
if (!tryInitializeContainerRId.Succeeded)
{
throw tryInitializeContainerRId.Exception.InnerException;
}
this.partitionKeyRangeCache = await this.monitoredContainer.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
}
PartitionKeyRange partitionKeyRange = await this.partitionKeyRangeCache.TryGetPartitionKeyRangeByIdAsync(
this.lazyContainerRid.Result.Result,
lease.CurrentLeaseToken,
NoOpTrace.Singleton);
if (partitionKeyRange != null)
{
lease.FeedRange = new FeedRangeEpk(partitionKeyRange.ToRange());
}
}
return(await this.leaseUpdater.UpdateLeaseAsync(
lease,
lease.Id,
this.requestOptionsFactory.GetPartitionKey(lease.Id, lease.PartitionKey),
serverLease =>
{
if (serverLease.Owner != oldOwner)
{
DefaultTrace.TraceInformation("{0} lease token was taken over by owner '{1}'", lease.CurrentLeaseToken, serverLease.Owner);
throw new LeaseLostException(lease);
}
serverLease.Owner = this.options.HostName;
serverLease.Properties = lease.Properties;
return serverLease;
}).ConfigureAwait(false));
}
/// <summary>
/// Get the next set of results from the cosmos service
/// </summary>
/// <param name="cancellationToken">(Optional) <see cref="CancellationToken"/> representing request cancellation.</param>
/// <returns>A query response from cosmos service</returns>
public override async Task <ResponseMessage> ReadNextAsync(CancellationToken cancellationToken = default(CancellationToken))
{
cancellationToken.ThrowIfCancellationRequested();
if (this.compositeContinuationToken == null)
{
PartitionKeyRangeCache pkRangeCache = await this.clientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
this.containerRid = await this.container.GetRIDAsync(cancellationToken);
this.compositeContinuationToken = await StandByFeedContinuationToken.CreateAsync(this.containerRid, this.continuationToken, pkRangeCache.TryGetOverlappingRangesAsync);
}
(CompositeContinuationToken currentRangeToken, string rangeId) = await this.compositeContinuationToken.GetCurrentTokenAsync();
this.partitionKeyRangeId = rangeId;
this.continuationToken = currentRangeToken.Token;
ResponseMessage response = await this.NextResultSetDelegateAsync(this.continuationToken, this.partitionKeyRangeId, this.maxItemCount, this.changeFeedOptions, cancellationToken);
if (await this.ShouldRetryFailureAsync(response, cancellationToken))
{
cancellationToken.ThrowIfCancellationRequested();
(CompositeContinuationToken currentRangeTokenForRetry, string rangeIdForRetry) = await this.compositeContinuationToken.GetCurrentTokenAsync();
currentRangeToken = currentRangeTokenForRetry;
this.partitionKeyRangeId = rangeIdForRetry;
this.continuationToken = currentRangeToken.Token;
response = await this.NextResultSetDelegateAsync(this.continuationToken, this.partitionKeyRangeId, this.maxItemCount, this.changeFeedOptions, cancellationToken);
}
// Change Feed read uses Etag for continuation
string responseContinuationToken = response.Headers.ETag;
bool hasMoreResults = response.StatusCode != HttpStatusCode.NotModified;
if (!hasMoreResults)
{
currentRangeToken.Token = responseContinuationToken;
// Current Range is done, push it to the end
this.compositeContinuationToken.MoveToNextToken();
}
else if (response.IsSuccessStatusCode)
{
currentRangeToken.Token = responseContinuationToken;
}
// Send to the user the composite state for all ranges
response.Headers.ContinuationToken = this.compositeContinuationToken.ToString();
return(response);
}
private async Task <ShouldRetryResult> ShouldRetryInternalAsync(
HttpStatusCode?statusCode,
SubStatusCodes?subStatusCode,
CancellationToken cancellationToken)
{
if (statusCode == HttpStatusCode.Gone)
{
this.retriesOn410++;
if (this.retriesOn410 > MaxRetryOn410)
{
return(ShouldRetryResult.NoRetry());
}
if (subStatusCode == SubStatusCodes.PartitionKeyRangeGone ||
subStatusCode == SubStatusCodes.CompletingSplit ||
subStatusCode == SubStatusCodes.CompletingPartitionMigration)
{
PartitionKeyRangeCache partitionKeyRangeCache = await this.container.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton);
string containerRid = await this.container.GetCachedRIDAsync(
forceRefresh : false,
NoOpTrace.Singleton,
cancellationToken : cancellationToken);
await partitionKeyRangeCache.TryGetOverlappingRangesAsync(
containerRid,
FeedRangeEpk.FullRange.Range,
NoOpTrace.Singleton,
forceRefresh : true);
return(ShouldRetryResult.RetryAfter(TimeSpan.Zero));
}
if (subStatusCode == SubStatusCodes.NameCacheIsStale)
{
return(ShouldRetryResult.RetryAfter(TimeSpan.Zero));
}
}
// Batch API can return 413 which means the response is bigger than 4Mb.
// Operations that exceed the 4Mb limit are returned as 413/3402, while the operations within the 4Mb limit will be 200
if (statusCode == HttpStatusCode.RequestEntityTooLarge &&
(int)subStatusCode == SubstatusCodeBatchResponseSizeExceeded)
{
return(ShouldRetryResult.RetryAfter(TimeSpan.Zero));
}
return(null);
}
internal static async Task <Tuple <bool, string> > TryResolveSessionTokenAsync(
DocumentServiceRequest request,
ISessionContainer sessionContainer,
PartitionKeyRangeCache partitionKeyRangeCache,
CollectionCache clientCollectionCache)
{
if (request == null)
{
throw new ArgumentNullException(nameof(request));
}
if (sessionContainer == null)
{
throw new ArgumentNullException(nameof(sessionContainer));
}
if (partitionKeyRangeCache == null)
{
throw new ArgumentNullException(nameof(partitionKeyRangeCache));
}
if (clientCollectionCache == null)
{
throw new ArgumentNullException(nameof(clientCollectionCache));
}
if (request.ResourceType.IsPartitioned())
{
(bool isSuccess, PartitionKeyRange partitionKeyRange) = await TryResolvePartitionKeyRangeAsync(
request : request,
sessionContainer : sessionContainer,
partitionKeyRangeCache : partitionKeyRangeCache,
clientCollectionCache : clientCollectionCache,
refreshCache : false);
if (isSuccess && sessionContainer is SessionContainer gatewaySessionContainer)
{
request.RequestContext.ResolvedPartitionKeyRange = partitionKeyRange;
string localSessionToken = gatewaySessionContainer.ResolvePartitionLocalSessionTokenForGateway(request, partitionKeyRange.Id);
if (!string.IsNullOrEmpty(localSessionToken))
{
return(new Tuple <bool, string>(true, localSessionToken));
}
}
}
return(new Tuple <bool, string>(false, null));
}
async Task <IEnumerable <string> > GetPartitionKeyRangesAsync(
FeedToken feedToken,
CancellationToken cancellationToken = default(CancellationToken))
{
if (feedToken is FeedTokenEPKRange feedTokenEPKRange)
{
PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
string containerRId = await this.GetRIDAsync(cancellationToken);
IReadOnlyList <Documents.PartitionKeyRange> partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync(containerRId, feedTokenEPKRange.CompleteRange, forceRefresh : false);
return(partitionKeyRanges.Select(partitionKeyRange => partitionKeyRange.Id));
}
if (feedToken is FeedTokenPartitionKeyRange feedTokenPartitionKeyRange)
{
if (feedTokenPartitionKeyRange.FeedTokenEPKRange != null)
{
return(await this.GetPartitionKeyRangesAsync(feedTokenPartitionKeyRange.FeedTokenEPKRange, cancellationToken));
}
return(new List <string>()
{
feedTokenPartitionKeyRange.PartitionKeyRangeId
});
}
if (feedToken is FeedTokenPartitionKey feedTokenPartitionKey)
{
CollectionRoutingMap collectionRoutingMap = await this.GetRoutingMapAsync(cancellationToken);
PartitionKeyDefinition partitionKeyDefinition = await this.GetPartitionKeyDefinitionAsync(cancellationToken);
PartitionKeyInternal partitionKeyInternal = feedTokenPartitionKey.PartitionKey.InternalKey;
string effectivePartitionKeyString = partitionKeyInternal.GetEffectivePartitionKeyString(partitionKeyDefinition);
string partitionKeyRangeId = collectionRoutingMap.GetRangeByEffectivePartitionKey(effectivePartitionKeyString).Id;
return(new List <string>()
{
partitionKeyRangeId
});
}
throw new ArgumentException(nameof(feedToken), ClientResources.FeedToken_UnrecognizedFeedToken);
}
internal override Task <CollectionRoutingMap> GetRoutingMapAsync(CancellationToken cancellationToken)
{
string collectionRID = null;
return(this.GetRIDAsync(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());
}
public async Task GetChangeFeedTokensAsyncReturnsOnePerPartitionKeyRange()
{
// Setting mock to have 3 ranges, to generate 3 tokens
MultiRangeMockDocumentClient documentClient = new MultiRangeMockDocumentClient();
using CosmosClient client = MockCosmosUtil.CreateMockCosmosClient();
Mock <CosmosClientContext> mockContext = new Mock <CosmosClientContext>();
mockContext.Setup(x => x.ClientOptions).Returns(MockCosmosUtil.GetDefaultConfiguration());
mockContext.Setup(x => x.DocumentClient).Returns(documentClient);
mockContext.Setup(x => x.SerializerCore).Returns(MockCosmosUtil.Serializer);
mockContext.Setup(x => x.Client).Returns(client);
mockContext.Setup(x => x.CreateLink(It.IsAny <string>(), It.IsAny <string>(), It.IsAny <string>())).Returns(UriFactory.CreateDocumentCollectionUri("test", "test").OriginalString);
DatabaseInternal db = new DatabaseInlineCore(mockContext.Object, "test");
ContainerInternal container = new ContainerInlineCore(mockContext.Object, db, "test");
IEnumerable <string> tokens = await container.GetChangeFeedTokensAsync();
Assert.AreEqual(3, tokens.Count());
PartitionKeyRangeCache pkRangeCache = await documentClient.GetPartitionKeyRangeCacheAsync();
foreach (string token in tokens)
{
// Validate that each token represents a StandByFeedContinuationToken with a single Range
List <CompositeContinuationToken> deserialized = JsonConvert.DeserializeObject <List <CompositeContinuationToken> >(token);
Assert.AreEqual(1, deserialized.Count);
CompositeContinuationToken compositeToken = deserialized[0];
IReadOnlyList <Documents.PartitionKeyRange> rangesForTheToken = await pkRangeCache.TryGetOverlappingRangesAsync("", compositeToken.Range);
// Token represents one range
Assert.AreEqual(1, rangesForTheToken.Count);
Assert.AreEqual(rangesForTheToken[0].MinInclusive, compositeToken.Range.Min);
Assert.AreEqual(rangesForTheToken[0].MaxExclusive, compositeToken.Range.Max);
}
}
protected override async Task <DocumentFeedResponse <CosmosElement> > ExecuteInternalAsync(CancellationToken token)
{
CollectionCache collectionCache = await this.Client.GetCollectionCacheAsync();
PartitionKeyRangeCache partitionKeyRangeCache = await this.Client.GetPartitionKeyRangeCacheAsync();
IDocumentClientRetryPolicy retryPolicyInstance = this.Client.ResetSessionTokenRetryPolicy.GetRequestPolicy();
retryPolicyInstance = new InvalidPartitionExceptionRetryPolicy(retryPolicyInstance);
if (base.ResourceTypeEnum.IsPartitioned())
{
retryPolicyInstance = new PartitionKeyRangeGoneRetryPolicy(
collectionCache,
partitionKeyRangeCache,
PathsHelper.GetCollectionPath(base.ResourceLink),
retryPolicyInstance);
}
return(await BackoffRetryUtility <DocumentFeedResponse <CosmosElement> > .ExecuteAsync(
async() =>
{
this.fetchExecutionRangeAccumulator.BeginFetchRange();
++this.retries;
Tuple <DocumentFeedResponse <CosmosElement>, string> responseAndPartitionIdentifier = await this.ExecuteOnceAsync(retryPolicyInstance, token);
DocumentFeedResponse <CosmosElement> response = responseAndPartitionIdentifier.Item1;
string partitionIdentifier = responseAndPartitionIdentifier.Item2;
if (!string.IsNullOrEmpty(response.ResponseHeaders[HttpConstants.HttpHeaders.QueryMetrics]))
{
this.fetchExecutionRangeAccumulator.EndFetchRange(
partitionIdentifier,
response.ActivityId,
response.Count,
this.retries);
response = new DocumentFeedResponse <CosmosElement>(
response,
response.Count,
response.Headers,
response.UseETagAsContinuation,
new Dictionary <string, QueryMetrics>
{
{
partitionIdentifier,
QueryMetrics.CreateFromDelimitedStringAndClientSideMetrics(
response.ResponseHeaders[HttpConstants.HttpHeaders.QueryMetrics],
response.ResponseHeaders[HttpConstants.HttpHeaders.IndexUtilization],
new ClientSideMetrics(
this.retries,
response.RequestCharge,
this.fetchExecutionRangeAccumulator.GetExecutionRanges()))
}
},
response.RequestStatistics,
response.DisallowContinuationTokenMessage,
response.ResponseLengthBytes);
}
this.retries = -1;
return response;
},
retryPolicyInstance,
token));
}
internal async Task <(string, ResponseMessage)> ReadNextInternalAsync(CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
if (this.compositeContinuationToken == null)
{
PartitionKeyRangeCache pkRangeCache = await this.clientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
this.containerRid = await this.container.GetRIDAsync(cancellationToken);
if (this.changeFeedStartFrom is ChangeFeedStartFromContinuation startFromContinuation)
{
this.compositeContinuationToken = await StandByFeedContinuationToken.CreateAsync(
this.containerRid,
startFromContinuation.Continuation,
pkRangeCache.TryGetOverlappingRangesAsync);
(CompositeContinuationToken token, string id) = await this.compositeContinuationToken.GetCurrentTokenAsync();
if (token.Token != null)
{
this.changeFeedStartFrom = ChangeFeedStartFrom.ContinuationToken(token.Token);
}
else
{
this.changeFeedStartFrom = ChangeFeedStartFrom.Beginning();
}
}
else
{
this.compositeContinuationToken = await StandByFeedContinuationToken.CreateAsync(
this.containerRid,
initialStandByFeedContinuationToken : null,
pkRangeCache.TryGetOverlappingRangesAsync);
}
}
(CompositeContinuationToken currentRangeToken, string rangeId) = await this.compositeContinuationToken.GetCurrentTokenAsync();
FeedRange feedRange = new FeedRangePartitionKeyRange(rangeId);
if (currentRangeToken.Token != null)
{
this.changeFeedStartFrom = new ChangeFeedStartFromContinuationAndFeedRange(currentRangeToken.Token, (FeedRangeInternal)feedRange);
}
else
{
this.changeFeedStartFrom = ChangeFeedStartFrom.Beginning(feedRange);
}
ResponseMessage response = await this.NextResultSetDelegateAsync(this.changeFeedOptions, cancellationToken);
if (await this.ShouldRetryFailureAsync(response, cancellationToken))
{
return(await this.ReadNextInternalAsync(cancellationToken));
}
if (response.IsSuccessStatusCode ||
response.StatusCode == HttpStatusCode.NotModified)
{
// Change Feed read uses Etag for continuation
currentRangeToken.Token = response.Headers.ETag;
}
return(rangeId, response);
}
private static async Task <Tuple <bool, PartitionKeyRange> > TryResolvePartitionKeyRangeAsync(DocumentServiceRequest request,
ISessionContainer sessionContainer,
PartitionKeyRangeCache partitionKeyRangeCache,
ClientCollectionCache clientCollectionCache,
bool refreshCache)
{
if (refreshCache)
{
request.ForceMasterRefresh = true;
request.ForceNameCacheRefresh = true;
}
PartitionKeyRange partitonKeyRange = null;
ContainerProperties collection = await clientCollectionCache.ResolveCollectionAsync(request, CancellationToken.None, NoOpTrace.Singleton);
string partitionKeyString = request.Headers[HttpConstants.HttpHeaders.PartitionKey];
if (partitionKeyString != null)
{
CollectionRoutingMap collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(collectionRid : collection.ResourceId,
previousValue : null,
request : request,
cancellationToken : CancellationToken.None,
NoOpTrace.Singleton);
if (refreshCache && collectionRoutingMap != null)
{
collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(collectionRid : collection.ResourceId,
previousValue : collectionRoutingMap,
request : request,
cancellationToken : CancellationToken.None,
NoOpTrace.Singleton);
}
partitonKeyRange = AddressResolver.TryResolveServerPartitionByPartitionKey(request: request,
partitionKeyString: partitionKeyString,
collectionCacheUptoDate: false,
collection: collection,
routingMap: collectionRoutingMap);
}
else if (request.PartitionKeyRangeIdentity != null)
{
PartitionKeyRangeIdentity partitionKeyRangeId = request.PartitionKeyRangeIdentity;
partitonKeyRange = await partitionKeyRangeCache.TryGetPartitionKeyRangeByIdAsync(collection.ResourceId,
partitionKeyRangeId.ToString(),
NoOpTrace.Singleton,
refreshCache);
}
if (partitonKeyRange == null)
{
if (refreshCache)
{
return(new Tuple <bool, PartitionKeyRange>(false, null));
}
// need to refresh cache. Maybe split happened.
return(await GatewayStoreModel.TryResolvePartitionKeyRangeAsync(request : request,
sessionContainer : sessionContainer,
partitionKeyRangeCache : partitionKeyRangeCache,
clientCollectionCache : clientCollectionCache,
refreshCache : true));
}
return(new Tuple <bool, PartitionKeyRange>(true, partitonKeyRange));
}
public void SetCaches(PartitionKeyRangeCache partitionKeyRangeCache,
ClientCollectionCache clientCollectionCache)
{
this.clientCollectionCache = clientCollectionCache;
this.partitionKeyRangeCache = partitionKeyRangeCache;
}
public async Task <IReadOnlyList <FeedRange> > GetFeedRangesAsync(
CosmosDiagnosticsContext diagnosticsContext,
ITrace trace,
CancellationToken cancellationToken = default)
{
PartitionKeyRangeCache partitionKeyRangeCache = await this.ClientContext.DocumentClient.GetPartitionKeyRangeCacheAsync();
string containerRId;
using (diagnosticsContext.CreateScope(nameof(GetCachedRIDAsync)))
{
containerRId = await this.GetCachedRIDAsync(
forceRefresh : false,
cancellationToken);
}
IReadOnlyList <PartitionKeyRange> partitionKeyRanges;
using (diagnosticsContext.CreateScope(nameof(partitionKeyRangeCache.TryGetOverlappingRangesAsync)))
{
partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync(
containerRId,
ContainerCore.allRanges,
forceRefresh : true);
}
if (partitionKeyRanges == null)
{
string refreshedContainerRId;
using (diagnosticsContext.CreateScope("GetRIDAsyncForceRefresh"))
{
refreshedContainerRId = await this.GetCachedRIDAsync(
forceRefresh : true,
cancellationToken);
}
if (string.Equals(containerRId, refreshedContainerRId))
{
throw CosmosExceptionFactory.CreateInternalServerErrorException(
$"Container rid {containerRId} did not have a partition key range after refresh",
diagnosticsContext: diagnosticsContext);
}
using (diagnosticsContext.CreateScope(nameof(partitionKeyRangeCache.TryGetOverlappingRangesAsync)))
{
partitionKeyRanges = await partitionKeyRangeCache.TryGetOverlappingRangesAsync(
containerRId,
ContainerCore.allRanges,
forceRefresh : true);
}
if (partitionKeyRanges == null)
{
throw CosmosExceptionFactory.CreateInternalServerErrorException(
$"Container rid {containerRId} returned partitionKeyRanges null after Container RID refresh",
diagnosticsContext: diagnosticsContext);
}
}
List <FeedRange> feedTokens = new List <FeedRange>(partitionKeyRanges.Count);
foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges)
{
feedTokens.Add(new FeedRangeEpk(partitionKeyRange.ToRange()));
}
return(feedTokens);
}
protected override async Task <FeedResponse <dynamic> > ExecuteInternalAsync(CancellationToken cancellationToken)
{
CollectionCache collectionCache = await this.Client.GetCollectionCacheAsync();
PartitionKeyRangeCache partitionKeyRangeCache = await this.Client.GetPartitionKeyRangeCache();
IDocumentClientRetryPolicy retryPolicyInstance = this.Client.RetryPolicy.GetRequestPolicy();
retryPolicyInstance = new InvalidPartitionExceptionRetryPolicy(collectionCache, retryPolicyInstance);
if (base.ResourceTypeEnum.IsPartitioned())
{
retryPolicyInstance = new PartitionKeyRangeGoneRetryPolicy(
collectionCache,
partitionKeyRangeCache,
PathsHelper.GetCollectionPath(base.ResourceLink),
retryPolicyInstance);
}
return(await BackoffRetryUtility <FeedResponse <dynamic> > .ExecuteAsync(
async() =>
{
this.fetchExecutionRangeAccumulator.BeginFetchRange();
++this.retries;
this.fetchSchedulingMetrics.Start();
this.fetchExecutionRangeAccumulator.BeginFetchRange();
FeedResponse <dynamic> response = await this.ExecuteOnceAsync(retryPolicyInstance, cancellationToken);
this.fetchSchedulingMetrics.Stop();
this.fetchExecutionRangeAccumulator.EndFetchRange(response.Count, this.retries);
if (!string.IsNullOrEmpty(response.Headers[HttpConstants.HttpHeaders.QueryMetrics]))
{
this.fetchExecutionRangeAccumulator.EndFetchRange(response.Count, this.retries);
response = new FeedResponse <dynamic>(
response,
response.Count,
response.Headers,
response.UseETagAsContinuation,
new Dictionary <string, QueryMetrics>
{
{
singlePartitionKeyId,
QueryMetrics.CreateFromDelimitedStringAndClientSideMetrics(
response.Headers[HttpConstants.HttpHeaders.QueryMetrics],
new ClientSideMetrics(
this.retries,
response.RequestCharge,
this.fetchExecutionRangeAccumulator.GetExecutionRanges(),
string.IsNullOrEmpty(response.ResponseContinuation) ? new List <Tuple <string, SchedulingTimeSpan> >()
{
new Tuple <string, SchedulingTimeSpan>(singlePartitionKeyId, this.fetchSchedulingMetrics.Elapsed)
} : new List <Tuple <string, SchedulingTimeSpan> >()),
Guid.Parse(response.ActivityId))
}
},
response.RequestStatistics,
response.DisallowContinuationTokenMessage,
response.ResponseLengthBytes);
}
this.retries = -1;
return response;
},
retryPolicyInstance,
cancellationToken));
}
public override async Task <ResponseMessage> SendAsync(
RequestMessage request,
CancellationToken cancellationToken)
{
using (ITrace childTrace = request.Trace.StartChild(this.FullHandlerName, TraceComponent.RequestHandler, Tracing.TraceLevel.Info))
{
request.Trace = childTrace;
ResponseMessage response = null;
string originalContinuation = request.Headers.ContinuationToken;
try
{
RntdbEnumerationDirection rntdbEnumerationDirection = RntdbEnumerationDirection.Forward;
if (request.Properties.TryGetValue(HttpConstants.HttpHeaders.EnumerationDirection, out object direction))
{
rntdbEnumerationDirection = (byte)direction == (byte)RntdbEnumerationDirection.Reverse ? RntdbEnumerationDirection.Reverse : RntdbEnumerationDirection.Forward;
}
request.Headers.Remove(HttpConstants.HttpHeaders.IsContinuationExpected);
request.Headers.Add(HttpConstants.HttpHeaders.IsContinuationExpected, bool.TrueString);
if (!request.Properties.TryGetValue(HandlerConstants.StartEpkString, out object startEpk))
{
startEpk = PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey;
}
if (!request.Properties.TryGetValue(HandlerConstants.EndEpkString, out object endEpk))
{
endEpk = PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey;
}
startEpk ??= PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey;
endEpk ??= PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey;
List <Range <string> > providedRanges = new List <Range <string> >
{
new Range <string>(
(string)startEpk,
(string)endEpk,
isMinInclusive: true,
isMaxInclusive: false)
};
DocumentServiceRequest serviceRequest = request.ToDocumentServiceRequest();
PartitionKeyRangeCache routingMapProvider = await this.client.DocumentClient.GetPartitionKeyRangeCacheAsync();
CollectionCache collectionCache = await this.client.DocumentClient.GetCollectionCacheAsync(NoOpTrace.Singleton);
ContainerProperties collectionFromCache =
await collectionCache.ResolveCollectionAsync(serviceRequest, CancellationToken.None);
//direction is not expected to change between continuations.
Range <string> rangeFromContinuationToken =
this.partitionRoutingHelper.ExtractPartitionKeyRangeFromContinuationToken(serviceRequest.Headers, out List <CompositeContinuationToken> suppliedTokens);
ResolvedRangeInfo resolvedRangeInfo =
await this.partitionRoutingHelper.TryGetTargetRangeFromContinuationTokenRangeAsync(
providedPartitionKeyRanges : providedRanges,
routingMapProvider : routingMapProvider,
collectionRid : collectionFromCache.ResourceId,
rangeFromContinuationToken : rangeFromContinuationToken,
suppliedTokens : suppliedTokens,
direction : rntdbEnumerationDirection);
if (serviceRequest.IsNameBased && resolvedRangeInfo.ResolvedRange == null && resolvedRangeInfo.ContinuationTokens == null)
{
serviceRequest.ForceNameCacheRefresh = true;
collectionFromCache = await collectionCache.ResolveCollectionAsync(serviceRequest, CancellationToken.None);
resolvedRangeInfo = await this.partitionRoutingHelper.TryGetTargetRangeFromContinuationTokenRangeAsync(
providedPartitionKeyRanges : providedRanges,
routingMapProvider : routingMapProvider,
collectionRid : collectionFromCache.ResourceId,
rangeFromContinuationToken : rangeFromContinuationToken,
suppliedTokens : suppliedTokens,
direction : rntdbEnumerationDirection);
}
if (resolvedRangeInfo.ResolvedRange == null && resolvedRangeInfo.ContinuationTokens == null)
{
return(((DocumentClientException) new NotFoundException(
$"{DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture)}: Was not able to get queryRoutingInfo even after resolve collection async with force name cache refresh to the following collectionRid: {collectionFromCache.ResourceId} with the supplied tokens: {JsonConvert.SerializeObject(suppliedTokens)}")
).ToCosmosResponseMessage(request));
}
serviceRequest.RouteTo(new PartitionKeyRangeIdentity(collectionFromCache.ResourceId, resolvedRangeInfo.ResolvedRange.Id));
response = await base.SendAsync(request, cancellationToken);
if (!response.IsSuccessStatusCode)
{
this.SetOriginalContinuationToken(request, response, originalContinuation);
}
else
{
if (!await this.partitionRoutingHelper.TryAddPartitionKeyRangeToContinuationTokenAsync(
response.Headers.CosmosMessageHeaders,
providedPartitionKeyRanges: providedRanges,
routingMapProvider: routingMapProvider,
collectionRid: collectionFromCache.ResourceId,
resolvedRangeInfo: resolvedRangeInfo,
direction: rntdbEnumerationDirection))
{
return(((DocumentClientException) new NotFoundException(
$"{DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture)}: Call to TryAddPartitionKeyRangeToContinuationTokenAsync failed to the following collectionRid: {collectionFromCache.ResourceId} with the supplied tokens: {JsonConvert.SerializeObject(suppliedTokens)}")
).ToCosmosResponseMessage(request));
}
}
return(response);
}
catch (DocumentClientException ex)
{
ResponseMessage errorResponse = ex.ToCosmosResponseMessage(request);
this.SetOriginalContinuationToken(request, errorResponse, originalContinuation);
return(errorResponse);
}
catch (CosmosException ex)
{
ResponseMessage errorResponse = ex.ToCosmosResponseMessage(request);
this.SetOriginalContinuationToken(request, errorResponse, originalContinuation);
return(errorResponse);
}
catch (AggregateException ex)
{
this.SetOriginalContinuationToken(request, response, originalContinuation);
// TODO: because the SDK underneath this path uses ContinueWith or task.Result we need to catch AggregateExceptions here
// in order to ensure that underlying DocumentClientExceptions get propagated up correctly. Once all ContinueWith and .Result
// is removed this catch can be safely removed.
AggregateException innerExceptions = ex.Flatten();
Exception docClientException = innerExceptions.InnerExceptions.FirstOrDefault(innerEx => innerEx is DocumentClientException);
if (docClientException != null)
{
return(((DocumentClientException)docClientException).ToCosmosResponseMessage(request));
}
throw;
}
}
}
