/// <summary>
/// Drains documents from this execution context.
/// </summary>
/// <param name="maxElements">The maximum number of documents to drains.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A task that when awaited on returns a DoucmentFeedResponse of results.</returns>
public override async Task <IReadOnlyList <CosmosElement> > InternalDrainAsync(int maxElements, CancellationToken cancellationToken)
{
// In order to maintain the continuation token for the user we must drain with a few constraints
// 1) We fully drain from the left most partition before moving on to the next partition
// 2) We drain only full pages from the document producer so we aren't left with a partial page
// otherwise we would need to add to the continuation token how many items to skip over on that page.
// Only drain from the leftmost (current) document producer tree
ItemProducerTree currentItemProducerTree = this.PopCurrentItemProducerTree();
// This might be the first time we have seen this document producer tree so we need to buffer documents
if (currentItemProducerTree.Current == null)
{
await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken);
}
int itemsLeftInCurrentPage = currentItemProducerTree.ItemsLeftInCurrentPage;
// Only drain full pages or less if this is a top query.
List <CosmosElement> results = new List <CosmosElement>();
for (int i = 0; i < Math.Min(itemsLeftInCurrentPage, maxElements); i++)
{
results.Add(currentItemProducerTree.Current);
if (await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken))
{
break;
}
}
this.PushCurrentItemProducerTree(currentItemProducerTree);
// At this point the document producer tree should have internally called MoveNextPage, since we fully drained a page.
return(results);
}
/// <summary>
/// Function that is given to all the document producers to call on once they are done fetching.
/// This is so that the CosmosCrossPartitionQueryExecutionContext can aggregate metadata from them.
/// </summary>
/// <param name="producer">The document producer that just finished fetching.</param>
/// <param name="itemsBuffered">The number of items that the producer just fetched.</param>
/// <param name="resourceUnitUsage">The amount of RUs that the producer just consumed.</param>
/// <param name="queryMetrics">The query metrics that the producer just got back from the backend.</param>
/// <param name="responseLengthBytes">The length of the response the producer just got back in bytes.</param>
/// <param name="token">The cancellation token.</param>
/// <remarks>
/// This function is by nature a bit racy.
/// A query might be fully drained but a background task is still fetching documents so this will get called after the context is done.
/// </remarks>
private void OnItemProducerTreeCompleteFetching(
ItemProducerTree producer,
int itemsBuffered,
double resourceUnitUsage,
QueryMetrics queryMetrics,
long responseLengthBytes,
CancellationToken token)
{
// Update charge and states
this.requestChargeTracker.AddCharge(resourceUnitUsage);
Interlocked.Add(ref this.totalBufferedItems, itemsBuffered);
this.IncrementResponseLengthBytes(responseLengthBytes);
this.partitionedQueryMetrics.Add(Tuple.Create(producer.PartitionKeyRange.Id, queryMetrics));
// Adjust the producer page size so that we reach the optimal page size.
producer.PageSize = Math.Min((long)(producer.PageSize * DynamicPageSizeAdjustmentFactor), this.actualMaxPageSize);
// Adjust Max Degree Of Parallelism if necessary
// (needs to wait for comparable task scheduler refactor).
// Fetch again if necessary
if (producer.HasMoreBackendResults)
{
// 4mb is the max response size
long expectedResponseSize = Math.Min(producer.PageSize, 4 * 1024 * 1024);
if (this.CanPrefetch && this.FreeItemSpace > expectedResponseSize)
{
this.TryScheduleFetch(producer);
}
}
}
/// <summary>
/// Initializes a new instance of the ItemProducerTreeComparableTask class.
/// </summary>
/// <param name="producer">The producer to fetch from.</param>
/// <param name="taskPriorityFunction">The callback to determine the fetch priority of the document producer.</param>
public ItemProducerTreeComparableTask(
ItemProducerTree producer,
Func <ItemProducerTree, int> taskPriorityFunction)
: base(taskPriorityFunction(producer))
{
this.producer = producer;
}
/// <summary>
/// Tries to schedule a fetch from the document producer tree.
/// </summary>
/// <param name="itemProducerTree">The document producer tree to schedule a fetch for.</param>
/// <returns>Whether or not the fetch was successfully scheduled.</returns>
private bool TryScheduleFetch(ItemProducerTree itemProducerTree)
{
return(this.comparableTaskScheduler.TryQueueTask(
new ItemProducerTreeComparableTask(
itemProducerTree,
this.fetchPrioirtyFunction),
default(TimeSpan)));
}
/// <summary>
/// A helper to move next and set the failure response if one is received
/// </summary>
/// <param name="itemProducerTree">The item producer tree</param>
/// <param name="cancellationToken">The cancellation token</param>
/// <returns>True if it move next failed. It can fail from an error or hitting the end of the tree</returns>
protected async Task<bool> MoveNextHelperAsync(ItemProducerTree itemProducerTree, CancellationToken cancellationToken)
{
(bool successfullyMovedNext, QueryResponseCore? failureResponse) moveNextResponse = await itemProducerTree.MoveNextAsync(cancellationToken);
if (moveNextResponse.failureResponse != null)
{
this.FailureResponse = moveNextResponse.failureResponse;
}
return moveNextResponse.successfullyMovedNext;
}
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
// In order to maintain the continuation token for the user we must drain with a few constraints
// 1) We fully drain from the left most partition before moving on to the next partition
// 2) We drain only full pages from the document producer so we aren't left with a partial page
// otherwise we would need to add to the continuation token how many items to skip over on that page.
// Only drain from the leftmost (current) document producer tree
ItemProducerTree currentItemProducerTree = this.PopCurrentItemProducerTree();
List <CosmosElement> results = new List <CosmosElement>();
try
{
(bool gotNextPage, QueryResponseCore? failureResponse) = await currentItemProducerTree.TryMoveNextPageAsync(cancellationToken);
if (failureResponse != null)
{
return(failureResponse.Value);
}
if (gotNextPage)
{
int itemsLeftInCurrentPage = currentItemProducerTree.ItemsLeftInCurrentPage;
// Only drain full pages or less if this is a top query.
currentItemProducerTree.TryMoveNextDocumentWithinPage();
int numberOfItemsToDrain = Math.Min(itemsLeftInCurrentPage, maxElements);
for (int i = 0; i < numberOfItemsToDrain; i++)
{
results.Add(currentItemProducerTree.Current);
currentItemProducerTree.TryMoveNextDocumentWithinPage();
}
}
}
finally
{
this.PushCurrentItemProducerTree(currentItemProducerTree);
}
return(QueryResponseCore.CreateSuccess(
result: results,
requestCharge: this.requestChargeTracker.GetAndResetCharge(),
activityId: null,
responseLengthBytes: this.GetAndResetResponseLengthBytes(),
disallowContinuationTokenMessage: null,
continuationToken: this.ContinuationToken,
diagnostics: this.GetAndResetDiagnostics()));
}
/// <summary>
/// After a split you need to maintain the continuation tokens for all the child document producers until a condition is met.
/// For example lets say that a document producer is at continuation X and it gets split,
/// then the children each get continuation X, but since you only drain from one of them at a time you are left with the first child having
/// continuation X + delta and the second child having continuation X (draw this out if you are following along).
/// At this point you have the answer the question: "Which continuation token do you return to the user?".
/// Let's say you return X, then when you come back to the first child you will be repeating work, thus returning some documents more than once.
/// Let's say you return X + delta, then you fine when you return to the first child, but when you get to the second child you don't have a continuation token
/// meaning that you will be repeating all the document for the second partition up until X and again you will be returning some documents more than once.
/// Thus you have to return the continuation token for both children.
/// Both this means you are returning more than 1 continuation token for the rest of the query.
/// Well a naive optimization is to flush the continuation for a child partition once you are done draining from it, which isn't bad for a parallel query,
/// but if you have an order by query you might not be done with a producer until the end of the query.
/// The next optimization for a parallel query is to flush the continuation token the moment you start reading from a child partition.
/// This works for a parallel query, but breaks for an order by query.
/// The final realization is that for an order by query you are only choosing between multiple child partitions when their is a tie,
/// so the key is that you can dump the continuation token the moment you come across a new order by item.
/// For order by queries that is determined by the order by field and for parallel queries that is the moment you come by a new rid (which is any document, since rids are unique within a partition).
/// So by passing an equality comparer to the document producers they can determine whether they are still "active".
/// </summary>
/// <returns>
/// Returns all document producers whose continuation token you have to return.
/// Only during a split will this list contain more than 1 item.
/// </returns>
public IEnumerable <ItemProducer> GetActiveItemProducers()
{
ItemProducerTree current = this.itemProducerForest.Peek().CurrentItemProducerTree;
if (current.HasMoreResults && !current.IsActive)
{
// If the current document producer tree has more results, but isn't active.
// then we still want to emit it, since it won't get picked up in the below for loop.
yield return(current.Root);
}
foreach (ItemProducerTree itemProducerTree in this.itemProducerForest)
{
foreach (ItemProducer itemProducer in itemProducerTree.GetActiveItemProducers())
{
yield return(itemProducer);
}
}
}
/// <summary>
/// Function that is given to all the document producers to call on once they are done fetching.
/// This is so that the CosmosCrossPartitionQueryExecutionContext can aggregate metadata from them.
/// </summary>
/// <param name="producer">The document producer that just finished fetching.</param>
/// <param name="itemsBuffered">The number of items that the producer just fetched.</param>
/// <param name="resourceUnitUsage">The amount of RUs that the producer just consumed.</param>
/// <param name="queryMetrics">The query metrics that the producer just got back from the backend.</param>
/// <param name="responseLengthBytes">The length of the response the producer just got back in bytes.</param>
/// <param name="token">The cancellation token.</param>
/// <remarks>
/// This function is by nature a bit racy.
/// A query might be fully drained but a background task is still fetching documents so this will get called after the context is done.
/// </remarks>
private void OnItemProducerTreeCompleteFetching(
ItemProducerTree producer,
int itemsBuffered,
double resourceUnitUsage,
QueryMetrics queryMetrics,
long responseLengthBytes,
CancellationToken token)
{
// Update charge and states
this.requestChargeTracker.AddCharge(resourceUnitUsage);
Interlocked.Add(ref this.totalBufferedItems, itemsBuffered);
this.IncrementResponseLengthBytes(responseLengthBytes);
this.partitionedQueryMetrics.Add(Tuple.Create(producer.PartitionKeyRange.Id, queryMetrics));
// Adjust the producer page size so that we reach the optimal page size.
producer.PageSize = Math.Min((long)(producer.PageSize * DynamicPageSizeAdjustmentFactor), this.actualMaxPageSize);
// Adjust Max Degree Of Paralleism if neccesary
// (needs to wait for comparable task scheudler refactor).
// Fetch again if necessary
if (producer.HasMoreBackendResults)
{
// 4mb is the max reponse size
long expectedResponseSize = Math.Min(producer.PageSize, 4 * 1024 * 1024);
if (this.CanPrefetch && this.FreeItemSpace > expectedResponseSize)
{
this.TryScheduleFetch(producer);
}
}
this.TraceVerbose(string.Format(
CultureInfo.InvariantCulture,
"Id: {0}, size: {1}, resourceUnitUsage: {2}, taskScheduler.CurrentRunningTaskCount: {3}",
producer.PartitionKeyRange.Id,
itemsBuffered,
resourceUnitUsage,
this.comparableTaskScheduler.CurrentRunningTaskCount,
this.queryContext.CorrelatedActivityId));
}
/// <summary>
/// Function that is given to all the document producers to call on once they are done fetching.
/// This is so that the CosmosCrossPartitionQueryExecutionContext can aggregate metadata from them.
/// </summary>
/// <param name="producer">The document producer that just finished fetching.</param>
/// <param name="itemsBuffered">The number of items that the producer just fetched.</param>
/// <param name="resourceUnitUsage">The amount of RUs that the producer just consumed.</param>
/// <param name="diagnostics">The query metrics that the producer just got back from the backend.</param>
/// <param name="responseLengthBytes">The length of the response the producer just got back in bytes.</param>
/// <param name="token">The cancellation token.</param>
/// <remarks>
/// This function is by nature a bit racy.
/// A query might be fully drained but a background task is still fetching documents so this will get called after the context is done.
/// </remarks>
private void OnItemProducerTreeCompleteFetching(
ItemProducerTree producer,
int itemsBuffered,
double resourceUnitUsage,
IReadOnlyCollection <QueryPageDiagnostics> diagnostics,
long responseLengthBytes,
CancellationToken token)
{
// Update charge and states
this.requestChargeTracker.AddCharge(resourceUnitUsage);
Interlocked.Add(ref this.totalBufferedItems, itemsBuffered);
this.IncrementResponseLengthBytes(responseLengthBytes);
// Add the pages to the concurrent bag to safely merge all the list together.
foreach (QueryPageDiagnostics diagnosticPage in diagnostics)
{
this.diagnosticsPages.Add(diagnosticPage);
}
// Adjust the producer page size so that we reach the optimal page size.
producer.PageSize = Math.Min((long)(producer.PageSize * DynamicPageSizeAdjustmentFactor), this.actualMaxPageSize);
// Adjust Max Degree Of Parallelism if necessary
// (needs to wait for comparable task scheduler refactor).
// Fetch again if necessary
if (producer.HasMoreBackendResults)
{
// 4mb is the max response size
long expectedResponseSize = Math.Min(producer.PageSize, 4 * 1024 * 1024);
if (this.CanPrefetch && this.FreeItemSpace > expectedResponseSize)
{
this.TryScheduleFetch(producer);
}
}
}
/// <summary>
/// Drains a page of documents from this context.
/// </summary>
/// <param name="maxElements">The maximum number of elements.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A task that when awaited on return a page of documents.</returns>
public override async Task <IList <CosmosElement> > InternalDrainAsync(int maxElements, CancellationToken cancellationToken)
{
//// In order to maintain the continuation token for the user we must drain with a few constraints
//// 1) We always drain from the partition, which has the highest priority item first
//// 2) If multiple partitions have the same priority item then we drain from the left most first
//// otherwise we would need to keep track of how many of each item we drained from each partition
//// (just like parallel queries).
//// Visually that look the following case where we have three partitions that are numbered and store letters.
//// For teaching purposes I have made each item a tuple of the following form:
//// <item stored in partition, partition number>
//// So that duplicates across partitions are distinct, but duplicates within partitions are indistinguishable.
//// |-------| |-------| |-------|
//// | <a,1> | | <a,2> | | <a,3> |
//// | <a,1> | | <b,2> | | <c,3> |
//// | <a,1> | | <b,2> | | <c,3> |
//// | <d,1> | | <c,2> | | <c,3> |
//// | <d,1> | | <e,2> | | <f,3> |
//// | <e,1> | | <h,2> | | <j,3> |
//// | <f,1> | | <i,2> | | <k,3> |
//// |-------| |-------| |-------|
//// Now the correct drain order in this case is:
//// <a,1>,<a,1>,<a,1>,<a,2>,<a,3>,<b,2>,<b,2>,<c,2>,<c,3>,<c,3>,<c,3>,
//// <d,1>,<d,1>,<e,1>,<e,2>,<f,1>,<f,3>,<h,2>,<i,2>,<j,3>,<k,3>
//// In more mathematical terms
//// 1) <x, y> always comes before <z, y> where x < z
//// 2) <i, j> always come before <i, k> where j < k
List <CosmosElement> results = new List <CosmosElement>();
while (!this.IsDone && results.Count < maxElements)
{
// Only drain from the highest priority document producer
// We need to pop and push back the document producer tree, since the priority changes according to the sort order.
ItemProducerTree currentItemProducerTree = this.PopCurrentItemProducerTree();
OrderByQueryResult orderByQueryResult = new OrderByQueryResult(currentItemProducerTree.Current);
// Only add the payload, since other stuff is garbage from the caller's perspective.
results.Add(orderByQueryResult.Payload);
// If we are at the beginning of the page and seeing an rid from the previous page we should increment the skip count
// due to the fact that JOINs can make a document appear multiple times and across continuations, so we don't want to
// surface this more than needed. More information can be found in the continuation token docs.
if (this.ShouldIncrementSkipCount(currentItemProducerTree.CurrentItemProducerTree.Root))
{
++this.skipCount;
}
else
{
this.skipCount = 0;
}
this.previousRid = orderByQueryResult.Rid;
if (await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken))
{
break;
}
this.PushCurrentItemProducerTree(currentItemProducerTree);
}
return(results);
}
/// <summary>
/// Initializes cross partition query execution context by initializing the necessary document producers.
/// </summary>
/// <param name="collectionRid">The collection to drain from.</param>
/// <param name="partitionKeyRanges">The partitions to target.</param>
/// <param name="initialPageSize">The page size to start the document producers off with.</param>
/// <param name="querySpecForInit">The query specification for the rewritten query.</param>
/// <param name="targetRangeToContinuationMap">Map from partition to it's corresponding continuation token.</param>
/// <param name="deferFirstPage">Whether or not we should defer the fetch of the first page from each partition.</param>
/// <param name="filter">The filter to inject in the predicate.</param>
/// <param name="filterCallback">The callback used to filter each partition.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>A task to await on.</returns>
protected async Task InitializeAsync(
string collectionRid,
IReadOnlyList <PartitionKeyRange> partitionKeyRanges,
int initialPageSize,
SqlQuerySpec querySpecForInit,
Dictionary <string, string> targetRangeToContinuationMap,
bool deferFirstPage,
string filter,
Func <ItemProducerTree, Task> filterCallback,
CancellationToken token)
{
List <ItemProducerTree> itemProducerTrees = new List <ItemProducerTree>();
foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges)
{
string initialContinuationToken = (targetRangeToContinuationMap != null && targetRangeToContinuationMap.ContainsKey(partitionKeyRange.Id)) ? targetRangeToContinuationMap[partitionKeyRange.Id] : null;
ItemProducerTree itemProducerTree = new ItemProducerTree(
this.queryContext,
querySpecForInit,
partitionKeyRange,
this.OnItemProducerTreeCompleteFetching,
this.itemProducerForest.Comparer as IComparer <ItemProducerTree>,
this.equalityComparer,
deferFirstPage,
collectionRid,
initialPageSize,
initialContinuationToken)
{
Filter = filter
};
// Prefetch if necessary, and populate consume queue.
if (this.CanPrefetch)
{
this.TryScheduleFetch(itemProducerTree);
}
itemProducerTrees.Add(itemProducerTree);
}
// Using loop fission so that we can load the document producers in parallel
foreach (ItemProducerTree itemProducerTree in itemProducerTrees)
{
if (!deferFirstPage)
{
(bool successfullyMovedNext, QueryResponseCore? failureResponse)response = await itemProducerTree.MoveNextIfNotSplitAsync(token);
if (response.failureResponse != null)
{
// Set the failure so on drain it can be returned.
this.FailureResponse = response.failureResponse;
// No reason to enqueue the rest of the itemProducerTrees since there is a failure.
break;
}
}
if (filterCallback != null)
{
await filterCallback(itemProducerTree);
}
if (itemProducerTree.HasMoreResults)
{
this.itemProducerForest.Enqueue(itemProducerTree);
}
}
}
/// <summary>
/// Pushes a document producer back to the queue.
/// </summary>
public void PushCurrentItemProducerTree(ItemProducerTree itemProducerTree)
{
this.itemProducerForest.Enqueue(itemProducerTree);
}
/// <summary>
/// When resuming an order by query we need to filter the document producers.
/// </summary>
/// <param name="producer">The producer to filter down.</param>
/// <param name="sortOrders">The sort orders.</param>
/// <param name="continuationToken">The continuation token.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A task to await on.</returns>
private async Task FilterAsync(
ItemProducerTree producer,
SortOrder[] sortOrders,
OrderByContinuationToken continuationToken,
CancellationToken cancellationToken)
{
// When we resume a query on a partition there is a possibility that we only read a partial page from the backend
// meaning that will we repeat some documents if we didn't do anything about it.
// The solution is to filter all the documents that come before in the sort order, since we have already emitted them to the client.
// The key is to seek until we get an order by value that matches the order by value we left off on.
// Once we do that we need to seek to the correct _rid within the term,
// since there might be many documents with the same order by value we left off on.
foreach (ItemProducerTree tree in producer)
{
if (!ResourceId.TryParse(continuationToken.Rid, out ResourceId continuationRid))
{
throw new CosmosException(
statusCode: HttpStatusCode.BadRequest,
message: $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.");
}
Dictionary <string, ResourceId> resourceIds = new Dictionary <string, ResourceId>();
int itemToSkip = continuationToken.SkipCount;
bool continuationRidVerified = false;
while (true)
{
OrderByQueryResult orderByResult = new OrderByQueryResult(tree.Current);
// Throw away documents until it matches the item from the continuation token.
int cmp = 0;
for (int i = 0; i < sortOrders.Length; ++i)
{
cmp = ItemComparer.Instance.Compare(
continuationToken.OrderByItems[i].Item,
orderByResult.OrderByItems[i].Item);
if (cmp != 0)
{
cmp = sortOrders[i] != SortOrder.Descending ? cmp : -cmp;
break;
}
}
if (cmp < 0)
{
// We might have passed the item due to deletions and filters.
break;
}
if (cmp == 0)
{
ResourceId rid;
if (!resourceIds.TryGetValue(orderByResult.Rid, out rid))
{
if (!ResourceId.TryParse(orderByResult.Rid, out rid))
{
throw new CosmosException(
statusCode: HttpStatusCode.BadRequest,
message: $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context~TryParse.");
}
resourceIds.Add(orderByResult.Rid, rid);
}
if (!continuationRidVerified)
{
if (continuationRid.Database != rid.Database || continuationRid.DocumentCollection != rid.DocumentCollection)
{
throw new CosmosException(
statusCode: HttpStatusCode.BadRequest,
message: $"Invalid Rid in the continuation token {continuationToken.CompositeContinuationToken.Token} for OrderBy~Context.");
}
continuationRidVerified = true;
}
// Once the item matches the order by items from the continuation tokens
// We still need to remove all the documents that have a lower rid in the rid sort order.
// If there is a tie in the sort order the documents should be in _rid order in the same direction as the first order by field.
// So if it's ORDER BY c.age ASC, c.name DESC the _rids are ASC
// If ti's ORDER BY c.age DESC, c.name DESC the _rids are DESC
cmp = continuationRid.Document.CompareTo(rid.Document);
if (sortOrders[0] == SortOrder.Descending)
{
cmp = -cmp;
}
// We might have passed the item due to deletions and filters.
// We also have a skip count for JOINs
if (cmp < 0 || (cmp == 0 && itemToSkip-- <= 0))
{
break;
}
}
(bool successfullyMovedNext, QueryResponse failureResponse)moveNextResponse = await tree.MoveNextAsync(cancellationToken);
if (!moveNextResponse.successfullyMovedNext)
{
if (moveNextResponse.failureResponse != null)
{
this.FailureResponse = moveNextResponse.failureResponse;
}
break;
}
}
}
}
/// <summary>
/// Initializes cross partition query execution context by initializing the necessary document producers.
/// </summary>
/// <param name="collectionRid">The collection to drain from.</param>
/// <param name="partitionKeyRanges">The partitions to target.</param>
/// <param name="initialPageSize">The page size to start the document producers off with.</param>
/// <param name="querySpecForInit">The query specification for the rewritten query.</param>
/// <param name="targetRangeToContinuationMap">Map from partition to it's corresponding continuation token.</param>
/// <param name="deferFirstPage">Whether or not we should defer the fetch of the first page from each partition.</param>
/// <param name="filter">The filter to inject in the predicate.</param>
/// <param name="filterCallback">The callback used to filter each partition.</param>
/// <param name="token">The cancellation token.</param>
/// <returns>A task to await on.</returns>
protected async Task InitializeAsync(
string collectionRid,
IReadOnlyList <PartitionKeyRange> partitionKeyRanges,
int initialPageSize,
SqlQuerySpec querySpecForInit,
Dictionary <string, string> targetRangeToContinuationMap,
bool deferFirstPage,
string filter,
Func <ItemProducerTree, Task> filterCallback,
CancellationToken token)
{
CollectionCache collectionCache = await this.queryContext.QueryClient.GetCollectionCacheAsync();
this.TraceInformation(string.Format(
CultureInfo.InvariantCulture,
"parallel~contextbase.initializeasync, queryspec {0}, maxbuffereditemcount: {1}, target partitionkeyrange count: {2}, maximumconcurrencylevel: {3}, documentproducer initial page size {4}",
JsonConvert.SerializeObject(querySpecForInit, DefaultJsonSerializationSettings.Value),
this.actualMaxBufferedItemCount,
partitionKeyRanges.Count,
this.comparableTaskScheduler.MaximumConcurrencyLevel,
initialPageSize));
List <ItemProducerTree> itemProducerTrees = new List <ItemProducerTree>();
foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges)
{
string initialContinuationToken = (targetRangeToContinuationMap != null && targetRangeToContinuationMap.ContainsKey(partitionKeyRange.Id)) ? targetRangeToContinuationMap[partitionKeyRange.Id] : null;
ItemProducerTree itemProducerTree = new ItemProducerTree(
this.queryContext,
querySpecForInit,
partitionKeyRange,
this.OnItemProducerTreeCompleteFetching,
this.itemProducerForest.Comparer as IComparer <ItemProducerTree>,
this.equalityComparer,
deferFirstPage,
collectionRid,
initialPageSize,
initialContinuationToken)
{
Filter = filter
};
// Prefetch if necessary, and populate consume queue.
if (this.CanPrefetch)
{
this.TryScheduleFetch(itemProducerTree);
}
itemProducerTrees.Add(itemProducerTree);
}
// Using loop fisson so that we can load the document producers in parallel
foreach (ItemProducerTree itemProducerTree in itemProducerTrees)
{
if (!deferFirstPage)
{
(bool successfullyMovedNext, QueryResponse failureResponse)response = await itemProducerTree.MoveNextIfNotSplitAsync(token);
if (response.failureResponse != null)
{
// Set the failure so on drain it can be returned.
this.FailureResponse = response.failureResponse;
// No reason to enqueue the rest of the itemProducerTrees since there is a failure.
break;
}
}
if (filterCallback != null)
{
await filterCallback(itemProducerTree);
}
if (itemProducerTree.HasMoreResults)
{
this.itemProducerForest.Enqueue(itemProducerTree);
}
}
}
/// <summary>
/// Initializes cross partition query execution context by initializing the necessary document producers.
/// </summary>
/// <param name="collectionRid">The collection to drain from.</param>
/// <param name="partitionKeyRanges">The partitions to target.</param>
/// <param name="initialPageSize">The page size to start the document producers off with.</param>
/// <param name="querySpecForInit">The query specification for the rewritten query.</param>
/// <param name="targetRangeToContinuationMap">Map from partition to it's corresponding continuation token.</param>
/// <param name="deferFirstPage">Whether or not we should defer the fetch of the first page from each partition.</param>
/// <param name="filter">The filter to inject in the predicate.</param>
/// <param name="tryFilterAsync">The callback used to filter each partition.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A task to await on.</returns>
protected async Task <TryCatch> TryInitializeAsync(
string collectionRid,
IReadOnlyList <PartitionKeyRange> partitionKeyRanges,
int initialPageSize,
SqlQuerySpec querySpecForInit,
IReadOnlyDictionary <string, string> targetRangeToContinuationMap,
bool deferFirstPage,
string filter,
Func <ItemProducerTree, Task <TryCatch> > tryFilterAsync,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
List <ItemProducerTree> itemProducerTrees = new List <ItemProducerTree>();
foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges)
{
string initialContinuationToken;
if (targetRangeToContinuationMap != null)
{
if (!targetRangeToContinuationMap.TryGetValue(partitionKeyRange.Id, out initialContinuationToken))
{
initialContinuationToken = null;
}
}
else
{
initialContinuationToken = null;
}
ItemProducerTree itemProducerTree = new ItemProducerTree(
this.queryContext,
querySpecForInit,
partitionKeyRange,
this.OnItemProducerTreeCompleteFetching,
this.itemProducerForest.Comparer as IComparer <ItemProducerTree>,
this.equalityComparer,
this.testSettings,
deferFirstPage,
collectionRid,
initialPageSize,
initialContinuationToken)
{
Filter = filter
};
// Prefetch if necessary, and populate consume queue.
if (this.CanPrefetch)
{
this.TryScheduleFetch(itemProducerTree);
}
itemProducerTrees.Add(itemProducerTree);
}
// Using loop fission so that we can load the document producers in parallel
foreach (ItemProducerTree itemProducerTree in itemProducerTrees)
{
if (!deferFirstPage)
{
while (true)
{
(bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationToken);
if (failureResponse.HasValue)
{
return(TryCatch.FromException(
new CosmosException(
statusCode: failureResponse.Value.StatusCode,
subStatusCode: (int)failureResponse.Value.SubStatusCode.GetValueOrDefault(0),
message: failureResponse.Value.ErrorMessage,
activityId: failureResponse.Value.ActivityId,
requestCharge: failureResponse.Value.RequestCharge)));
}
if (!movedToNextPage)
{
break;
}
if (itemProducerTree.IsAtBeginningOfPage)
{
break;
}
if (itemProducerTree.TryMoveNextDocumentWithinPage())
{
break;
}
}
}
if (tryFilterAsync != null)
{
TryCatch tryFilter = await tryFilterAsync(itemProducerTree);
if (!tryFilter.Succeeded)
{
return(tryFilter);
}
}
this.itemProducerForest.Enqueue(itemProducerTree);
}
return(TryCatch.FromResult());
}
/// <summary>
/// Pushes a document producer back to the queue.
/// </summary>
public void PushCurrentItemProducerTree(ItemProducerTree itemProducerTree)
{
itemProducerTree.UpdatePriority();
this.itemProducerForest.Enqueue(itemProducerTree);
}
/// <summary>
/// Initializes cross partition query execution context by initializing the necessary document producers.
/// </summary>
/// <param name="collectionRid">The collection to drain from.</param>
/// <param name="partitionKeyRanges">The partitions to target.</param>
/// <param name="initialPageSize">The page size to start the document producers off with.</param>
/// <param name="querySpecForInit">The query specification for the rewritten query.</param>
/// <param name="targetRangeToContinuationMap">Map from partition to it's corresponding continuation token.</param>
/// <param name="deferFirstPage">Whether or not we should defer the fetch of the first page from each partition.</param>
/// <param name="filter">The filter to inject in the predicate.</param>
/// <param name="tryFilterAsync">The callback used to filter each partition.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A task to await on.</returns>
protected async Task<TryCatch<bool>> TryInitializeAsync(
string collectionRid,
IReadOnlyList<PartitionKeyRange> partitionKeyRanges,
int initialPageSize,
SqlQuerySpec querySpecForInit,
IReadOnlyDictionary<string, string> targetRangeToContinuationMap,
bool deferFirstPage,
string filter,
Func<ItemProducerTree, Task<TryCatch<bool>>> tryFilterAsync,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
List<ItemProducerTree> itemProducerTrees = new List<ItemProducerTree>();
foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges)
{
string initialContinuationToken;
if (targetRangeToContinuationMap != null)
{
if (!targetRangeToContinuationMap.TryGetValue(partitionKeyRange.Id, out initialContinuationToken))
{
initialContinuationToken = null;
}
}
else
{
initialContinuationToken = null;
}
ItemProducerTree itemProducerTree = new ItemProducerTree(
this.queryContext,
querySpecForInit,
partitionKeyRange,
this.OnItemProducerTreeCompleteFetching,
this.itemProducerForest.Comparer as IComparer<ItemProducerTree>,
this.equalityComparer,
deferFirstPage,
collectionRid,
initialPageSize,
initialContinuationToken)
{
Filter = filter
};
// Prefetch if necessary, and populate consume queue.
if (this.CanPrefetch)
{
this.TryScheduleFetch(itemProducerTree);
}
itemProducerTrees.Add(itemProducerTree);
}
// Using loop fission so that we can load the document producers in parallel
foreach (ItemProducerTree itemProducerTree in itemProducerTrees)
{
if (!deferFirstPage)
{
(bool successfullyMovedNext, QueryResponseCore? failureResponse) = await itemProducerTree.MoveNextIfNotSplitAsync(cancellationToken);
if (failureResponse != null)
{
// Set the failure so on drain it can be returned.
this.FailureResponse = failureResponse;
// No reason to enqueue the rest of the itemProducerTrees since there is a failure.
break;
}
}
if (tryFilterAsync != null)
{
TryCatch<bool> tryFilter = await tryFilterAsync(itemProducerTree);
if (!tryFilter.Succeeded)
{
return tryFilter;
}
}
if (itemProducerTree.HasMoreResults)
{
this.itemProducerForest.Enqueue(itemProducerTree);
}
}
return TryCatch<bool>.FromResult(true);
}
