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