Query: Fixes query performance regression and adds mocked query benchmark

Microsoft.Azure.Cosmos/src/Pagination/Page.cs

@@ -29,6 +29,9 @@ protected Page(
             this.ActivityId = activityId;
             this.State = state;
 
+#if DEBUG
+            // Only do the additional header validation on debug.
+            // This causes a significant impact to performance and is not necessary for release scenarios.
             if (additionalHeaders != null)
             {
                 foreach (string key in additionalHeaders.Keys)
@@ -39,6 +42,7 @@ protected Page(
                     }
                 }
             }
+#endif
 
             this.AdditionalHeaders = additionalHeaders == null ? EmptyDictionary : additionalHeaders.ToImmutableDictionary();
         }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Benchmarks/IItemBenchmark.cs

@@ -23,5 +23,7 @@ public interface IItemBenchmark
         public Task DeleteItemNotExists();
 
         public Task ReadFeed();
+
+        public Task QuerySinglePage();
     }
 }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Benchmarks/MockedItemBenchmark.cs

@@ -108,5 +108,12 @@ public async Task UpsertItem()
         {
             await this.CurrentBenchmark.UpsertItem();
         }
+
+        [Benchmark]
+        [BenchmarkCategory("GateBenchmark")]
+        public async Task QuerySinglePage()
+        {
+            await this.CurrentBenchmark.QuerySinglePage();
+        }
     }
 }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Benchmarks/MockedItemOfTBenchmark.cs

