public async Task TestOperationCanceledExceptionAsync()
{
int seed = (int)(DateTime.UtcNow - new DateTime(1970, 1, 1)).TotalSeconds;
Random rand = new Random(seed);
IDocumentClientRetryPolicy retryPolicy = new MockRetryPolicy(rand);
ComparableTaskScheduler scheduler = new ComparableTaskScheduler(1);
DocumentProducer <int> producer = new DocumentProducer <int>(
scheduler,
(continuation, pageSize) => null,
new PartitionKeyRange {
Id = "test", MinInclusive = "", MaxExclusive = "ff"
},
p => 0,
(request, token) =>
{
scheduler.Stop();
throw new Exception();
},
() => retryPolicy,
(produer, size, ru, queryMetrics, token, length) => { },
Guid.NewGuid()
)
{
PageSize = 1000
};
await producer.MoveNextAsync(new CancellationTokenSource().Token);
}
private Task <bool> TryMoveNextProducerAsync(DocumentProducer <object> producer, CancellationToken cancellationToken)
{
return(base.TryMoveNextProducerAsync(
producer,
(currentProducer) => this.RepairParallelContext(currentProducer),
cancellationToken));
}
public override async Task <FeedResponse <object> > DrainAsync(int maxElements, CancellationToken token)
{
List <object> result = new List <object>();
List <Uri> replicaUris = new List <Uri>();
ClientSideRequestStatistics requestStats = new ClientSideRequestStatistics();
while (!this.IsDone)
{
DocumentProducer <object> currentDocumentProducer = base.DocumentProducers[this.currentDocumentProducerIndex];
if (currentDocumentProducer.IsAtContinuationBoundary)
{
if (maxElements - result.Count < currentDocumentProducer.ItemsTillNextContinuationBoundary)
{
break;
}
base.CurrentContinuationTokens[currentDocumentProducer] = currentDocumentProducer.ResponseContinuation;
result.Add(currentDocumentProducer.Current);
if (currentDocumentProducer.RequestStatistics != null)
{
replicaUris.AddRange(currentDocumentProducer.RequestStatistics.ContactedReplicas);
}
while (currentDocumentProducer.ItemsTillNextContinuationBoundary > 1)
{
bool hasMoreResults = await currentDocumentProducer.MoveNextAsync(token);
Debug.Assert(hasMoreResults, "Expect hasMoreResults be true.");
result.Add(currentDocumentProducer.Current);
}
}
if (!await this.TryMoveNextProducerAsync(currentDocumentProducer, token))
{
++this.currentDocumentProducerIndex;
if (this.currentDocumentProducerIndex < base.DocumentProducers.Count &&
!base.CurrentContinuationTokens.ContainsKey(base.DocumentProducers[this.currentDocumentProducerIndex]))
{
base.CurrentContinuationTokens[base.DocumentProducers[this.currentDocumentProducerIndex]] = null;
}
continue;
}
if (maxElements >= int.MaxValue && result.Count > this.ActualMaxBufferedItemCount)
{
break;
}
}
this.ReduceTotalBufferedItems(result.Count);
requestStats.ContactedReplicas.AddRange(replicaUris);
return(new FeedResponse <object>(result, result.Count, this.ResponseHeaders, requestStats, this.GetAndResetResponseLengthBytes()));
}
public override async Task <FeedResponse <object> > DrainAsync(int maxElements, CancellationToken token)
{
List <object> result = new List <object>();
List <Uri> replicaUris = new List <Uri>();
ClientSideRequestStatistics requestStats = new ClientSideRequestStatistics();
if (maxElements >= int.MaxValue)
{
maxElements = this.ActualMaxBufferedItemCount;
}
while (!this.IsDone && result.Count < maxElements)
{
this.UpdateCurrentDocumentProducer();
DefaultTrace.TraceVerbose(string.Format(
CultureInfo.InvariantCulture,
"{0}, CorrelatedActivityId: {1}, ActivityId: {2} | OrderBy~Context.DrainAsync, currentDocumentProducer.Id: {3}, documentProducerConsumeQueue.Count: {4}",
DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture),
this.CorrelatedActivityId,
this.ActivityId,
this.currentDocumentProducer.TargetRange.Id,
this.documentProducerConsumerQueue.Count));
OrderByQueryResult orderByResult = (OrderByQueryResult)this.currentDocumentProducer.Current;
result.Add(orderByResult.Payload);
if (this.currentDocumentProducer.RequestStatistics != null)
{
replicaUris.AddRange(this.currentDocumentProducer.RequestStatistics.ContactedReplicas);
}
if (this.ShouldIncrementSkipCount(orderByResult.Rid))
{
++this.skipCount;
}
else
{
this.skipCount = 0;
}
this.previousRid = orderByResult.Rid;
if (!await this.TryMoveNextProducerAsync(
this.currentDocumentProducer,
null,
cancellationToken: token))
{
this.currentDocumentProducer = null;
}
}
this.ReduceTotalBufferedItems(result.Count);
requestStats.ContactedReplicas.AddRange(replicaUris);
return(new FeedResponse <object>(result, result.Count, this.ResponseHeaders, requestStats, this.GetAndResetResponseLengthBytes()));
}
/// <summary>
/// Gets whether or not we should increment the skip count based on the rid of the document.
/// </summary>
/// <param name="currentDocumentProducer">The current document producer.</param>
/// <returns>Whether or not we should increment the skip count.</returns>
private bool ShouldIncrementSkipCount(DocumentProducer currentDocumentProducer)
{
// If we are not at the begining of the page and we saw the same rid again.
return(!currentDocumentProducer.IsAtBeginningOfPage &&
string.Equals(
this.previousRid,
new OrderByQueryResult(currentDocumentProducer.Current).Rid,
StringComparison.Ordinal));
}
private Task <bool> TryMoveNextProducerAsync(
DocumentProducer <OrderByQueryResult> producer,
Dictionary <string, OrderByContinuationToken> targetRangeToOrderByContinuationMap,
CancellationToken cancellationToken)
{
return(base.TryMoveNextProducerAsync(
producer,
(currentProducer) => this.RepairOrderByContext(currentProducer, targetRangeToOrderByContinuationMap),
cancellationToken));
}
private async Task <DocumentProducer <object> > RepairParallelContext(DocumentProducer <object> producer)
{
List <PartitionKeyRange> replacementRanges = await base.GetReplacementRanges(producer.TargetRange, this.collectionRid);
await base.RepairContextAsync(
this.collectionRid,
this.currentDocumentProducerIndex,
this.taskPriorityFunc,
replacementRanges,
this.QuerySpec,
() => this.PopulateDocumentProducerPositionCache(this.currentDocumentProducerIndex));
return(base.DocumentProducers[this.currentDocumentProducerIndex]);
}
private async Task <DocumentProducer <OrderByQueryResult> > RepairOrderByContext(
DocumentProducer <OrderByQueryResult> parentProducer,
Dictionary <string, OrderByContinuationToken> targetRangeToOrderByContinuationMap)
{
List <PartitionKeyRange> replacementRanges = await base.GetReplacementRanges(parentProducer.TargetRange, this.collectionRid);
string parentRangeId = parentProducer.TargetRange.Id;
int indexOfCurrentDocumentProducer = base.DocumentProducers.BinarySearch(
parentProducer,
Comparer <DocumentProducer <OrderByQueryResult> > .Create(
(producer1, producer2) => string.CompareOrdinal(producer1.TargetRange.MinInclusive, producer2.TargetRange.MinInclusive)));
Debug.Assert(indexOfCurrentDocumentProducer >= 0, "Index of a producer in the Producers list can't be < 0");
// default filter is "true", since it gets used when we replace the FormatPlaceHolder and if there is no parent filter, then the query becomes
// SELECT * FROM c where blah and true
string parentFilter;
if (!this.filters.TryGetValue(parentRangeId, out parentFilter))
{
parentFilter = True;
}
replacementRanges.ForEach(pkr => this.filters.Add(pkr.Id, parentFilter));
await base.RepairContextAsync(
this.collectionRid,
indexOfCurrentDocumentProducer,
TaskPriorityFunc,
replacementRanges,
new SqlQuerySpec(
this.QuerySpec.QueryText.Replace(FormatPlaceHolder, parentFilter),
this.QuerySpec.Parameters));
this.filters.Remove(parentRangeId);
if (targetRangeToOrderByContinuationMap != null && targetRangeToOrderByContinuationMap.ContainsKey(parentRangeId))
{
for (int index = 0; index < replacementRanges.Count; ++index)
{
targetRangeToOrderByContinuationMap[replacementRanges[index].Id] = targetRangeToOrderByContinuationMap[parentRangeId];
targetRangeToOrderByContinuationMap[replacementRanges[index].Id].CompositeToken.Range = replacementRanges[index].ToRange();
}
targetRangeToOrderByContinuationMap.Remove(parentRangeId);
}
return(base.DocumentProducers[indexOfCurrentDocumentProducer]);
}
private void UpdateCurrentDocumentProducer()
{
if (this.documentProducerConsumerQueue.Count > 0)
{
if (this.currentDocumentProducer == null)
{
this.currentDocumentProducer = this.documentProducerConsumerQueue.Dequeue();
}
else if (this.documentProducerConsumerQueue.Comparer.Compare(this.currentDocumentProducer, this.documentProducerConsumerQueue.Peek()) > 0)
{
this.documentProducerConsumerQueue.Enqueue(this.currentDocumentProducer);
this.currentDocumentProducer = this.documentProducerConsumerQueue.Dequeue();
}
}
}
protected async Task <bool> TryMoveNextProducerAsync(
DocumentProducer <T> producer,
Func <DocumentProducer <T>, Task <DocumentProducer <T> > > producerRepairCallback,
CancellationToken cancellationToken)
{
bool movedNext = false;
DocumentProducer <T> currentProducer = producer;
while (true)
{
bool needRefreshedPartitionKeyRangeCache = false;
try
{
movedNext = await currentProducer.MoveNextAsync(cancellationToken);
}
catch (DocumentClientException ex)
{
if (!(needRefreshedPartitionKeyRangeCache = base.NeedPartitionKeyRangeCacheRefresh(ex)))
{
throw;
}
}
if (needRefreshedPartitionKeyRangeCache)
{
currentProducer = await producerRepairCallback(currentProducer);
}
else
{
break;
}
}
if (!movedNext)
{
this.CurrentContinuationTokens.Remove(currentProducer);
}
return(movedNext);
}
protected async Task RepairContextAsync(
string collectionRid,
int currentDocumentProducerIndex,
Func <DocumentProducer <T>, int> taskPriorityFunc,
IReadOnlyList <PartitionKeyRange> replacementRanges,
SqlQuerySpec querySpecForRepair,
Action callback = null)
{
CollectionCache collectionCache = await this.Client.GetCollectionCacheAsync();
INameValueCollection requestHeaders = await this.CreateCommonHeadersAsync(this.GetFeedOptions(null));
this.DocumentProducers.Capacity = this.DocumentProducers.Count + replacementRanges.Count - 1;
DocumentProducer <T> replacedDocumentProducer = this.DocumentProducers[currentDocumentProducerIndex];
DefaultTrace.TraceInformation(string.Format(
CultureInfo.InvariantCulture,
"{0}, CorrelatedActivityId: {5} | Parallel~ContextBase.RepairContextAsync, MaxBufferedItemCount: {1}, Replacement PartitionKeyRange Count: {2}, MaximumConcurrencyLevel: {3}, DocumentProducer Initial Page Size {4}",
DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture),
this.actualMaxBufferedItemCount,
replacementRanges.Count,
this.TaskScheduler.MaximumConcurrencyLevel,
replacedDocumentProducer.PageSize,
this.CorrelatedActivityId));
int index = currentDocumentProducerIndex + 1;
foreach (PartitionKeyRange range in replacementRanges)
{
this.DocumentProducers.Insert(
index++,
new DocumentProducer <T>(
this.TaskScheduler,
(continuationToken, pageSize) =>
{
INameValueCollection headers = requestHeaders.Clone();
headers[HttpConstants.HttpHeaders.Continuation] = continuationToken;
headers[HttpConstants.HttpHeaders.PageSize] = pageSize.ToString(CultureInfo.InvariantCulture);
return(this.CreateDocumentServiceRequest(
headers,
querySpecForRepair,
range,
collectionRid));
},
range,
taskPriorityFunc,
this.ExecuteRequestAsync <T>,
() => new NonRetriableInvalidPartitionExceptionRetryPolicy(collectionCache, this.Client.RetryPolicy.GetRequestPolicy()),
this.OnDocumentProducerCompleteFetching,
this.CorrelatedActivityId,
replacedDocumentProducer.PageSize,
replacedDocumentProducer.CurrentBackendContinuationToken));
}
this.DocumentProducers.RemoveAt(currentDocumentProducerIndex);
if (callback != null)
{
callback();
}
if (this.ShouldPrefetch)
{
for (int i = 0; i < replacementRanges.Count; i++)
{
this.DocumentProducers[i + currentDocumentProducerIndex].TryScheduleFetch();
}
}
if (this.CurrentContinuationTokens.Remove(replacedDocumentProducer))
{
for (int i = 0; i < replacementRanges.Count; ++i)
{
this.CurrentContinuationTokens[this.DocumentProducers[currentDocumentProducerIndex + i]] = replacedDocumentProducer.CurrentBackendContinuationToken;
}
}
}
/// <summary>
/// Initializes a new instance of the DocumentProducerTree class.
/// </summary>
/// <param name="partitionKeyRange">The partition key range.</param>
/// <param name="createRequestFunc">Callback to create a request.</param>
/// <param name="executeRequestFunc">Callback to execute a request.</param>
/// <param name="createRetryPolicyFunc">Callback to create a retry policy.</param>
/// <param name="produceAsyncCompleteCallback">Callback to invoke once a fetch finishes.</param>
/// <param name="documentProducerTreeComparer">Comparer to determine, which tree to produce from.</param>
/// <param name="equalityComparer">Comparer to see if we need to return the continuation token for a partition.</param>
/// <param name="client">The client</param>
/// <param name="deferFirstPage">Whether or not to defer fetching the first page.</param>
/// <param name="collectionRid">The collection to drain from.</param>
/// <param name="initialPageSize">The initial page size.</param>
/// <param name="initialContinuationToken">The initial continuation token.</param>
public DocumentProducerTree(
PartitionKeyRange partitionKeyRange,
Func <PartitionKeyRange, string, int, DocumentServiceRequest> createRequestFunc,
Func <DocumentServiceRequest, CancellationToken, Task <FeedResponse <CosmosElement> > > executeRequestFunc,
Func <IDocumentClientRetryPolicy> createRetryPolicyFunc,
Action <DocumentProducerTree, int, double, QueryMetrics, long, CancellationToken> produceAsyncCompleteCallback,
IComparer <DocumentProducerTree> documentProducerTreeComparer,
IEqualityComparer <CosmosElement> equalityComparer,
IDocumentQueryClient client,
bool deferFirstPage,
string collectionRid,
long initialPageSize = 50,
string initialContinuationToken = null)
{
if (documentProducerTreeComparer == null)
{
throw new ArgumentNullException($"{nameof(documentProducerTreeComparer)}");
}
if (createRequestFunc == null)
{
throw new ArgumentNullException($"{nameof(createRequestFunc)}");
}
if (executeRequestFunc == null)
{
throw new ArgumentNullException($"{nameof(executeRequestFunc)}");
}
if (createRetryPolicyFunc == null)
{
throw new ArgumentNullException($"{nameof(createRetryPolicyFunc)}");
}
if (produceAsyncCompleteCallback == null)
{
throw new ArgumentNullException($"{nameof(produceAsyncCompleteCallback)}");
}
if (documentProducerTreeComparer == null)
{
throw new ArgumentNullException($"{nameof(documentProducerTreeComparer)}");
}
if (equalityComparer == null)
{
throw new ArgumentNullException($"{nameof(equalityComparer)}");
}
if (client == null)
{
throw new ArgumentNullException($"{nameof(client)}");
}
if (string.IsNullOrEmpty(collectionRid))
{
throw new ArgumentException($"{nameof(collectionRid)} can not be null or empty.");
}
this.root = new DocumentProducer(
partitionKeyRange,
createRequestFunc,
executeRequestFunc,
createRetryPolicyFunc,
(documentProducer, itemsBuffered, resourceUnitUsage, queryMetrics, requestLength, token) => produceAsyncCompleteCallback(this, itemsBuffered, resourceUnitUsage, queryMetrics, requestLength, token),
equalityComparer,
initialPageSize,
initialContinuationToken);
this.children = new PriorityQueue <DocumentProducerTree>(documentProducerTreeComparer, true);
this.deferFirstPage = deferFirstPage;
this.client = client;
this.collectionRid = collectionRid;
this.createDocumentProducerTreeCallback = DocumentProducerTree.CreateDocumentProducerTreeCallback(
createRequestFunc,
executeRequestFunc,
createRetryPolicyFunc,
produceAsyncCompleteCallback,
documentProducerTreeComparer,
equalityComparer,
client,
deferFirstPage,
collectionRid,
initialPageSize);
this.executeWithSplitProofingSemaphore = new SemaphoreSlim(1, 1);
}
private OrderByContinuationToken CreateCrossPartitionOrderByContinuationToken(DocumentProducer <OrderByQueryResult> documentProducer)
{
OrderByQueryResult orderByResult = documentProducer.Current;
string filter = null;
this.filters.TryGetValue(documentProducer.TargetRange.Id, out filter);
return(new OrderByContinuationToken
{
CompositeToken = new CompositeContinuationToken
{
Token = documentProducer.PreviousResponseContinuation,
Range = documentProducer.TargetRange.ToRange(),
},
OrderByItems = orderByResult.OrderByItems,
Rid = orderByResult.Rid,
SkipCount = this.ShouldIncrementSkipCount(orderByResult.Rid) ? this.skipCount + 1 : 0,
Filter = filter
});
}
private async Task FilterAsync(
DocumentProducer <OrderByQueryResult> producer,
SortOrder[] sortOrders,
OrderByContinuationToken continuationToken,
CancellationToken cancellationToken)
{
ResourceId continuationRid;
if (!ResourceId.TryParse(continuationToken.Rid, out continuationRid))
{
DefaultTrace.TraceWarning(
string.Format(
CultureInfo.InvariantCulture,
"{0}, CorrelatedActivityId: {1}, ActivityId: {2} | Invalid Rid in the continuation token {3} for OrderBy~Context.",
DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture),
this.CorrelatedActivityId,
this.ActivityId,
continuationToken.CompositeToken.Token));
throw new BadRequestException(RMResources.InvalidContinuationToken);
}
Dictionary <string, ResourceId> resourceIds = new Dictionary <string, ResourceId>();
int itemToSkip = continuationToken.SkipCount;
bool continuationRidVerified = false;
while (true)
{
OrderByQueryResult orderByResult = (OrderByQueryResult)producer.Current;
int cmp = 0;
for (int i = 0; i < sortOrders.Length; ++i)
{
cmp = ItemComparer.Instance.Compare(
continuationToken.OrderByItems[i].GetItem(),
orderByResult.OrderByItems[i].GetItem());
if (cmp != 0)
{
cmp = sortOrders[i] != SortOrder.Descending ? cmp : -cmp;
break;
}
}
if (cmp < 0)
{
break;
}
if (cmp == 0)
{
ResourceId rid;
if (!resourceIds.TryGetValue(orderByResult.Rid, out rid))
{
if (!ResourceId.TryParse(orderByResult.Rid, out rid))
{
DefaultTrace.TraceWarning(
string.Format(
CultureInfo.InvariantCulture,
"{0}, CorrelatedActivityId: {1}, ActivityId: {2} | Invalid Rid in the continuation token {3} for OrderBy~Context.",
DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture),
this.CorrelatedActivityId,
this.ActivityId,
continuationToken.CompositeToken.Token));
throw new BadRequestException(RMResources.InvalidContinuationToken);
}
resourceIds.Add(orderByResult.Rid, rid);
}
if (!continuationRidVerified)
{
if (continuationRid.Database != rid.Database || continuationRid.DocumentCollection != rid.DocumentCollection)
{
DefaultTrace.TraceWarning(
string.Format(
CultureInfo.InvariantCulture,
"{0}, CorrelatedActivityId: {1}, ActivityId: {2} | Invalid Rid in the continuation token {3} for OrderBy~Context.",
DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture),
this.CorrelatedActivityId,
this.ActivityId,
continuationToken.CompositeToken.Token));
throw new BadRequestException(RMResources.InvalidContinuationToken);
}
continuationRidVerified = true;
}
cmp = continuationRid.Document.CompareTo(rid.Document);
if (sortOrders[sortOrders.Length - 1] == SortOrder.Descending)
{
cmp = -cmp;
}
if (cmp < 0 || (cmp == 0 && itemToSkip-- <= 0))
{
break;
}
}
if (!await producer.MoveNextAsync(cancellationToken))
{
break;
}
}
}
private async Task InitializeAsync(
string collectionRid,
List <Range <string> > queryRanges,
List <PartitionKeyRange> partitionKeyRanges,
SortOrder[] sortOrders,
string[] orderByExpressions,
int initialPageSize,
string requestContinuation,
CancellationToken cancellationToken)
{
this.InitializationSchedulingMetrics.Start();
try
{
OrderByContinuationToken[] suppliedContinuationTokens = this.ValidateAndExtractContinuationTokens(requestContinuation, sortOrders, orderByExpressions);
Dictionary <string, OrderByContinuationToken> targetRangeToOrderByContinuationMap = null;
base.DocumentProducers.Capacity = partitionKeyRanges.Count;
if (suppliedContinuationTokens == null)
{
await base.InitializeAsync(
collectionRid,
queryRanges,
TaskPriorityFunc,
partitionKeyRanges,
initialPageSize,
new SqlQuerySpec(this.QuerySpec.QueryText.Replace(FormatPlaceHolder, True), this.QuerySpec.Parameters),
null,
cancellationToken);
}
else
{
RangeFilterInitializationInfo[] orderByInfos = this.GetPartitionKeyRangesInitializationInfo(
suppliedContinuationTokens,
partitionKeyRanges,
sortOrders,
orderByExpressions,
out targetRangeToOrderByContinuationMap);
Debug.Assert(targetRangeToOrderByContinuationMap != null, "If targetRangeToOrderByContinuationMap can't be null is valid continuation is supplied");
this.currentOrderByItems = suppliedContinuationTokens[0].OrderByItems;
// For ascending order-by, left of target partition has filter expression > value,
// right of target partition has filter expression >= value,
// and target partition takes the previous filter from continuation (or true if no continuation)
foreach (RangeFilterInitializationInfo info in orderByInfos)
{
if (info.StartIndex > info.EndIndex)
{
continue;
}
PartialReadOnlyList <PartitionKeyRange> partialRanges =
new PartialReadOnlyList <PartitionKeyRange>(partitionKeyRanges, info.StartIndex, info.EndIndex - info.StartIndex + 1);
Task initTask = base.InitializeAsync(
collectionRid,
queryRanges,
TaskPriorityFunc,
partialRanges,
initialPageSize,
new SqlQuerySpec(
this.QuerySpec.QueryText.Replace(FormatPlaceHolder, info.Filter),
this.QuerySpec.Parameters),
targetRangeToOrderByContinuationMap.ToDictionary(kvp => kvp.Key, kvp => kvp.Value.CompositeToken.Token),
cancellationToken);
foreach (PartitionKeyRange range in partialRanges)
{
this.filters[range.Id] = info.Filter;
}
await initTask;
}
}
// The Foreach loop below is an optimization for the following While loop. The While loop is made single-threaded as the base.DocumentProducers object can change during runtime due to split. Even though the While loop is single threaded, the Foreach loop below makes the producers fetch doucments concurrently. If any of the producers fails to produce due to split (i.e., encounters PartitionKeyRangeGoneException), then the while loop below will take out the failed document producers and replace it approprite ones and then call TryScheduleFetch() on them.
foreach (var producer in base.DocumentProducers)
{
producer.TryScheduleFetch();
}
// Fetch one item from each of the producers to initialize the priority-queue. "TryMoveNextProducerAsync()" has
// a side-effect that, if Split is encountered while trying to move, related parent producer will be taken out and child
// producers will be added to "base.DocumentProducers".
for (int index = 0; index < base.DocumentProducers.Count; ++index)
{
DocumentProducer <OrderByQueryResult> producer = base.DocumentProducers[index];
if (await this.TryMoveNextProducerAsync(
producer,
targetRangeToOrderByContinuationMap,
cancellationToken: cancellationToken))
{
producer = base.DocumentProducers[index];
OrderByContinuationToken continuationToken =
(targetRangeToOrderByContinuationMap != null && targetRangeToOrderByContinuationMap.ContainsKey(producer.TargetRange.Id)) ?
targetRangeToOrderByContinuationMap[producer.TargetRange.Id] : null;
if (continuationToken != null)
{
await this.FilterAsync(producer, sortOrders, continuationToken, cancellationToken);
}
this.documentProducerConsumerQueue.Enqueue(producer);
}
}
}
finally
{
this.InitializationSchedulingMetrics.Stop();
DefaultTrace.TraceInformation(string.Format(
CultureInfo.InvariantCulture,
"{0}, CorrelatedActivityId: {1} | OrderBy~Context.InitializeAsync",
DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture),
this.CorrelatedActivityId));
}
}
private void OnDocumentProducerCompleteFetching(
DocumentProducer <T> producer,
int size,
double resourceUnitUsage,
QueryMetrics queryMetrics,
long responseLengthBytes,
CancellationToken token)
{
// Update charge and states
this.chargeTracker.AddCharge(resourceUnitUsage);
Interlocked.Add(ref this.totalBufferedItems, size);
Interlocked.Increment(ref this.totalRequestRoundTrips);
this.IncrementResponseLengthBytes(responseLengthBytes);
this.partitionedQueryMetrics.Add(Tuple.Create(producer.TargetRange.Id, queryMetrics));
//Check to see if we can buffer more item
long countToAdd = size - this.FreeItemSpace;
if (countToAdd > 0 &&
this.actualMaxBufferedItemCount < MaxixmumDynamicMaxBufferedItemCountValue - countToAdd)
{
DefaultTrace.TraceVerbose(string.Format(
CultureInfo.InvariantCulture,
"{0}, CorrelatedActivityId: {4} | Id: {1}, increasing MaxBufferedItemCount {2} by {3}.",
DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture),
producer.TargetRange.Id,
this.actualMaxBufferedItemCount,
countToAdd,
this.CorrelatedActivityId));
countToAdd += this.actualMaxBufferedItemCount;
}
// Adjust max DoP if necessary
this.AdjustTaskSchedulerMaximumConcurrencyLevel();
// Fetch again if necessary
if (!producer.FetchedAll)
{
if (producer.PageSize < this.actualMaxPageSize)
{
producer.PageSize = Math.Min((long)(producer.PageSize * DynamicPageSizeAdjustmentFactor), this.actualMaxPageSize);
Debug.Assert(producer.PageSize >= 0 && producer.PageSize <= int.MaxValue, string.Format("producer.PageSize is invalid at {0}", producer.PageSize));
}
if (this.ShouldPrefetch &&
this.FreeItemSpace - producer.NormalizedPageSize > 0)
{
producer.TryScheduleFetch();
}
}
DefaultTrace.TraceVerbose(string.Format(
CultureInfo.InvariantCulture,
"{0}, CorrelatedActivityId: {5} | Id: {1}, size: {2}, resourceUnitUsage: {3}, taskScheduler.CurrentRunningTaskCount: {4}",
DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture),
producer.TargetRange.Id,
size,
resourceUnitUsage,
this.TaskScheduler.CurrentRunningTaskCount,
this.CorrelatedActivityId));
}
public async Task ConcurrentMoveNextTryScheduleTestAsync()
{
int seed = (int)(DateTime.UtcNow - new DateTime(1970, 1, 1)).TotalSeconds;
Random rand = new Random(seed);
int maxValue = 100;
int trials = 1000;
int maxTicks = 100;
IEnumerable <int> expectedValues = Enumerable.Range(1, maxValue);
IDocumentClientRetryPolicy retryPolicy = new MockRetryPolicy(rand);
ComparableTaskScheduler scheduler = new ComparableTaskScheduler(1);
for (int trial = 0; trial < trials; ++trial)
{
DocumentProducer <int> producer = new DocumentProducer <int>(
scheduler,
(continuation, pageSize) => DocumentServiceRequest.Create(
OperationType.Query,
"/dbs/db/colls/coll",
ResourceType.Document,
new MemoryStream(Encoding.UTF8.GetBytes(continuation)),
AuthorizationTokenType.PrimaryMasterKey,
new StringKeyValueCollection
{
{ HttpConstants.HttpHeaders.Continuation, continuation }
}),
new PartitionKeyRange {
Id = "test", MinInclusive = "", MaxExclusive = "ff"
},
p => 0,
(request, token) =>
{
if (rand.Next(4) == 0)
{
throw new Exception();
}
if (rand.Next(10) == 0)
{
return(Task.FromResult(new FeedResponse <int>(new int[] { }, 0, request.Headers)));
}
using (StreamReader reader = new StreamReader(request.Body))
{
int value = int.Parse(reader.ReadToEnd()) + 1;
INameValueCollection headers = new StringKeyValueCollection
{
{ HttpConstants.HttpHeaders.Continuation, value >= maxValue? null : value.ToString(CultureInfo.InvariantCulture) }
};
return(Task.FromResult(new FeedResponse <int>(new int[] { value }, 1, headers)));
}
},
() => retryPolicy,
(produer, size, ru, queryMetrics, token, length) => { },
Guid.NewGuid(),
1000,
"0");
Timer timer = new Timer(
(state) => producer.TryScheduleFetch(TimeSpan.FromTicks(rand.Next(maxTicks))),
null,
TimeSpan.FromTicks(rand.Next(maxTicks)),
TimeSpan.FromTicks(rand.Next(maxTicks)));
List <int> actualValues = new List <int>();
CancellationTokenSource tokenSource = new CancellationTokenSource(5000);
while (await producer.MoveNextAsync(tokenSource.Token))
{
actualValues.Add(producer.Current);
}
Assert.AreEqual(
string.Join(", ", expectedValues),
string.Join(", ", actualValues),
string.Format(CultureInfo.InvariantCulture, "seed: {0}", seed));
}
}
