public async Task TestMoveNextWithEmptyPagesAsync(string initialContinuationToken) { int maxPageSize = 5; List <MockPartitionResponse[]> mockResponses = MockQueryFactory.GetAllCombinationWithEmptyPage(); foreach (MockPartitionResponse[] mockResponse in mockResponses) { Mock <CosmosQueryClient> mockQueryClient = new Mock <CosmosQueryClient>(); IList <ToDoItem> allItems = MockQueryFactory.GenerateAndMockResponse( mockQueryClient, isOrderByQuery: false, sqlQuerySpec: MockQueryFactory.DefaultQuerySpec, containerRid: MockQueryFactory.DefaultCollectionRid, initContinuationToken: initialContinuationToken, maxPageSize: maxPageSize, mockResponseForSinglePartition: mockResponse, cancellationTokenForMocks: this.cancellationToken); CosmosQueryContext context = MockQueryFactory.CreateContext( mockQueryClient.Object); ItemProducerTree itemProducerTree = new ItemProducerTree( queryContext: context, querySpecForInit: MockQueryFactory.DefaultQuerySpec, partitionKeyRange: mockResponse[0].PartitionKeyRange, produceAsyncCompleteCallback: MockItemProducerFactory.DefaultTreeProduceAsyncCompleteDelegate, itemProducerTreeComparer: new ParallelItemProducerTreeComparer(), equalityComparer: CosmosElementEqualityComparer.Value, testSettings: new TestInjections(simulate429s: false, simulateEmptyPages: false), deferFirstPage: true, collectionRid: MockQueryFactory.DefaultCollectionRid, initialPageSize: maxPageSize, initialContinuationToken: initialContinuationToken); Assert.IsTrue(itemProducerTree.HasMoreResults); List <ToDoItem> itemsRead = new List <ToDoItem>(); while ((await itemProducerTree.TryMoveNextPageAsync(this.cancellationToken)).movedToNextPage) { while (itemProducerTree.TryMoveNextDocumentWithinPage()) { Assert.IsTrue(itemProducerTree.HasMoreResults); string jsonValue = itemProducerTree.Current.ToString(); ToDoItem item = JsonConvert.DeserializeObject <ToDoItem>(jsonValue); itemsRead.Add(item); } } Assert.IsFalse(itemProducerTree.HasMoreResults); Assert.AreEqual(allItems.Count, itemsRead.Count); CollectionAssert.AreEqual(itemsRead, allItems.ToList(), new ToDoItemComparer()); } }
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())); }
public async Task TestMoveNextWithEmptyPagesAndSplitAsync(string initialContinuationToken) { int maxPageSize = 5; List <MockPartitionResponse[]> mockResponsesScenario = MockQueryFactory.GetSplitScenarios(); foreach (MockPartitionResponse[] mockResponse in mockResponsesScenario) { Mock <CosmosQueryClient> mockQueryClient = new Mock <CosmosQueryClient>(); IList <ToDoItem> allItems = MockQueryFactory.GenerateAndMockResponse( mockQueryClient, isOrderByQuery: false, sqlQuerySpec: MockQueryFactory.DefaultQuerySpec, containerRid: MockQueryFactory.DefaultCollectionRid, initContinuationToken: initialContinuationToken, maxPageSize: maxPageSize, mockResponseForSinglePartition: mockResponse, cancellationTokenForMocks: this.cancellationToken); CosmosQueryContext context = MockQueryFactory.CreateContext( mockQueryClient.Object); ItemProducerTree itemProducerTree = new ItemProducerTree( context, MockQueryFactory.DefaultQuerySpec, mockResponse[0].PartitionKeyRange, MockItemProducerFactory.DefaultTreeProduceAsyncCompleteDelegate, new ParallelItemProducerTreeComparer(), CosmosElementEqualityComparer.Value, new TestInjections(simulate429s: false, simulateEmptyPages: false), true, MockQueryFactory.DefaultCollectionRid, maxPageSize, initialContinuationToken: initialContinuationToken); Assert.IsTrue(itemProducerTree.HasMoreResults); List <ToDoItem> itemsRead = new List <ToDoItem>(); while ((await itemProducerTree.TryMoveNextPageAsync(this.cancellationToken)).movedToNextPage) { while (itemProducerTree.TryMoveNextDocumentWithinPage()) { Assert.IsTrue(itemProducerTree.HasMoreResults); if (itemProducerTree.Current != null) { string jsonValue = itemProducerTree.Current.ToString(); ToDoItem item = JsonConvert.DeserializeObject <ToDoItem>(jsonValue); itemsRead.Add(item); } } itemProducerTree.UpdatePriority(); } Assert.IsFalse(itemProducerTree.HasMoreResults); Assert.AreEqual(allItems.Count, itemsRead.Count); List <ToDoItem> exepected = allItems.OrderBy(x => x.id).ToList(); List <ToDoItem> actual = itemsRead.OrderBy(x => x.id).ToList(); CollectionAssert.AreEqual(exepected, actual, new ToDoItemComparer()); } }
public async Task TestItemProducerTreeWithFailure() { int callBackCount = 0; Mock <CosmosQueryContext> mockQueryContext = new Mock <CosmosQueryContext>(); SqlQuerySpec sqlQuerySpec = new SqlQuerySpec("Select * from t"); PartitionKeyRange partitionKeyRange = new PartitionKeyRange { Id = "0", MinInclusive = "A", MaxExclusive = "B" }; void produceAsyncCompleteCallback( ItemProducerTree producer, int itemsBuffered, double resourceUnitUsage, IReadOnlyCollection <QueryPageDiagnostics> queryPageDiagnostics, long responseLengthBytes, CancellationToken token) { callBackCount++; } Mock <IComparer <ItemProducerTree> > comparer = new Mock <IComparer <ItemProducerTree> >(); Mock <IEqualityComparer <CosmosElement> > cosmosElementComparer = new Mock <IEqualityComparer <CosmosElement> >(); CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); IReadOnlyList <CosmosElement> cosmosElements = new List <CosmosElement>() { new Mock <CosmosElement>(CosmosElementType.Object).Object }; QueryPageDiagnostics diagnostics = new QueryPageDiagnostics( partitionKeyRangeId: "0", queryMetricText: "SomeRandomQueryMetricText", indexUtilizationText: null, requestDiagnostics: new PointOperationStatistics( Guid.NewGuid().ToString(), System.Net.HttpStatusCode.OK, subStatusCode: SubStatusCodes.Unknown, requestCharge: 42, errorMessage: null, method: HttpMethod.Post, requestUri: new Uri("http://localhost.com"), requestSessionToken: null, responseSessionToken: null, clientSideRequestStatistics: null), schedulingStopwatch: new SchedulingStopwatch()); IReadOnlyCollection <QueryPageDiagnostics> pageDiagnostics = new List <QueryPageDiagnostics>() { diagnostics }; mockQueryContext.Setup(x => x.ContainerResourceId).Returns("MockCollectionRid"); mockQueryContext.Setup(x => x.ExecuteQueryAsync( sqlQuerySpec, It.IsAny <string>(), It.IsAny <PartitionKeyRangeIdentity>(), It.IsAny <bool>(), It.IsAny <int>(), It.IsAny <SchedulingStopwatch>(), cancellationTokenSource.Token)).Returns( Task.FromResult(QueryResponseCore.CreateSuccess( result: cosmosElements, requestCharge: 42, activityId: "AA470D71-6DEF-4D61-9A08-272D8C9ABCFE", diagnostics: pageDiagnostics, responseLengthBytes: 500, disallowContinuationTokenMessage: null, continuationToken: "TestToken"))); ItemProducerTree itemProducerTree = new ItemProducerTree( queryContext: mockQueryContext.Object, querySpecForInit: sqlQuerySpec, partitionKeyRange: partitionKeyRange, produceAsyncCompleteCallback: produceAsyncCompleteCallback, itemProducerTreeComparer: comparer.Object, equalityComparer: cosmosElementComparer.Object, testSettings: new TestInjections(simulate429s: false, simulateEmptyPages: false), deferFirstPage: false, collectionRid: "collectionRid", initialContinuationToken: null, initialPageSize: 50); // Buffer to success responses await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token); await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token); diagnostics = new QueryPageDiagnostics( partitionKeyRangeId: "0", queryMetricText: null, indexUtilizationText: null, requestDiagnostics: new PointOperationStatistics( Guid.NewGuid().ToString(), System.Net.HttpStatusCode.InternalServerError, subStatusCode: SubStatusCodes.Unknown, requestCharge: 10.2, errorMessage: "Error message", method: HttpMethod.Post, requestUri: new Uri("http://localhost.com"), requestSessionToken: null, responseSessionToken: null, clientSideRequestStatistics: null), schedulingStopwatch: new SchedulingStopwatch()); pageDiagnostics = new List <QueryPageDiagnostics>() { diagnostics }; // Buffer a failure mockQueryContext.Setup(x => x.ExecuteQueryAsync( sqlQuerySpec, It.IsAny <string>(), It.IsAny <PartitionKeyRangeIdentity>(), It.IsAny <bool>(), It.IsAny <int>(), It.IsAny <SchedulingStopwatch>(), cancellationTokenSource.Token)).Returns( Task.FromResult(QueryResponseCore.CreateFailure( statusCode: HttpStatusCode.InternalServerError, subStatusCodes: null, errorMessage: "Error message", requestCharge: 10.2, activityId: Guid.NewGuid().ToString(), diagnostics: pageDiagnostics))); await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token); // First item should be a success { (bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationTokenSource.Token); Assert.IsTrue(movedToNextPage); Assert.IsNull(failureResponse); Assert.IsTrue(itemProducerTree.TryMoveNextDocumentWithinPage()); Assert.IsFalse(itemProducerTree.TryMoveNextDocumentWithinPage()); Assert.IsTrue(itemProducerTree.HasMoreResults); } // Second item should be a success { (bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationTokenSource.Token); Assert.IsTrue(movedToNextPage); Assert.IsNull(failureResponse); Assert.IsTrue(itemProducerTree.TryMoveNextDocumentWithinPage()); Assert.IsFalse(itemProducerTree.TryMoveNextDocumentWithinPage()); Assert.IsTrue(itemProducerTree.HasMoreResults); } // Third item should be a failure { (bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationTokenSource.Token); Assert.IsFalse(movedToNextPage); Assert.IsNotNull(failureResponse); Assert.IsFalse(itemProducerTree.HasMoreResults); } // Try to buffer after failure. It should return the previous cached failure and not try to buffer again. mockQueryContext.Setup(x => x.ExecuteQueryAsync( sqlQuerySpec, It.IsAny <string>(), It.IsAny <PartitionKeyRangeIdentity>(), It.IsAny <bool>(), It.IsAny <int>(), It.IsAny <SchedulingStopwatch>(), cancellationTokenSource.Token)). Throws(new Exception("Previous buffer failed. Operation should return original failure and not try again")); await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token); Assert.IsFalse(itemProducerTree.HasMoreResults); }
/// <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()); }
protected async Task <TryCatch> TryInitializeAsync( string collectionRid, int initialPageSize, SqlQuerySpec querySpecForInit, IReadOnlyDictionary <PartitionKeyRange, string> targetRangeToContinuationMap, bool deferFirstPage, string filter, Func <ItemProducerTree, Task <TryCatch> > tryFilterAsync, CancellationToken cancellationToken) { if (collectionRid == null) { throw new ArgumentNullException(nameof(collectionRid)); } if (initialPageSize < 0) { throw new ArgumentOutOfRangeException(nameof(initialPageSize)); } if (querySpecForInit == null) { throw new ArgumentNullException(nameof(querySpecForInit)); } if (targetRangeToContinuationMap == null) { throw new ArgumentNullException(nameof(targetRangeToContinuationMap)); } cancellationToken.ThrowIfCancellationRequested(); List <ItemProducerTree> itemProducerTrees = new List <ItemProducerTree>(); foreach (KeyValuePair <PartitionKeyRange, string> rangeAndContinuationToken in targetRangeToContinuationMap) { PartitionKeyRange partitionKeyRange = rangeAndContinuationToken.Key; string continuationToken = rangeAndContinuationToken.Value; ItemProducerTree itemProducerTree = new ItemProducerTree( this.queryContext, querySpecForInit, partitionKeyRange, this.OnItemProducerTreeCompleteFetching, this.itemProducerForest.Comparer, this.equalityComparer, this.testSettings, deferFirstPage, collectionRid, initialPageSize, continuationToken) { 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( failureResponse.Value.CosmosException)); } 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> /// 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 <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 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 (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(); try { if (!currentItemProducerTree.HasMoreResults) { // This means there are no more items to drain break; } 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; this.previousOrderByItems = orderByQueryResult.OrderByItems; if (!currentItemProducerTree.TryMoveNextDocumentWithinPage()) { while (true) { (bool movedToNextPage, QueryResponseCore? failureResponse) = await currentItemProducerTree.TryMoveNextPageAsync(cancellationToken); if (!movedToNextPage) { if (failureResponse.HasValue) { // TODO: We can buffer this failure so that the user can still get the pages we already got. return(failureResponse.Value); } break; } if (currentItemProducerTree.IsAtBeginningOfPage) { break; } if (currentItemProducerTree.TryMoveNextDocumentWithinPage()) { break; } } } } 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())); }