Skip to content

Commit

Permalink
Add tests + fix issue with underreported RUs
Browse files Browse the repository at this point in the history
  • Loading branch information
@ccurrens
ccurrens committed Apr 1, 2021
1 parent 0d86f64 commit 7622574
Show file tree
Hide file tree
Showing 4 changed files with 66 additions and 18 deletions.
29 changes: 15 additions & 14 deletions ...src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,8 @@ private async ValueTask<bool> MoveNextAsync_Initialize_FromBeginningAsync(
}
else
{
QueryPage page = uninitializedEnumerator.Current.Result.Page;

if (!uninitializedEnumerator.Current.Result.Enumerator.MoveNext())
{
// Page was empty
Expand All @@ -146,12 +148,12 @@ private async ValueTask<bool> MoveNextAsync_Initialize_FromBeginningAsync(
this.Current = TryCatch<QueryPage>.FromResult(
new QueryPage(
documents: EmptyPage,
requestCharge: 0,
activityId: Guid.NewGuid().ToString(),
responseLengthInBytes: 0,
cosmosQueryExecutionInfo: default,
disallowContinuationTokenMessage: default,
additionalHeaders: default,
requestCharge: page.RequestCharge,
activityId: string.IsNullOrEmpty(page.ActivityId) ? Guid.NewGuid().ToString() : page.ActivityId,
responseLengthInBytes: page.ResponseLengthInBytes,
cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo,
disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage,
additionalHeaders: page.AdditionalHeaders,
state: null));
this.returnedFinalPage = true;
return true;
Expand All @@ -162,7 +164,6 @@ private async ValueTask<bool> MoveNextAsync_Initialize_FromBeginningAsync(
this.enumerators.Enqueue(uninitializedEnumerator);
}

QueryPage page = uninitializedEnumerator.Current.Result.Page;
// Just return an empty page with the stats
this.Current = TryCatch<QueryPage>.FromResult(
new QueryPage(
Expand Down Expand Up @@ -215,6 +216,7 @@ private async ValueTask<bool> MoveNextAsync_Initialize_FilterAsync(
}

(bool doneFiltering, int itemsLeftToSkip, TryCatch<OrderByQueryPage> monadicQueryByPage) = filterMonad.Result;
QueryPage page = uninitializedEnumerator.Current.Result.Page;
if (doneFiltering)
{
if (uninitializedEnumerator.Current.Result.Enumerator.Current != null)
Expand All @@ -228,12 +230,12 @@ private async ValueTask<bool> MoveNextAsync_Initialize_FilterAsync(
this.Current = TryCatch<QueryPage>.FromResult(
new QueryPage(
documents: EmptyPage,
requestCharge: 0,
activityId: Guid.NewGuid().ToString(),
responseLengthInBytes: 0,
cosmosQueryExecutionInfo: default,
disallowContinuationTokenMessage: default,
additionalHeaders: default,
requestCharge: page.RequestCharge,
activityId: string.IsNullOrEmpty(page.ActivityId) ? Guid.NewGuid().ToString() : page.ActivityId,
responseLengthInBytes: page.ResponseLengthInBytes,
cosmosQueryExecutionInfo: page.CosmosQueryExecutionInfo,
disallowContinuationTokenMessage: page.DisallowContinuationTokenMessage,
additionalHeaders: page.AdditionalHeaders,
state: null));
this.returnedFinalPage = true;
return true;
Expand Down Expand Up @@ -267,7 +269,6 @@ private async ValueTask<bool> MoveNextAsync_Initialize_FilterAsync(
}
}

QueryPage page = uninitializedEnumerator.Current.Result.Page;
// Just return an empty page with the stats
this.Current = TryCatch<QueryPage>.FromResult(
new QueryPage(
Expand Down
7 changes: 4 additions & 3 deletions Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/SkipEmptyPageQueryPipelineStage.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,9 +84,6 @@ public async ValueTask<bool> MoveNextAsync(ITrace trace)
QueryPage sourcePage = tryGetSourcePage.Result;
if (sourcePage.Documents.Count == 0)
{
this.cumulativeRequestCharge += sourcePage.RequestCharge;
this.cumulativeResponseLengthInBytes += sourcePage.ResponseLengthInBytes;
this.cumulativeAdditionalHeaders = sourcePage.AdditionalHeaders;
if (sourcePage.State == null)
{
QueryPage queryPage = new QueryPage(
Expand All @@ -105,6 +102,10 @@ public async ValueTask<bool> MoveNextAsync(ITrace trace)
return true;
}

this.cumulativeRequestCharge += sourcePage.RequestCharge;
this.cumulativeResponseLengthInBytes += sourcePage.ResponseLengthInBytes;
this.cumulativeAdditionalHeaders = sourcePage.AdditionalHeaders;

return await this.MoveNextAsync();
}

Expand Down
6 changes: 5 additions & 1 deletion Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -632,6 +632,10 @@ public Task<TryCatch<QueryPage>> MonadicQueryAsync(
queryState = default;
}

ImmutableDictionary<string, string>.Builder additionalHeaders = ImmutableDictionary.CreateBuilder<string, string>();
additionalHeaders.Add("x-ms-documentdb-partitionkeyrangeid", "0");
additionalHeaders.Add("x-ms-test-header", "true");

return Task.FromResult(
TryCatch<QueryPage>.FromResult(
new QueryPage(
Expand All @@ -641,7 +645,7 @@ public Task<TryCatch<QueryPage>> MonadicQueryAsync(
responseLengthInBytes: 1337,
cosmosQueryExecutionInfo: default,
disallowContinuationTokenMessage: default,
additionalHeaders: default,
additionalHeaders: additionalHeaders.ToImmutable(),
state: queryState)));
}
}
Expand Down
42 changes: 42 additions & 0 deletions ...crosoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -241,12 +241,54 @@ FROM c

QueryPage queryPage = tryGetQueryPage.Result;
documents.AddRange(queryPage.Documents);

Assert.AreNotEqual(0, queryPage.RequestCharge);
}

Assert.AreEqual(numItems, documents.Count);
Assert.IsTrue(documents.OrderBy(document => ((CosmosObject)document)["_ts"]).ToList().SequenceEqual(documents));
}

[TestMethod]
public async Task TestDrain_IncludesResponseHeadersInQueryPage()
{
IDocumentContainer documentContainer = await CreateDocumentContainerAsync(10);

TryCatch<IQueryPipelineStage> monadicCreate = OrderByCrossPartitionQueryPipelineStage.MonadicCreate(
documentContainer: documentContainer,
sqlQuerySpec: new SqlQuerySpec(@"
SELECT c._rid AS _rid, [{""item"": c._ts}] AS orderByItems, c AS payload
FROM c
WHERE {documentdb-formattableorderbyquery-filter}
ORDER BY c._ts"),
targetRanges: await documentContainer.GetFeedRangesAsync(
trace: NoOpTrace.Singleton,
cancellationToken: default),
partitionKey: null,
orderByColumns: new List<OrderByColumn>()
{
new OrderByColumn("c._ts", SortOrder.Ascending)
},
queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10),
maxConcurrency: 10,
cancellationToken: default,
continuationToken: null);
Assert.IsTrue(monadicCreate.Succeeded);
IQueryPipelineStage queryPipelineStage = monadicCreate.Result;

while (await queryPipelineStage.MoveNextAsync())
{
TryCatch<QueryPage> tryGetQueryPage = queryPipelineStage.Current;
if (tryGetQueryPage.Failed)
{
Assert.Fail(tryGetQueryPage.Exception.ToString());
}

QueryPage queryPage = tryGetQueryPage.Result;
Assert.IsTrue(queryPage.AdditionalHeaders.Count > 0);
}
}

[TestMethod]
[DataRow(false, false, false, DisplayName = "Use State: false, Allow Splits: false, Allow Merges: false")]
[DataRow(false, false, true, DisplayName = "Use State: false, Allow Splits: false, Allow Merges: true")]
Expand Down

0 comments on commit 7622574

Please sign in to comment.