@@ -116,7 +116,28 @@ public async Task DeleteItemNotExists()
 
         public async Task ReadFeed()
         {
-            FeedIterator<ToDoActivity> resultIterator = this.BenchmarkHelper.TestContainer.GetItemQueryIterator<ToDoActivity>();
+            using FeedIterator<ToDoActivity> resultIterator = this.BenchmarkHelper.TestContainer.GetItemQueryIterator<ToDoActivity>();
+            while (resultIterator.HasMoreResults)
+            {
+                FeedResponse<ToDoActivity> response = await resultIterator.ReadNextAsync();
+                if (response.StatusCode != HttpStatusCode.OK || response.Resource.Count() == 0)
+                {
+                    throw new Exception();
+                }
+
+                this.BenchmarkHelper.IncludeDiagnosticToStringHelper(response.Diagnostics);
+            }
+        }
+
+        public async Task QuerySinglePage()
+        {
+            using FeedIterator<ToDoActivity> resultIterator = this.BenchmarkHelper.TestContainer.GetItemQueryIterator<ToDoActivity>(
+                "select * from T",
+                requestOptions: new QueryRequestOptions()
+                {
+                    PartitionKey = new PartitionKey("dummyValue"),
+                });
+
             while (resultIterator.HasMoreResults)
             {
                 FeedResponse<ToDoActivity> response = await resultIterator.ReadNextAsync();

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Benchmarks/MockedItemStreamBenchmark.cs

@@ -136,7 +136,25 @@ public async Task DeleteItemNotExists()
 
         public async Task ReadFeed()
         {
-            FeedIterator streamIterator = this.benchmarkHelper.TestContainer.GetItemQueryStreamIterator();
+            using FeedIterator streamIterator = this.benchmarkHelper.TestContainer.GetItemQueryStreamIterator();
+            while (streamIterator.HasMoreResults)
+            {
+                ResponseMessage response = await streamIterator.ReadNextAsync();
+                if (response.StatusCode != HttpStatusCode.OK)
+                {
+                    throw new Exception();
+                }
+            }
+        }
+
+        public async Task QuerySinglePage()
+        {
+            using FeedIterator streamIterator = this.benchmarkHelper.TestContainer.GetItemQueryStreamIterator(
+                "select * from T",
+                requestOptions: new QueryRequestOptions()
+                {
+                    PartitionKey = new PartitionKey("dummyValue"),
+                });
             while (streamIterator.HasMoreResults)
             {
                 ResponseMessage response = await streamIterator.ReadNextAsync();

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Mocks/MockDocumentClient.cs

@@ -22,6 +22,8 @@ namespace Microsoft.Azure.Cosmos.Performance.Tests
     using Microsoft.CodeAnalysis.CSharp.Syntax;
     using System.IO;
     using Microsoft.Azure.Cosmos.Tracing;
+    using Microsoft.Azure.Cosmos.Query.Core.QueryPlan;
+    using Newtonsoft.Json;
 
     internal class MockDocumentClient : DocumentClient, ICosmosAuthorizationTokenProvider
     {
@@ -103,6 +105,12 @@ public static string GenerateRandomKey()
             return Convert.ToBase64String(randomEntries);
         }
 
+        private static readonly Task<QueryPartitionProvider> SingletonQueryPartitionProvider = Task.FromResult(
+            new QueryPartitionProvider(
+                JsonConvert.DeserializeObject<Dictionary<string, object>>("{\"maxSqlQueryInputLength\":262144,\"maxJoinsPerSqlQuery\":5,\"maxLogicalAndPerSqlQuery\":500,\"maxLogicalOrPerSqlQuery\":500,\"maxUdfRefPerSqlQuery\":10,\"maxInExpressionItemsCount\":16000,\"queryMaxInMemorySortDocumentCount\":500,\"maxQueryRequestTimeoutFraction\":0.9,\"sqlAllowNonFiniteNumbers\":false,\"sqlAllowAggregateFunctions\":true,\"sqlAllowSubQuery\":true,\"sqlAllowScalarSubQuery\":true,\"allowNewKeywords\":true,\"sqlAllowLike\":true,\"sqlAllowGroupByClause\":true,\"maxSpatialQueryCells\":12,\"spatialMaxGeometryPointCount\":256,\"sqlDisableOptimizationFlags\":0,\"sqlAllowTop\":true,\"enableSpatialIndexing\":true}")));
+
+        internal override Task<QueryPartitionProvider> QueryPartitionProvider => SingletonQueryPartitionProvider;
+
         internal override IRetryPolicyFactory ResetSessionTokenRetryPolicy => new RetryPolicy(
             this.globalEndpointManager.Object,
             new ConnectionPolicy(), 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Mocks/MockRequestHelper.cs

@@ -123,6 +123,46 @@ public static DocumentServiceResponse GetDocumentServiceResponse(DocumentService
         /// <returns>A <see cref="StoreResponse"/> instance.</returns>
         public static StoreResponse GetStoreResponse(DocumentServiceRequest request)
         {
+            if (request.ResourceType == ResourceType.Document &&
+               request.OperationType == OperationType.Query)
+            {
+                StoreResponseNameValueCollection queryHeaders = new StoreResponseNameValueCollection()
+                {
+                    ActivityId = Guid.NewGuid().ToString(),
+                    BackendRequestDurationMilliseconds = "1.42",
+                    CurrentReplicaSetSize = "1",
+                    CurrentWriteQuorum = "1",
+                    CurrentResourceQuotaUsage = "documentSize=0;documentsSize=1;documentsCount=1;collectionSize=1;",
+                    GlobalCommittedLSN = "-1",
+                    LSN = "2540",
+                    LocalLSN = "2540",
+                    LastStateChangeUtc = "Wed, 18 Aug 2021 20:30:05.117 GMT",
+                    MaxResourceQuota = "documentSize=10240;documentsSize=10485760;documentsCount=-1;collectionSize=10485760;",
+                    NumberOfReadRegions = "0",
+                    OwnerFullName = "dbs/f4ac3cfd-dd38-4adb-b2d2-be97b3efbd1b/colls/2a926112-a26e-4935-ac6a-66df269c890d",
+                    OwnerId = "GHRtAJahWQ4=",
+                    PartitionKeyRangeId = "0",
+                    PendingPKDelete = "false",
+                    QueryExecutionInfo = "{\"reverseRidEnabled\":false,\"reverseIndexScan\":false}",
+                    QueryMetrics = "totalExecutionTimeInMs=0.78;queryCompileTimeInMs=0.26;queryLogicalPlanBuildTimeInMs=0.04;queryPhysicalPlanBuildTimeInMs=0.18;queryOptimizationTimeInMs=0.01;VMExecutionTimeInMs=0.04;indexLookupTimeInMs=0.00;documentLoadTimeInMs=0.02;systemFunctionExecuteTimeInMs=0.00;userFunctionExecuteTimeInMs=0.00;retrievedDocumentCount=1;retrievedDocumentSize=671;outputDocumentCount=1;outputDocumentSize=720;writeOutputTimeInMs=0.00;indexUtilizationRatio=1.00",
+                    QuorumAckedLSN = "2540" ,
+                    QuorumAckedLocalLSN = "2540",
+                    RequestCharge = "2.27" ,
+                    SchemaVersion = "1.12",
+                    ServerVersion = " version=2.14.0.0",
+                    SessionToken = "0:-1#2540",
+                    TransportRequestID = "2",
+                    XPRole = "0",
+                };
+
+                return new StoreResponse()
+                {
+                    ResponseBody = new MemoryStream(MockRequestHelper.testItemFeedResponsePayload, 0, MockRequestHelper.testItemFeedResponsePayload.Length, writable: false, publiclyVisible: true),
+                    Status = (int)System.Net.HttpStatusCode.OK,
+                    Headers = queryHeaders,
+                };
+            }
+
             StoreResponseNameValueCollection headers = MockRequestHelper.GenerateTestHeaders();
 
             if (request.OperationType == OperationType.Read)