/// <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); }