From 1dbfc8a4e21789a95fddcffdecd1f2f753d788ec Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 19 Oct 2022 13:45:20 -0700 Subject: [PATCH 01/37] Added tests to test different aspects of merge/split support with OptimisticDirectExecution pipeline. Tests check for gone exception handling, pipeline switching etc. --- ...misticDirectExecutionQueryBaselineTests.cs | 358 ++++++++++++------ 1 file changed, 238 insertions(+), 120 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index c2ac15064d..e9f25184b2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -35,52 +35,120 @@ public class OptimisticDirectExecutionQueryBaselineTests : BaselineTests mockDocumentContainer = new Mock(); - - TryCatch monadicCreate = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( - documentContainer: mockDocumentContainer.Object, - sqlQuerySpec: new SqlQuerySpec("SELECT VALUE COUNT(1) FROM c"), - targetRange: FeedRangeEpk.FullRange, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), - partitionKey: null, - cancellationToken: default, - continuationToken: null); - Assert.IsTrue(monadicCreate.Succeeded); + public void PositiveOptimisticDirectExecutionOutput() + { + List testVariations = new List + { + CreateInput( + @"Partition Key + Value and Distinct", + "SELECT DISTINCT c.key FROM c", + true, + @"/pk", + @"/value"), + + CreateInput( + @"Partition Key + Value and Min Aggregate", + "SELECT VALUE MIN(c.key) FROM c", + true, + @"/pk", + @"/value"), + + CreateInput( + @"Partition Key + Value Fields", + "SELECT c.key FROM c", + true, + @"/pk", + @"/value"), + }; + this.ExecuteTestSuite(testVariations); } [TestMethod] [Owner("akotalwar")] - public void TestPipelineSingleContinuationToken() + public void NegativeOptimisticDirectExecutionOutput() { - Mock mockDocumentContainer = new Mock(); + List testVariations = new List + { + CreateInput( + @"Null Partition Key Value", + "SELECT * FROM c", + false, + @"/pk", + Cosmos.PartitionKey.Null), + CreateInput( + @"None Partition Key Value", + "SELECT * FROM c", + false, + @"/pk", + Cosmos.PartitionKey.None), + + CreateInput( + @"C# Null Partition Key Value", + "SELECT * FROM c", + false, + @"/pk", + null), + }; + this.ExecuteTestSuite(testVariations); + } + + // This test confirms that TestInjection.EnableOptimisticDirectExection is set to false from default. + // Check test "TestPipelineForDistributedQueryAsync" to understand why this is done + [TestMethod] + public async Task TestDefaultTestInjectionSettings() + { + TestInjections testInjection = new TestInjections(simulate429s: false, simulateEmptyPages: false); + + Assert.AreEqual(testInjection.EnableOptimisticDirectExecution, false); + } + + [TestMethod] + [Owner("akotalwar")] + public async Task TestPipelineNullContinuationToken() + { + int numItems = 10; + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); + string query = "SELECT * FROM c"; + + TryCatch monadicCreate = await ODEMonadicCreate( + inMemoryCollection, + query, + FeedRangeEpk.FullRange, + null); + + Assert.IsTrue(monadicCreate.Succeeded); + } + + [TestMethod] + [Owner("akotalwar")] + public async Task TestPipelineSingleContinuationToken() + { + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(5, multiPartition: false); + Range range = new Documents.Routing.Range("A", "B", isMinInclusive: true, isMaxInclusive: false); ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( token: "asdf", - range: new Documents.Routing.Range("A", "B", true, false)); + range: range); OptimisticDirectExecutionContinuationToken token = new OptimisticDirectExecutionContinuationToken(parallelContinuationToken); CosmosElement cosmosElementContinuationToken = OptimisticDirectExecutionContinuationToken.ToCosmosElement(token); + string query = "SELECT * FROM c"; + + TryCatch monadicCreate = await ODEMonadicCreate(inMemoryCollection, + query, + new FeedRangeEpk(range), + cosmosElementContinuationToken); - TryCatch monadicCreate = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( - documentContainer: mockDocumentContainer.Object, - sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), - targetRange: new FeedRangeEpk(new Documents.Routing.Range(min: "A", max: "B", isMinInclusive: true, isMaxInclusive: false)), - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), - partitionKey: null, - cancellationToken: default, - continuationToken: cosmosElementContinuationToken); Assert.IsTrue(monadicCreate.Succeeded); } // test checks that the pipeline can take a query to the backend and returns its associated document(s). [TestMethod] - public async Task TestPipelineForBackendDocumentsAsync() + public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() { int numItems = 10; string query = "SELECT VALUE COUNT(1) FROM c"; - IDocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems); + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); int documentCountInSinglePartition = 0; @@ -92,16 +160,16 @@ public async Task TestPipelineForBackendDocumentsAsync() documentCountInSinglePartition += Int32.Parse(tryGetPage.Result.Documents[0].ToString()); } - Assert.AreEqual(documentCountInSinglePartition, 4); + Assert.AreEqual(documentCountInSinglePartition, 10); } // test checks that the pipeline can take a query to the backend and returns its associated document(s) + continuation token. [TestMethod] - public async Task TestPipelineForContinuationTokenAsync() + public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() { int numItems = 100; string query = "SELECT * FROM c"; - IDocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems); + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); List documents = new List(); int continuationTokenCount = 0; @@ -125,42 +193,117 @@ public async Task TestPipelineForContinuationTokenAsync() continuationTokenCount++; } - Assert.AreEqual(continuationTokenCount, 2); - Assert.AreEqual(documents.Count, 17); + Assert.AreEqual(continuationTokenCount, 10); + Assert.AreEqual(documents.Count, 100); } - // This test confirms that TestInjection.EnableOptimisticDirectExection is set to false from default. - // Check test "TestPipelineForDistributedQueryAsync" to understand why this is done + // test to check if pipeline handles a 410 exception properly and returns all the documents. + // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing + // TODO: With the addition of the merge/split supprt, this queryPipelineStage should be able to return all documents regardless of a gone exception happening [TestMethod] - public async Task TestDefaultTestInjectionSettings() + public async Task TestPipelineForGoneExceptionAsync() { - TestInjections testInjection = new TestInjections(simulate429s: false, simulateEmptyPages: false); + int numItems = 100; + string query = "SELECT * FROM c"; + List documents = new List(); + string errorMessage = $"Epk Range: Partition does not exist at the given range."; + CosmosException goneException = new CosmosException( + message: errorMessage, + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: Guid.NewGuid().ToString(), + requestCharge: default); + + int moveNextAsyncCounter = 0; + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( + numItems, + multiPartition: true, + failureConfigs: new FlakyDocumentContainer.FailureConfigs( + inject429s: false, + injectEmptyPages: false, + shouldReturnFailure: ()=> Task.FromResult(moveNextAsyncCounter == 0 ? null: goneException))); - Assert.AreEqual(testInjection.EnableOptimisticDirectExecution, false); + IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) + { + moveNextAsyncCounter++; + try + { + TryCatch tryGetPage = queryPipelineStage.Current; + tryGetPage.ThrowIfFailed(); + + documents.AddRange(tryGetPage.Result.Documents); + } + catch + { + // check if gone exception is handled properly + Assert.IsTrue(queryPipelineStage.Current.Failed); + Assert.AreEqual(queryPipelineStage.Current.InnerMostException.Message, errorMessage); + Assert.AreEqual(((CosmosException)queryPipelineStage.Current.InnerMostException).StatusCode, System.Net.HttpStatusCode.Gone); + break; + } + } + Assert.AreEqual(documents.Count, 10); + } + + // Start with ODE pipeline and then switch to parallel in a cross partition scenario. + // This is to simulate whether a fallback solution would work and provide client with all the required documents + [TestMethod] + public async Task TestPipelineSwitchForContinuationTokenOnMultiplePartitionsAsync() + { + int numItems = 100; + string query = "SELECT * FROM c"; + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: true); + IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + List documents = new List(); + int continuationTokenCount = 0; + + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) + { + TryCatch tryGetPage = queryPipelineStage.Current; + tryGetPage.ThrowIfFailed(); + + documents.AddRange(tryGetPage.Result.Documents); + + if (tryGetPage.Result.State == null) + { + break; + } + else + { + queryPipelineStage = await CreateParallelCrossPartitionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); + } + + continuationTokenCount++; + } + + Assert.AreEqual(continuationTokenCount, 15); + Assert.AreEqual(documents.Count, 100); } // The reason we have the below test is to show the missing capabilities of the OptimisticDirectExecution pipeline. // Currently this pipeline cannot handle distributed queries as it does not have the logic to sum up the values it gets from the backend in partial results. // This functionality is available for other pipelines such as the ParallelCrossPartitionQueryPipelineStage as evident below [TestMethod] - public async Task TestPipelineForDistributedQueryAsync() + public async Task TestPipelinesForDistributedQueryAsync() { int numItems = 100; string query = "SELECT VALUE COUNT(1) FROM c"; - IDocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems); + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: true); IQueryPipelineStage optimisticDirectExecutionQueryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); IQueryPipelineStage parallelQueryPipelineStage = await CreateParallelCrossPartitionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); int documentCountOptimisticPipeline = 0; int documentCountParallelPipeline = 0; List queryPipelineStages = new List - { + { optimisticDirectExecutionQueryPipelineStage, parallelQueryPipelineStage }; - List documentPipelinesCount = new List - { + List documentPipelinesCount = new List + { documentCountOptimisticPipeline, documentCountParallelPipeline }; @@ -192,28 +335,50 @@ public async Task TestPipelineForDistributedQueryAsync() int countDifference = documentCountParallelPipeline - documentCountOptimisticPipeline; Assert.AreNotEqual(documentCountOptimisticPipeline, documentCountParallelPipeline, countDifference.ToString()); - Assert.AreEqual(documentCountOptimisticPipeline, 17); + Assert.AreEqual(documentCountOptimisticPipeline, 24); Assert.AreEqual(documentCountParallelPipeline, 100); - Assert.AreEqual(countDifference, 83); + Assert.AreEqual(countDifference, 76); Assert.AreEqual(documentCountParallelPipeline, numItems); } - private static async Task CreateOptimisticDirectExecutionPipelineStateAsync(IDocumentContainer documentContainer, string query, CosmosElement continuationToken) + private static async Task> ODEMonadicCreate(DocumentContainer documentContainer, string query, FeedRangeEpk targetRange, CosmosElement continuationToken) { - List targetRanges = await documentContainer.GetFeedRangesAsync( - trace: NoOpTrace.Singleton, - cancellationToken: default); - FeedRangeEpk firstRange = targetRanges[0]; - TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: new SqlQuerySpec(query), - targetRange: firstRange, + targetRange: targetRange, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, cancellationToken: default, continuationToken: continuationToken); + return monadicQueryPipelineStage; + } + + private static async Task CreateOptimisticDirectExecutionPipelineStateAsync(DocumentContainer documentContainer, string query, CosmosElement continuationToken) + { + List targetRanges = await documentContainer.GetFeedRangesAsync( + trace: NoOpTrace.Singleton, + cancellationToken: default); + FeedRangeEpk firstRange = null; + + if (targetRanges.Count > 1) + { + for (int i = 0; i < targetRanges.Capacity; i++) + { + if (targetRanges[i].Range.Min == "") + { + firstRange = targetRanges[i]; + } + } + } + else + { + firstRange = targetRanges[0]; + } + + TryCatch monadicQueryPipelineStage = await ODEMonadicCreate(documentContainer, query, firstRange, continuationToken); + Assert.IsTrue(monadicQueryPipelineStage.Succeeded); IQueryPipelineStage queryPipelineStage = monadicQueryPipelineStage.Result; @@ -226,6 +391,13 @@ private static async Task CreateParallelCrossPartitionPipel trace: NoOpTrace.Singleton, cancellationToken: default); + if (continuationToken is CosmosObject) + { + ((CosmosObject)continuationToken).TryGetValue("OptimisticDirectExecutionToken", out CosmosElement parallelContinuationToken); + CosmosArray cosmosElementParallelContinuationToken = CosmosArray.Create(parallelContinuationToken); + continuationToken = cosmosElementParallelContinuationToken; + } + TryCatch monadicQueryPipelineStage = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: new SqlQuerySpec(query), @@ -243,8 +415,9 @@ private static async Task CreateParallelCrossPartitionPipel return queryPipelineStage; } - private static async Task CreateDocumentContainerAsync( + private static async Task CreateDocumentContainerAsync( int numItems, + bool multiPartition, FlakyDocumentContainer.FailureConfigs failureConfigs = null) { PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() @@ -265,14 +438,18 @@ private static async Task CreateDocumentContainerAsync( DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); - for (int i = 0; i < 3; i++) + for (int i = 0; i < 2; i++) { IReadOnlyList ranges = await documentContainer.GetFeedRangesAsync( trace: NoOpTrace.Singleton, cancellationToken: default); - foreach (FeedRangeInternal range in ranges) + + if (multiPartition) { - await documentContainer.SplitAsync(range, cancellationToken: default); + foreach (FeedRangeInternal range in ranges) + { + await documentContainer.SplitAsync(range, cancellationToken: default); + } } await documentContainer.RefreshProviderAsync(NoOpTrace.Singleton, cancellationToken: default); @@ -295,66 +472,6 @@ private static async Task CreateDocumentContainerAsync( return documentContainer; } - [TestMethod] - [Owner("akotalwar")] - public void PositiveOptimisticDirectExecutionOutput() - { - List testVariations = new List - { - CreateInput( - @"Partition Key + Value and Distinct", - "SELECT DISTINCT c.key FROM c", - true, - @"/pk", - @"/value"), - - CreateInput( - @"Partition Key + Value and Min Aggregate", - "SELECT VALUE MIN(c.key) FROM c", - true, - @"/pk", - @"/value"), - - CreateInput( - @"Partition Key + Value Fields", - "SELECT c.key FROM c", - true, - @"/pk", - @"/value"), - }; - this.ExecuteTestSuite(testVariations); - } - - [TestMethod] - [Owner("akotalwar")] - public void NegativeOptimisticDirectExecutionOutput() - { - List testVariations = new List - { - CreateInput( - @"Null Partition Key Value", - "SELECT * FROM c", - false, - @"/pk", - Cosmos.PartitionKey.Null), - - CreateInput( - @"None Partition Key Value", - "SELECT * FROM c", - false, - @"/pk", - Cosmos.PartitionKey.None), - - CreateInput( - @"C# Null Partition Key Value", - "SELECT * FROM c", - false, - @"/pk", - null), - }; - this.ExecuteTestSuite(testVariations); - } - private static OptimisticDirectExecutionTestInput CreateInput( string description, string query, @@ -377,7 +494,7 @@ private static OptimisticDirectExecutionTestInput CreateInput( { return new OptimisticDirectExecutionTestInput(description, query, new SqlQuerySpec(query), expectedOptimisticDirectExecution, partitionKeyPath, partitionKeyValue); } - + private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(string querySpecJsonString, PartitionKeyDefinition pkDefinition) { TryCatch tryGetQueryPlan = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfo( @@ -446,7 +563,7 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect returnResultsInDeterministicOrder: null, forcePassthrough: true, testInjections: queryRequestOptions.TestSettings); - + IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( documentContainer, cosmosQueryContextCore, @@ -458,13 +575,14 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect { Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); } - else { + else + { Assert.AreNotEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); } - + Assert.IsNotNull(queryPipelineStage); Assert.IsTrue(result); - + return new OptimisticDirectExecutionTestOutput(input.ExpectedOptimisticDirectExecution); } } @@ -532,7 +650,7 @@ public override void SerializeAsXml(XmlWriter xmlWriter) } } - xmlWriter.WriteEndElement(); + xmlWriter.WriteEndElement(); if (this.PartitionKeyDefinition != null) { xmlWriter.WriteElementString( From 2fc0b89e55ddcbebc59a9277f75ed500fd1558ef Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 26 Oct 2022 15:59:32 -0700 Subject: [PATCH 02/37] Added gone exception simulation tests. --- ...misticDirectExecutionQueryBaselineTests.cs | 99 +++++++++++++++++-- 1 file changed, 92 insertions(+), 7 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index e9f25184b2..bef57f1588 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -146,7 +146,7 @@ public async Task TestPipelineSingleContinuationToken() [TestMethod] public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() { - int numItems = 10; + int numItems = 100; string query = "SELECT VALUE COUNT(1) FROM c"; DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); @@ -160,7 +160,7 @@ public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() documentCountInSinglePartition += Int32.Parse(tryGetPage.Result.Documents[0].ToString()); } - Assert.AreEqual(documentCountInSinglePartition, 10); + Assert.AreEqual(documentCountInSinglePartition, 100); } // test checks that the pipeline can take a query to the backend and returns its associated document(s) + continuation token. @@ -199,9 +199,58 @@ public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() // test to check if pipeline handles a 410 exception properly and returns all the documents. // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing - // TODO: With the addition of the merge/split supprt, this queryPipelineStage should be able to return all documents regardless of a gone exception happening + // TODO: With the addition of the merge/split support, this queryPipelineStage should be able to return all documents regardless of a gone exception happening [TestMethod] - public async Task TestPipelineForGoneExceptionAsync() + public async Task TestPipelineForGoneExceptionSinglePartitionAsync() + { + int numItems = 100; + string query = "SELECT * FROM c"; + List documents = new List(); + string errorMessage = $"Epk Range: Partition does not exist at the given range."; + CosmosException goneException = new CosmosException( + message: errorMessage, + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: Guid.NewGuid().ToString(), + requestCharge: default); + + int moveNextAsyncCounter = 0; + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( + numItems, + multiPartition: false, + failureConfigs: new FlakyDocumentContainer.FailureConfigs( + inject429s: false, + injectEmptyPages: false, + shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter != 1 ? null : goneException))); + + IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) + { + moveNextAsyncCounter++; + try + { + TryCatch tryGetPage = queryPipelineStage.Current; + tryGetPage.ThrowIfFailed(); + + documents.AddRange(tryGetPage.Result.Documents); + } + catch + { + // check if gone exception is handled properly + Assert.IsTrue(queryPipelineStage.Current.Failed); + Assert.AreEqual(queryPipelineStage.Current.InnerMostException.Message, errorMessage); + Assert.AreEqual(((CosmosException)queryPipelineStage.Current.InnerMostException).StatusCode, System.Net.HttpStatusCode.Gone); + break; + } + } + // Once fallback plan is implemented, this test should be able to return all 100 documents + Assert.AreEqual(documents.Count, 10); + } + + // test finds out pipeline response to a gone exception on multiple partitions + [TestMethod] + public async Task TestPipelineForGoneExceptionMultiplePartitionsAsync() { int numItems = 100; string query = "SELECT * FROM c"; @@ -221,7 +270,7 @@ public async Task TestPipelineForGoneExceptionAsync() failureConfigs: new FlakyDocumentContainer.FailureConfigs( inject429s: false, injectEmptyPages: false, - shouldReturnFailure: ()=> Task.FromResult(moveNextAsyncCounter == 0 ? null: goneException))); + shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter != 1 ? null : goneException))); IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); @@ -244,9 +293,45 @@ public async Task TestPipelineForGoneExceptionAsync() break; } } + // Once fallback plan is implemented, this test should be able to return all 100 documents Assert.AreEqual(documents.Count, 10); } + // Start with ODE pipeline and then switch to parallel in a single partition scenario. + // This is to simulate whether a fallback solution would work and provide client with all the required documents + [TestMethod] + public async Task TestPipelineSwitchForContinuationTokenOnSinglePartitionAsync() + { + int numItems = 100; + string query = "SELECT * FROM c"; + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); + IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + List documents = new List(); + int continuationTokenCount = 0; + + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) + { + TryCatch tryGetPage = queryPipelineStage.Current; + tryGetPage.ThrowIfFailed(); + + documents.AddRange(tryGetPage.Result.Documents); + + if (tryGetPage.Result.State == null) + { + break; + } + else + { + queryPipelineStage = await CreateParallelCrossPartitionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); + } + + continuationTokenCount++; + } + + Assert.AreEqual(continuationTokenCount, 10); + Assert.AreEqual(documents.Count, 100); + } + // Start with ODE pipeline and then switch to parallel in a cross partition scenario. // This is to simulate whether a fallback solution would work and provide client with all the required documents [TestMethod] @@ -713,8 +798,8 @@ public override Task> GetTargetPartitionKeyRangesAsync(s { return Task.FromResult(new List{new PartitionKeyRange() { - MinInclusive = PartitionKeyHash.V2.Hash("abc").ToString(), - MaxExclusive = PartitionKeyHash.V2.Hash("def").ToString() + MinInclusive = PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, + MaxExclusive = PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey } }); } From 36d18bb9a191e953b8586b9100f6adbe09d7750f Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Fri, 4 Nov 2022 11:50:48 -0700 Subject: [PATCH 03/37] Added new tests and improved test infra --- ...misticDirectExecutionQueryBaselineTests.cs | 533 +++++++----------- 1 file changed, 205 insertions(+), 328 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index bef57f1588..5c65ea9356 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -40,25 +40,25 @@ public void PositiveOptimisticDirectExecutionOutput() List testVariations = new List { CreateInput( - @"Partition Key + Value and Distinct", - "SELECT DISTINCT c.key FROM c", - true, - @"/pk", - @"/value"), + description: @"Partition Key + Value and Distinct", + query: "SELECT DISTINCT c.key FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: @"/value"), CreateInput( - @"Partition Key + Value and Min Aggregate", - "SELECT VALUE MIN(c.key) FROM c", - true, - @"/pk", - @"/value"), + description: @"Partition Key + Value and Min Aggregate", + query: "SELECT VALUE MIN(c.key) FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: @"/value"), CreateInput( - @"Partition Key + Value Fields", - "SELECT c.key FROM c", - true, - @"/pk", - @"/value"), + description: @"Partition Key + Value Fields", + query: "SELECT c.key FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: @"/value"), }; this.ExecuteTestSuite(testVariations); } @@ -70,25 +70,25 @@ public void NegativeOptimisticDirectExecutionOutput() List testVariations = new List { CreateInput( - @"Null Partition Key Value", - "SELECT * FROM c", - false, - @"/pk", - Cosmos.PartitionKey.Null), + description: @"Null Partition Key Value", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: Cosmos.PartitionKey.Null), CreateInput( - @"None Partition Key Value", - "SELECT * FROM c", - false, - @"/pk", - Cosmos.PartitionKey.None), + description: @"None Partition Key Value", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: Cosmos.PartitionKey.None), CreateInput( - @"C# Null Partition Key Value", - "SELECT * FROM c", - false, - @"/pk", - null), + description: @"C# Null Partition Key Value", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: null), }; this.ExecuteTestSuite(testVariations); } @@ -105,41 +105,28 @@ public async Task TestDefaultTestInjectionSettings() [TestMethod] [Owner("akotalwar")] - public async Task TestPipelineNullContinuationToken() + public async Task TestMonadicCreateODEPipeline() { int numItems = 10; - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); + bool multiPartition = false; string query = "SELECT * FROM c"; - TryCatch monadicCreate = await ODEMonadicCreate( - inMemoryCollection, - query, - FeedRangeEpk.FullRange, - null); + // null continuation token + Assert.IsTrue(await HasMonadicCreateSucceeded(numItems, multiPartition, query, targetRange: FeedRangeEpk.FullRange, continuationToken: null)); - Assert.IsTrue(monadicCreate.Succeeded); - } + // default continuation token + Assert.IsTrue(await HasMonadicCreateSucceeded(numItems, multiPartition, query, targetRange: FeedRangeEpk.FullRange, continuationToken: default)); - [TestMethod] - [Owner("akotalwar")] - public async Task TestPipelineSingleContinuationToken() - { - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(5, multiPartition: false); - Range range = new Documents.Routing.Range("A", "B", isMinInclusive: true, isMaxInclusive: false); + Range range = new Documents.Routing.Range("A", "B", isMinInclusive: true, isMaxInclusive: false); ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( token: "asdf", range: range); OptimisticDirectExecutionContinuationToken token = new OptimisticDirectExecutionContinuationToken(parallelContinuationToken); CosmosElement cosmosElementContinuationToken = OptimisticDirectExecutionContinuationToken.ToCosmosElement(token); - string query = "SELECT * FROM c"; - - TryCatch monadicCreate = await ODEMonadicCreate(inMemoryCollection, - query, - new FeedRangeEpk(range), - cosmosElementContinuationToken); - - Assert.IsTrue(monadicCreate.Succeeded); + + // single continuation token + Assert.IsTrue(await HasMonadicCreateSucceeded(numItems, multiPartition, query, targetRange: new FeedRangeEpk(range), continuationToken: cosmosElementContinuationToken)); } // test checks that the pipeline can take a query to the backend and returns its associated document(s). @@ -167,180 +154,129 @@ public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() [TestMethod] public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() { - int numItems = 100; - string query = "SELECT * FROM c"; - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); - IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); - List documents = new List(); - int continuationTokenCount = 0; - - while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) - { - TryCatch tryGetPage = queryPipelineStage.Current; - tryGetPage.ThrowIfFailed(); - - documents.AddRange(tryGetPage.Result.Documents); - - if (tryGetPage.Result.State == null) - { - break; - } - else - { - queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); - } - - continuationTokenCount++; - } + Assert.IsTrue(await this.GetResultsFromHydratingPipelineWithContinuation( + numItems: 100, + isMultiPartition: false, + isODE: true, + query: "SELECT * FROM c")); + } - Assert.AreEqual(continuationTokenCount, 10); - Assert.AreEqual(documents.Count, 100); + // test checks that ODE pipeline does not get invoked when CrossPartition query is run + [TestMethod] + public async Task TestParallelContTokenEvocationAsync() + { + Assert.IsTrue(await this.GetResultsFromHydratingPipelineWithContinuation( + numItems: 100, + isMultiPartition: true, + isODE: false, + query: "SELECT * FROM c")); } // test to check if pipeline handles a 410 exception properly and returns all the documents. - // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing - // TODO: With the addition of the merge/split support, this queryPipelineStage should be able to return all documents regardless of a gone exception happening [TestMethod] - public async Task TestPipelineForGoneExceptionSinglePartitionAsync() + public async Task TestPipelineForGoneExceptionOnSingleAndMultiplePartitionAsync() { - int numItems = 100; - string query = "SELECT * FROM c"; - List documents = new List(); - string errorMessage = $"Epk Range: Partition does not exist at the given range."; - CosmosException goneException = new CosmosException( - message: errorMessage, - statusCode: System.Net.HttpStatusCode.Gone, - subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, - activityId: Guid.NewGuid().ToString(), - requestCharge: default); + Assert.IsTrue(await ExecuteGoneExceptionOnODEPipeline(isMultiPartition: false)); - int moveNextAsyncCounter = 0; - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( - numItems, - multiPartition: false, - failureConfigs: new FlakyDocumentContainer.FailureConfigs( - inject429s: false, - injectEmptyPages: false, - shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter != 1 ? null : goneException))); - - IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + Assert.IsTrue(await ExecuteGoneExceptionOnODEPipeline(isMultiPartition: true)); + } - while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) + // The reason we have the below test is to show the missing capabilities of the OptimisticDirectExecution pipeline. + // Currently this pipeline cannot handle distributed queries as it does not have the logic to sum up the values it gets from the backend in partial results. + // This functionality is available for other pipelines such as the ParallelCrossPartitionQueryPipelineStage. + [TestMethod] + public async Task TestPipelinesForDistributedQueryAsync() + { + bool result = false; + try { - moveNextAsyncCounter++; - try - { - TryCatch tryGetPage = queryPipelineStage.Current; - tryGetPage.ThrowIfFailed(); - - documents.AddRange(tryGetPage.Result.Documents); - } - catch - { - // check if gone exception is handled properly - Assert.IsTrue(queryPipelineStage.Current.Failed); - Assert.AreEqual(queryPipelineStage.Current.InnerMostException.Message, errorMessage); - Assert.AreEqual(((CosmosException)queryPipelineStage.Current.InnerMostException).StatusCode, System.Net.HttpStatusCode.Gone); - break; - } + result = await this.GetResultsFromHydratingPipelineWithContinuation(100, false, true, "SELECT AVG(c) FROM c"); + } + catch + { + // Coming into catch because doc count Assert failed. Once aggregate logic is added, this should not go into Catch + Assert.IsFalse(result); } - // Once fallback plan is implemented, this test should be able to return all 100 documents - Assert.AreEqual(documents.Count, 10); } - // test finds out pipeline response to a gone exception on multiple partitions - [TestMethod] - public async Task TestPipelineForGoneExceptionMultiplePartitionsAsync() + private static async Task HasMonadicCreateSucceeded(int numItems, bool multiPartition, string query, FeedRangeEpk targetRange, CosmosElement continuationToken) { - int numItems = 100; - string query = "SELECT * FROM c"; - List documents = new List(); - string errorMessage = $"Epk Range: Partition does not exist at the given range."; - CosmosException goneException = new CosmosException( - message: errorMessage, - statusCode: System.Net.HttpStatusCode.Gone, - subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, - activityId: Guid.NewGuid().ToString(), - requestCharge: default); + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition); - int moveNextAsyncCounter = 0; - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( - numItems, - multiPartition: true, - failureConfigs: new FlakyDocumentContainer.FailureConfigs( - inject429s: false, - injectEmptyPages: false, - shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter != 1 ? null : goneException))); + TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( + documentContainer: inMemoryCollection, + sqlQuerySpec: new SqlQuerySpec(query), + targetRange: targetRange, + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), + partitionKey: null, + cancellationToken: default, + continuationToken: continuationToken); - IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + return monadicQueryPipelineStage.Succeeded; + } - while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) - { - moveNextAsyncCounter++; - try - { - TryCatch tryGetPage = queryPipelineStage.Current; - tryGetPage.ThrowIfFailed(); + private static async Task CreateOptimisticDirectExecutionPipelineStateAsync(DocumentContainer documentContainer, string query, CosmosElement continuationToken) + { + List targetRanges = await documentContainer.GetFeedRangesAsync( + trace: NoOpTrace.Singleton, + cancellationToken: default); - documents.AddRange(tryGetPage.Result.Documents); - } - catch - { - // check if gone exception is handled properly - Assert.IsTrue(queryPipelineStage.Current.Failed); - Assert.AreEqual(queryPipelineStage.Current.InnerMostException.Message, errorMessage); - Assert.AreEqual(((CosmosException)queryPipelineStage.Current.InnerMostException).StatusCode, System.Net.HttpStatusCode.Gone); - break; - } - } - // Once fallback plan is implemented, this test should be able to return all 100 documents - Assert.AreEqual(documents.Count, 10); + FeedRangeEpk firstRange = targetRanges[0]; + + TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: new SqlQuerySpec(query), + targetRange: firstRange, + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), + partitionKey: null, + cancellationToken: default, + continuationToken: continuationToken); + + Assert.IsTrue(monadicQueryPipelineStage.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicQueryPipelineStage.Result; + + return queryPipelineStage; } - // Start with ODE pipeline and then switch to parallel in a single partition scenario. - // This is to simulate whether a fallback solution would work and provide client with all the required documents - [TestMethod] - public async Task TestPipelineSwitchForContinuationTokenOnSinglePartitionAsync() + private static async Task CreateParallelCrossPartitionPipelineStateAsync(IDocumentContainer documentContainer, string query, CosmosElement continuationToken) { - int numItems = 100; - string query = "SELECT * FROM c"; - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); - IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); - List documents = new List(); - int continuationTokenCount = 0; + List targetRanges = await documentContainer.GetFeedRangesAsync( + trace: NoOpTrace.Singleton, + cancellationToken: default); - while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) + if (continuationToken is CosmosObject continuationTokenObject) { - TryCatch tryGetPage = queryPipelineStage.Current; - tryGetPage.ThrowIfFailed(); + bool canGetContinuationToken = continuationTokenObject.TryGetValue("OptimisticDirectExecutionToken", out CosmosElement parallelContinuationToken); + Assert.IsTrue(canGetContinuationToken); - documents.AddRange(tryGetPage.Result.Documents); + CosmosArray cosmosElementParallelContinuationToken = CosmosArray.Create(parallelContinuationToken); + continuationToken = cosmosElementParallelContinuationToken; + } - if (tryGetPage.Result.State == null) - { - break; - } - else - { - queryPipelineStage = await CreateParallelCrossPartitionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); - } + TryCatch monadicQueryPipelineStage = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: new SqlQuerySpec(query), + targetRanges: targetRanges, + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), + partitionKey: null, + maxConcurrency: 10, + prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, + cancellationToken: default, + continuationToken: continuationToken); - continuationTokenCount++; - } + Assert.IsTrue(monadicQueryPipelineStage.Succeeded); + IQueryPipelineStage queryPipelineStage = monadicQueryPipelineStage.Result; - Assert.AreEqual(continuationTokenCount, 10); - Assert.AreEqual(documents.Count, 100); + return queryPipelineStage; } - // Start with ODE pipeline and then switch to parallel in a cross partition scenario. - // This is to simulate whether a fallback solution would work and provide client with all the required documents - [TestMethod] - public async Task TestPipelineSwitchForContinuationTokenOnMultiplePartitionsAsync() + private async Task GetResultsFromHydratingPipelineWithContinuation(int numItems, bool isMultiPartition, bool isODE, string query) { - int numItems = 100; - string query = "SELECT * FROM c"; - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: true); - IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: isMultiPartition); + IQueryPipelineStage queryPipelineStage = isODE + ? await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null) + : await CreateParallelCrossPartitionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + List documents = new List(); int continuationTokenCount = 0; @@ -357,147 +293,81 @@ public async Task TestPipelineSwitchForContinuationTokenOnMultiplePartitionsAsyn } else { - queryPipelineStage = await CreateParallelCrossPartitionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); + if (isODE) + { + queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); + } + else { + List parallelTest = new List + { + CreateInput( + description: @"Cross partition continuation token", + query: query, + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: Cosmos.PartitionKey.Null, + continuationToken: tryGetPage.Result.State.Value) + }; + + this.ExecuteTestSuite(parallelTest); + } } continuationTokenCount++; } - Assert.AreEqual(continuationTokenCount, 15); - Assert.AreEqual(documents.Count, 100); + Assert.AreEqual(documents.Count, numItems); + return true; } - // The reason we have the below test is to show the missing capabilities of the OptimisticDirectExecution pipeline. - // Currently this pipeline cannot handle distributed queries as it does not have the logic to sum up the values it gets from the backend in partial results. - // This functionality is available for other pipelines such as the ParallelCrossPartitionQueryPipelineStage as evident below - [TestMethod] - public async Task TestPipelinesForDistributedQueryAsync() + // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing + // TODO: With the addition of the merge/split support, this queryPipelineStage should be able to return all documents regardless of a gone exception happening + private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) { int numItems = 100; - string query = "SELECT VALUE COUNT(1) FROM c"; - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: true); - IQueryPipelineStage optimisticDirectExecutionQueryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); - IQueryPipelineStage parallelQueryPipelineStage = await CreateParallelCrossPartitionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); - int documentCountOptimisticPipeline = 0; - int documentCountParallelPipeline = 0; - - List queryPipelineStages = new List - { - optimisticDirectExecutionQueryPipelineStage, - parallelQueryPipelineStage - }; + string query = "SELECT * FROM c"; + List documents = new List(); + string errorMessage = $"Epk Range: Partition does not exist at the given range."; + CosmosException goneException = new CosmosException( + message: errorMessage, + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: Guid.NewGuid().ToString(), + requestCharge: default); - List documentPipelinesCount = new List - { - documentCountOptimisticPipeline, - documentCountParallelPipeline - }; + int moveNextAsyncCounter = 0; + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( + numItems, + multiPartition: isMultiPartition, + failureConfigs: new FlakyDocumentContainer.FailureConfigs( + inject429s: false, + injectEmptyPages: false, + shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter != 1 ? null : goneException))); - for (int i = 0; i < queryPipelineStages.Count(); i++) + IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { - while (await queryPipelineStages[i].MoveNextAsync(NoOpTrace.Singleton)) + moveNextAsyncCounter++; + try { - TryCatch tryGetPage = queryPipelineStages[i].Current; + TryCatch tryGetPage = queryPipelineStage.Current; tryGetPage.ThrowIfFailed(); - documentPipelinesCount[i] += Int32.Parse(tryGetPage.Result.Documents[0].ToString()); - - if (tryGetPage.Result.State == null) - { - break; - } - else - { - queryPipelineStages[i] = queryPipelineStages[i].Equals(optimisticDirectExecutionQueryPipelineStage) - ? await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value) - : await CreateParallelCrossPartitionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); - } + documents.AddRange(tryGetPage.Result.Documents); } - } - - documentCountOptimisticPipeline = documentPipelinesCount[0]; - documentCountParallelPipeline = documentPipelinesCount[1]; - int countDifference = documentCountParallelPipeline - documentCountOptimisticPipeline; - - Assert.AreNotEqual(documentCountOptimisticPipeline, documentCountParallelPipeline, countDifference.ToString()); - Assert.AreEqual(documentCountOptimisticPipeline, 24); - Assert.AreEqual(documentCountParallelPipeline, 100); - Assert.AreEqual(countDifference, 76); - Assert.AreEqual(documentCountParallelPipeline, numItems); - } - - private static async Task> ODEMonadicCreate(DocumentContainer documentContainer, string query, FeedRangeEpk targetRange, CosmosElement continuationToken) - { - TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( - documentContainer: documentContainer, - sqlQuerySpec: new SqlQuerySpec(query), - targetRange: targetRange, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), - partitionKey: null, - cancellationToken: default, - continuationToken: continuationToken); - - return monadicQueryPipelineStage; - } - - private static async Task CreateOptimisticDirectExecutionPipelineStateAsync(DocumentContainer documentContainer, string query, CosmosElement continuationToken) - { - List targetRanges = await documentContainer.GetFeedRangesAsync( - trace: NoOpTrace.Singleton, - cancellationToken: default); - FeedRangeEpk firstRange = null; - - if (targetRanges.Count > 1) - { - for (int i = 0; i < targetRanges.Capacity; i++) + catch { - if (targetRanges[i].Range.Min == "") - { - firstRange = targetRanges[i]; - } + Assert.IsTrue(queryPipelineStage.Current.Failed); + Assert.AreEqual(queryPipelineStage.Current.InnerMostException.Message, errorMessage); + Assert.AreEqual(((CosmosException)queryPipelineStage.Current.InnerMostException).StatusCode, System.Net.HttpStatusCode.Gone); + break; } } - else - { - firstRange = targetRanges[0]; - } - TryCatch monadicQueryPipelineStage = await ODEMonadicCreate(documentContainer, query, firstRange, continuationToken); - - Assert.IsTrue(monadicQueryPipelineStage.Succeeded); - IQueryPipelineStage queryPipelineStage = monadicQueryPipelineStage.Result; - - return queryPipelineStage; - } - - private static async Task CreateParallelCrossPartitionPipelineStateAsync(IDocumentContainer documentContainer, string query, CosmosElement continuationToken) - { - List targetRanges = await documentContainer.GetFeedRangesAsync( - trace: NoOpTrace.Singleton, - cancellationToken: default); - - if (continuationToken is CosmosObject) - { - ((CosmosObject)continuationToken).TryGetValue("OptimisticDirectExecutionToken", out CosmosElement parallelContinuationToken); - CosmosArray cosmosElementParallelContinuationToken = CosmosArray.Create(parallelContinuationToken); - continuationToken = cosmosElementParallelContinuationToken; - } - - TryCatch monadicQueryPipelineStage = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( - documentContainer: documentContainer, - sqlQuerySpec: new SqlQuerySpec(query), - targetRanges: targetRanges, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), - partitionKey: null, - maxConcurrency: 10, - prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, - cancellationToken: default, - continuationToken: continuationToken); - - Assert.IsTrue(monadicQueryPipelineStage.Succeeded); - IQueryPipelineStage queryPipelineStage = monadicQueryPipelineStage.Result; + // Once fallback plan is implemented, this test should be able to return all 100 documents + Assert.AreEqual(documents.Count, 10); - return queryPipelineStage; + return true; } private static async Task CreateDocumentContainerAsync( @@ -523,7 +393,9 @@ private static async Task CreateDocumentContainerAsync( DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); - for (int i = 0; i < 2; i++) + int exponentPartitionKeyRanges = 2; // a value of 2 would lead to 4 partitions (2 * 2). 4 partitions are used because theyre easy to manage + demonstrates multi partition use case + + for (int i = 0; i < exponentPartitionKeyRanges; i++) { IReadOnlyList ranges = await documentContainer.GetFeedRangesAsync( trace: NoOpTrace.Singleton, @@ -562,12 +434,13 @@ private static OptimisticDirectExecutionTestInput CreateInput( string query, bool expectedOptimisticDirectExecution, string partitionKeyPath, - string partitionKeyValue) + string partitionKeyValue, + CosmosElement continuationToken = null) { PartitionKeyBuilder pkBuilder = new PartitionKeyBuilder(); pkBuilder.Add(partitionKeyValue); - return CreateInput(description, query, expectedOptimisticDirectExecution, partitionKeyPath, pkBuilder.Build()); + return CreateInput(description, query, expectedOptimisticDirectExecution, partitionKeyPath, pkBuilder.Build(), continuationToken); } private static OptimisticDirectExecutionTestInput CreateInput( @@ -575,9 +448,10 @@ private static OptimisticDirectExecutionTestInput CreateInput( string query, bool expectedOptimisticDirectExecution, string partitionKeyPath, - Cosmos.PartitionKey partitionKeyValue) + Cosmos.PartitionKey partitionKeyValue, + CosmosElement continuationToken = null) { - return new OptimisticDirectExecutionTestInput(description, query, new SqlQuerySpec(query), expectedOptimisticDirectExecution, partitionKeyPath, partitionKeyValue); + return new OptimisticDirectExecutionTestInput(description, query, new SqlQuerySpec(query), expectedOptimisticDirectExecution, partitionKeyPath, partitionKeyValue, continuationToken); } private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(string querySpecJsonString, PartitionKeyDefinition pkDefinition) @@ -586,7 +460,7 @@ private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(st querySpecJsonString: querySpecJsonString, partitionKeyDefinition: pkDefinition, requireFormattableOrderByQuery: true, - isContinuationExpected: false, + isContinuationExpected: true, allowNonValueAggregateQuery: true, hasLogicalPartitionKey: false, allowDCount: true, @@ -605,14 +479,14 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect // gets query context string databaseId = "db1234"; - string resourceLink = string.Format("dbs/{0}/colls", databaseId); + string resourceLink = $"dbs/{databaseId}/colls"; CosmosQueryContextCore cosmosQueryContextCore = new CosmosQueryContextCore( client: new TestCosmosQueryClient(), resourceTypeEnum: Documents.ResourceType.Document, operationType: Documents.OperationType.Query, resourceType: typeof(QueryResponseCore), resourceLink: resourceLink, - isContinuationExpected: false, + isContinuationExpected: true, allowNonValueAggregateQuery: true, useSystemPrefix: false, correlatedActivityId: Guid.NewGuid()); @@ -636,7 +510,7 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( sqlQuerySpec: sqlQuerySpec, - initialUserContinuationToken: null, + initialUserContinuationToken: input.ContinuationToken, initialFeedRange: null, maxConcurrency: queryRequestOptions.MaxConcurrency, maxItemCount: queryRequestOptions.MaxItemCount, @@ -697,6 +571,7 @@ public sealed class OptimisticDirectExecutionTestInput : BaselineTestInput internal bool ExpectedOptimisticDirectExecution { get; set; } internal PartitionKeyRangeIdentity PartitionKeyRangeId { get; set; } internal string Query { get; set; } + internal CosmosElement ContinuationToken { get; set; } internal OptimisticDirectExecutionTestInput( string description, @@ -704,7 +579,8 @@ internal OptimisticDirectExecutionTestInput( SqlQuerySpec sqlQuerySpec, bool expectedOptimisticDirectExecution, string partitionKeyPath, - Cosmos.PartitionKey partitionKeyValue) + Cosmos.PartitionKey partitionKeyValue, + CosmosElement continuationToken = null) : base(description) { this.PartitionKeyDefinition = new PartitionKeyDefinition() @@ -720,6 +596,7 @@ internal OptimisticDirectExecutionTestInput( this.ExpectedOptimisticDirectExecution = expectedOptimisticDirectExecution; this.Query = query; this.PartitionKeyValue = partitionKeyValue; + this.ContinuationToken = continuationToken; } public override void SerializeAsXml(XmlWriter xmlWriter) From 924b24a7bb9e62711ae6aaf1baf2cedb03f52358 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 9 Nov 2022 13:57:49 -0800 Subject: [PATCH 04/37] Removed ParalleContEvocation test. Fixed comments --- ...misticDirectExecutionQueryBaselineTests.cs | 165 ++++++++---------- 1 file changed, 71 insertions(+), 94 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 5c65ea9356..dcd1a6621c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -40,25 +40,25 @@ public void PositiveOptimisticDirectExecutionOutput() List testVariations = new List { CreateInput( - description: @"Partition Key + Value and Distinct", - query: "SELECT DISTINCT c.key FROM c", - expectedOptimisticDirectExecution: true, - partitionKeyPath: @"/pk", - partitionKeyValue: @"/value"), + description: @"Partition Key + Value and Distinct", + query: "SELECT DISTINCT c.key FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: @"value"), CreateInput( - description: @"Partition Key + Value and Min Aggregate", - query: "SELECT VALUE MIN(c.key) FROM c", - expectedOptimisticDirectExecution: true, - partitionKeyPath: @"/pk", - partitionKeyValue: @"/value"), + description: @"Partition Key + Value and Min Aggregate", + query: "SELECT VALUE MIN(c.key) FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: @"value"), CreateInput( - description: @"Partition Key + Value Fields", - query: "SELECT c.key FROM c", - expectedOptimisticDirectExecution: true, - partitionKeyPath: @"/pk", - partitionKeyValue: @"/value"), + description: @"Partition Key + Value Fields", + query: "SELECT c.key FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: @"value"), }; this.ExecuteTestSuite(testVariations); } @@ -70,25 +70,25 @@ public void NegativeOptimisticDirectExecutionOutput() List testVariations = new List { CreateInput( - description: @"Null Partition Key Value", - query: "SELECT * FROM c", - expectedOptimisticDirectExecution: false, - partitionKeyPath: @"/pk", - partitionKeyValue: Cosmos.PartitionKey.Null), + description: @"Null Partition Key Value", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: Cosmos.PartitionKey.Null), CreateInput( - description: @"None Partition Key Value", - query: "SELECT * FROM c", - expectedOptimisticDirectExecution: false, - partitionKeyPath: @"/pk", - partitionKeyValue: Cosmos.PartitionKey.None), + description: @"None Partition Key Value", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: Cosmos.PartitionKey.None), CreateInput( - description: @"C# Null Partition Key Value", - query: "SELECT * FROM c", - expectedOptimisticDirectExecution: false, - partitionKeyPath: @"/pk", - partitionKeyValue: null), + description: @"C# Null Partition Key Value", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: null), }; this.ExecuteTestSuite(testVariations); } @@ -112,21 +112,18 @@ public async Task TestMonadicCreateODEPipeline() string query = "SELECT * FROM c"; // null continuation token - Assert.IsTrue(await HasMonadicCreateSucceeded(numItems, multiPartition, query, targetRange: FeedRangeEpk.FullRange, continuationToken: null)); + Assert.IsTrue(await TryMonadicCreate(numItems, multiPartition, query, targetRange: FeedRangeEpk.FullRange, continuationToken: null)); // default continuation token - Assert.IsTrue(await HasMonadicCreateSucceeded(numItems, multiPartition, query, targetRange: FeedRangeEpk.FullRange, continuationToken: default)); + Assert.IsTrue(await TryMonadicCreate(numItems, multiPartition, query, targetRange: FeedRangeEpk.FullRange, continuationToken: default)); - Range range = new Documents.Routing.Range("A", "B", isMinInclusive: true, isMaxInclusive: false); - ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( - token: "asdf", - range: range); + CosmosElement cosmosElementContinuationToken = CosmosElement.Parse( + "{\"OptimisticDirectExecutionToken\":{\"token\":\"{\\\"resourceId\\\":\\\"AQAAAMmFOw8LAAAAAAAAAA==\\\",\\\"skipCount\\\":1}\"," + + "\"range\":{\"min\":\"\",\"max\":\"FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF\"}}}"); + Range range = new Documents.Routing.Range("", "FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF", isMinInclusive: true, isMaxInclusive: false); - OptimisticDirectExecutionContinuationToken token = new OptimisticDirectExecutionContinuationToken(parallelContinuationToken); - CosmosElement cosmosElementContinuationToken = OptimisticDirectExecutionContinuationToken.ToCosmosElement(token); - // single continuation token - Assert.IsTrue(await HasMonadicCreateSucceeded(numItems, multiPartition, query, targetRange: new FeedRangeEpk(range), continuationToken: cosmosElementContinuationToken)); + Assert.IsTrue(await TryMonadicCreate(numItems, multiPartition, query, targetRange: new FeedRangeEpk(range), continuationToken: cosmosElementContinuationToken)); } // test checks that the pipeline can take a query to the backend and returns its associated document(s). @@ -154,22 +151,14 @@ public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() [TestMethod] public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() { - Assert.IsTrue(await this.GetResultsFromHydratingPipelineWithContinuation( - numItems: 100, - isMultiPartition: false, - isODE: true, - query: "SELECT * FROM c")); - } + int numItems = 100; + int result = await this.GetResultsFromHydratingPipelineWithContinuation( + numItems: numItems, + isMultiPartition: false, + isODE: true, + query: "SELECT * FROM c"); - // test checks that ODE pipeline does not get invoked when CrossPartition query is run - [TestMethod] - public async Task TestParallelContTokenEvocationAsync() - { - Assert.IsTrue(await this.GetResultsFromHydratingPipelineWithContinuation( - numItems: 100, - isMultiPartition: true, - isODE: false, - query: "SELECT * FROM c")); + Assert.AreEqual(result, numItems); } // test to check if pipeline handles a 410 exception properly and returns all the documents. @@ -187,19 +176,18 @@ public async Task TestPipelineForGoneExceptionOnSingleAndMultiplePartitionAsync( [TestMethod] public async Task TestPipelinesForDistributedQueryAsync() { - bool result = false; - try - { - result = await this.GetResultsFromHydratingPipelineWithContinuation(100, false, true, "SELECT AVG(c) FROM c"); - } - catch - { - // Coming into catch because doc count Assert failed. Once aggregate logic is added, this should not go into Catch - Assert.IsFalse(result); - } + int numItems = 100; + int result = await this.GetResultsFromHydratingPipelineWithContinuation( + numItems: numItems, + isMultiPartition: false, + isODE: true, + query: "SELECT AVG(c) FROM c"); + + // TODO: These values will not equal each other until aggregate logic is added + Assert.AreNotEqual(result, numItems); } - private static async Task HasMonadicCreateSucceeded(int numItems, bool multiPartition, string query, FeedRangeEpk targetRange, CosmosElement continuationToken) + private static async Task TryMonadicCreate(int numItems, bool multiPartition, string query, FeedRangeEpk targetRange, CosmosElement continuationToken) { DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition); @@ -270,13 +258,11 @@ private static async Task CreateParallelCrossPartitionPipel return queryPipelineStage; } - private async Task GetResultsFromHydratingPipelineWithContinuation(int numItems, bool isMultiPartition, bool isODE, string query) + private async Task GetResultsFromHydratingPipelineWithContinuation(int numItems, bool isMultiPartition, bool isODE, string query) { DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: isMultiPartition); - IQueryPipelineStage queryPipelineStage = isODE - ? await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null) - : await CreateParallelCrossPartitionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); - + IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + List documents = new List(); int continuationTokenCount = 0; @@ -293,31 +279,13 @@ private async Task GetResultsFromHydratingPipelineWithContinuation(int num } else { - if (isODE) - { - queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); - } - else { - List parallelTest = new List - { - CreateInput( - description: @"Cross partition continuation token", - query: query, - expectedOptimisticDirectExecution: false, - partitionKeyPath: @"/pk", - partitionKeyValue: Cosmos.PartitionKey.Null, - continuationToken: tryGetPage.Result.State.Value) - }; - - this.ExecuteTestSuite(parallelTest); - } + queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); } continuationTokenCount++; } - Assert.AreEqual(documents.Count, numItems); - return true; + return documents.Count; } // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing @@ -359,14 +327,12 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa { Assert.IsTrue(queryPipelineStage.Current.Failed); Assert.AreEqual(queryPipelineStage.Current.InnerMostException.Message, errorMessage); - Assert.AreEqual(((CosmosException)queryPipelineStage.Current.InnerMostException).StatusCode, System.Net.HttpStatusCode.Gone); break; } } // Once fallback plan is implemented, this test should be able to return all 100 documents Assert.AreEqual(documents.Count, 10); - return true; } @@ -395,9 +361,11 @@ private static async Task CreateDocumentContainerAsync( int exponentPartitionKeyRanges = 2; // a value of 2 would lead to 4 partitions (2 * 2). 4 partitions are used because theyre easy to manage + demonstrates multi partition use case + IReadOnlyList ranges; + for (int i = 0; i < exponentPartitionKeyRanges; i++) { - IReadOnlyList ranges = await documentContainer.GetFeedRangesAsync( + ranges = await documentContainer.GetFeedRangesAsync( trace: NoOpTrace.Singleton, cancellationToken: default); @@ -412,6 +380,14 @@ private static async Task CreateDocumentContainerAsync( await documentContainer.RefreshProviderAsync(NoOpTrace.Singleton, cancellationToken: default); } + ranges = await documentContainer.GetFeedRangesAsync( + trace: NoOpTrace.Singleton, + cancellationToken: default); + + int rangeCount = multiPartition == true ? 4 : 1; + + Assert.AreEqual(rangeCount, ranges.Count); + for (int i = 0; i < numItems; i++) { // Insert an item @@ -528,6 +504,7 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect cosmosQueryContextCore, inputParameters, NoOpTrace.Singleton); + bool result = queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton).Result; if (input.ExpectedOptimisticDirectExecution) From 095a44d6e0003960a7b0dadeab206f0f2b854def Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 9 Nov 2022 14:48:40 -0800 Subject: [PATCH 05/37] Removed CreateParallelCrossPartitionPipelineStateAsync() as it is not being used anymore --- ...misticDirectExecutionQueryBaselineTests.cs | 36 +------------------ 1 file changed, 1 insertion(+), 35 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index dcd1a6621c..247487da75 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -155,7 +155,6 @@ public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() int result = await this.GetResultsFromHydratingPipelineWithContinuation( numItems: numItems, isMultiPartition: false, - isODE: true, query: "SELECT * FROM c"); Assert.AreEqual(result, numItems); @@ -180,7 +179,6 @@ public async Task TestPipelinesForDistributedQueryAsync() int result = await this.GetResultsFromHydratingPipelineWithContinuation( numItems: numItems, isMultiPartition: false, - isODE: true, query: "SELECT AVG(c) FROM c"); // TODO: These values will not equal each other until aggregate logic is added @@ -226,39 +224,7 @@ private static async Task CreateOptimisticDirectExecutionPi return queryPipelineStage; } - private static async Task CreateParallelCrossPartitionPipelineStateAsync(IDocumentContainer documentContainer, string query, CosmosElement continuationToken) - { - List targetRanges = await documentContainer.GetFeedRangesAsync( - trace: NoOpTrace.Singleton, - cancellationToken: default); - - if (continuationToken is CosmosObject continuationTokenObject) - { - bool canGetContinuationToken = continuationTokenObject.TryGetValue("OptimisticDirectExecutionToken", out CosmosElement parallelContinuationToken); - Assert.IsTrue(canGetContinuationToken); - - CosmosArray cosmosElementParallelContinuationToken = CosmosArray.Create(parallelContinuationToken); - continuationToken = cosmosElementParallelContinuationToken; - } - - TryCatch monadicQueryPipelineStage = ParallelCrossPartitionQueryPipelineStage.MonadicCreate( - documentContainer: documentContainer, - sqlQuerySpec: new SqlQuerySpec(query), - targetRanges: targetRanges, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), - partitionKey: null, - maxConcurrency: 10, - prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, - cancellationToken: default, - continuationToken: continuationToken); - - Assert.IsTrue(monadicQueryPipelineStage.Succeeded); - IQueryPipelineStage queryPipelineStage = monadicQueryPipelineStage.Result; - - return queryPipelineStage; - } - - private async Task GetResultsFromHydratingPipelineWithContinuation(int numItems, bool isMultiPartition, bool isODE, string query) + private async Task GetResultsFromHydratingPipelineWithContinuation(int numItems, bool isMultiPartition, string query) { DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: isMultiPartition); IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); From 9115fcfe344891cc81ac95680a7939aa1d80a49b Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Thu, 10 Nov 2022 01:37:56 -0800 Subject: [PATCH 06/37] Removed while loop in CreateDocumentContainerAsync() --- .../OptimisticDirectExecutionQueryBaselineTests.cs | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 247487da75..f3fee023f7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -358,14 +358,8 @@ private static async Task CreateDocumentContainerAsync( { // Insert an item CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); - while (true) - { - TryCatch monadicCreateRecord = await documentContainer.MonadicCreateItemAsync(item, cancellationToken: default); - if (monadicCreateRecord.Succeeded) - { - break; - } - } + TryCatch monadicCreateRecord = await documentContainer.MonadicCreateItemAsync(item, cancellationToken: default); + Assert.IsTrue(monadicCreateRecord.Succeeded); } return documentContainer; @@ -436,7 +430,6 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect // gets input parameters QueryRequestOptions queryRequestOptions = new QueryRequestOptions { - MaxBufferedItemCount = 7000, TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) }; From c0549096f9cfb359facb7fb223e0055a2ae5eb70 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Thu, 10 Nov 2022 16:56:12 -0800 Subject: [PATCH 07/37] Fixed comments. --- .../CosmosQueryExecutionContextFactory.cs | 4 +-- ...misticDirectExecutionQueryBaselineTests.cs | 28 ++++++++----------- 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 21b512166e..cf1c9c6b0d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -265,7 +265,7 @@ private static async Task> TryCreateCoreContextAsy } } - public static async Task> TryCreateFromPartitionedQueryExecutionInfoAsync( + private static async Task> TryCreateFromPartitionedQueryExecutionInfoAsync( DocumentContainer documentContainer, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, ContainerQueryProperties containerQueryProperties, @@ -536,7 +536,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx return targetRanges; } - public static void SetTestInjectionPipelineType(InputParameters inputParameters, string pipelineType) + private static void SetTestInjectionPipelineType(InputParameters inputParameters, string pipelineType) { TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats; if (responseStats != null) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index f3fee023f7..a57f3aff6e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -114,9 +114,6 @@ public async Task TestMonadicCreateODEPipeline() // null continuation token Assert.IsTrue(await TryMonadicCreate(numItems, multiPartition, query, targetRange: FeedRangeEpk.FullRange, continuationToken: null)); - // default continuation token - Assert.IsTrue(await TryMonadicCreate(numItems, multiPartition, query, targetRange: FeedRangeEpk.FullRange, continuationToken: default)); - CosmosElement cosmosElementContinuationToken = CosmosElement.Parse( "{\"OptimisticDirectExecutionToken\":{\"token\":\"{\\\"resourceId\\\":\\\"AQAAAMmFOw8LAAAAAAAAAA==\\\",\\\"skipCount\\\":1}\"," + "\"range\":{\"min\":\"\",\"max\":\"FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF\"}}}"); @@ -152,7 +149,7 @@ public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() { int numItems = 100; - int result = await this.GetResultsFromHydratingPipelineWithContinuation( + int result = await this.CreateOptimisticPipelineAndDrainAsync( numItems: numItems, isMultiPartition: false, query: "SELECT * FROM c"); @@ -176,12 +173,11 @@ public async Task TestPipelineForGoneExceptionOnSingleAndMultiplePartitionAsync( public async Task TestPipelinesForDistributedQueryAsync() { int numItems = 100; - int result = await this.GetResultsFromHydratingPipelineWithContinuation( + int result = await this.CreateOptimisticPipelineAndDrainAsync( numItems: numItems, isMultiPartition: false, query: "SELECT AVG(c) FROM c"); - // TODO: These values will not equal each other until aggregate logic is added Assert.AreNotEqual(result, numItems); } @@ -207,7 +203,7 @@ private static async Task CreateOptimisticDirectExecutionPi trace: NoOpTrace.Singleton, cancellationToken: default); - FeedRangeEpk firstRange = targetRanges[0]; + FeedRangeEpk firstRange = targetRanges[0]; // only one range is taken because ODE pipeline can only accept one range TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, @@ -215,8 +211,8 @@ private static async Task CreateOptimisticDirectExecutionPi targetRange: firstRange, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, - cancellationToken: default, - continuationToken: continuationToken); + continuationToken: continuationToken, + cancellationToken: default); Assert.IsTrue(monadicQueryPipelineStage.Succeeded); IQueryPipelineStage queryPipelineStage = monadicQueryPipelineStage.Result; @@ -224,7 +220,7 @@ private static async Task CreateOptimisticDirectExecutionPi return queryPipelineStage; } - private async Task GetResultsFromHydratingPipelineWithContinuation(int numItems, bool isMultiPartition, string query) + private async Task CreateOptimisticPipelineAndDrainAsync(int numItems, bool isMultiPartition, string query) { DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: isMultiPartition); IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); @@ -254,9 +250,9 @@ private async Task GetResultsFromHydratingPipelineWithContinuation(int numI return documents.Count; } - // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing - // TODO: With the addition of the merge/split support, this queryPipelineStage should be able to return all documents regardless of a gone exception happening - private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) + // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing + // TODO: With the addition of the merge/split support, this queryPipelineStage should be able to return all documents regardless of a gone exception happening + private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) { int numItems = 100; string query = "SELECT * FROM c"; @@ -266,7 +262,7 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa message: errorMessage, statusCode: System.Net.HttpStatusCode.Gone, subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, - activityId: Guid.NewGuid().ToString(), + activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", requestCharge: default); int moveNextAsyncCounter = 0; @@ -350,7 +346,7 @@ private static async Task CreateDocumentContainerAsync( trace: NoOpTrace.Singleton, cancellationToken: default); - int rangeCount = multiPartition == true ? 4 : 1; + int rangeCount = multiPartition ? 4 : 1; Assert.AreEqual(rangeCount, ranges.Count); @@ -516,7 +512,7 @@ internal OptimisticDirectExecutionTestInput( bool expectedOptimisticDirectExecution, string partitionKeyPath, Cosmos.PartitionKey partitionKeyValue, - CosmosElement continuationToken = null) + CosmosElement continuationToken) : base(description) { this.PartitionKeyDefinition = new PartitionKeyDefinition() From bfb41106fbb8fda9336d8f4330cc547512f5f396 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Fri, 11 Nov 2022 16:41:01 -0800 Subject: [PATCH 08/37] Updated ExecuteGoneExceptionOnODEPipeline() --- ...misticDirectExecutionQueryBaselineTests.cs | 66 +++++++++++-------- 1 file changed, 38 insertions(+), 28 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index bb4352555c..c8afaf813b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -40,22 +40,22 @@ public void PositiveOptimisticDirectExecutionOutput() List testVariations = new List { CreateInput( - description: @"Partition Key + Value and Distinct", - query: "SELECT DISTINCT c.key FROM c", + description: @"Single Partition Key and Distinct", + query: "SELECT DISTINCT c.age FROM c", expectedOptimisticDirectExecution: true, partitionKeyPath: @"/pk", partitionKeyValue: @"value"), CreateInput( - description: @"Partition Key + Value and Min Aggregate", - query: "SELECT VALUE MIN(c.key) FROM c", + description: @"Single Partition Key and Min Aggregate", + query: "SELECT VALUE MIN(c.age) FROM c", expectedOptimisticDirectExecution: true, partitionKeyPath: @"/pk", partitionKeyValue: @"value"), CreateInput( - description: @"Partition Key + Value Fields", - query: "SELECT c.key FROM c", + description: @"Single Partition Key and Value Field", + query: "SELECT c.age FROM c", expectedOptimisticDirectExecution: true, partitionKeyPath: @"/pk", partitionKeyValue: @"value"), @@ -121,6 +121,8 @@ public async Task TestMonadicCreateODEPipeline() // single continuation token Assert.IsTrue(await TryMonadicCreate(numItems, multiPartition, query, targetRange: new FeedRangeEpk(range), continuationToken: cosmosElementContinuationToken)); + + //TODO: Add non ODE continuation token case } // test checks that the pipeline can take a query to the backend and returns its associated document(s). @@ -141,7 +143,7 @@ public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() documentCountInSinglePartition += Int32.Parse(tryGetPage.Result.Documents[0].ToString()); } - Assert.AreEqual(documentCountInSinglePartition, 100); + Assert.AreEqual(100, documentCountInSinglePartition); } // test checks that the pipeline can take a query to the backend and returns its associated document(s) + continuation token. @@ -150,9 +152,10 @@ public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() { int numItems = 100; int result = await this.CreateOptimisticPipelineAndDrainAsync( - numItems: numItems, - isMultiPartition: false, - query: "SELECT * FROM c"); + numItems: numItems, + isMultiPartition: false, + query: "SELECT * FROM c", + expectedContinuationTokenCount: 10); Assert.AreEqual(result, numItems); } @@ -170,13 +173,14 @@ public async Task TestPipelineForGoneExceptionOnSingleAndMultiplePartitionAsync( // Currently this pipeline cannot handle distributed queries as it does not have the logic to sum up the values it gets from the backend in partial results. // This functionality is available for other pipelines such as the ParallelCrossPartitionQueryPipelineStage. [TestMethod] - public async Task TestPipelinesForDistributedQueryAsync() + public async Task TestPipelineForDistributedQueryAsync() { int numItems = 100; int result = await this.CreateOptimisticPipelineAndDrainAsync( numItems: numItems, isMultiPartition: false, - query: "SELECT AVG(c) FROM c"); + query: "SELECT AVG(c) FROM c", + expectedContinuationTokenCount: 0); Assert.AreNotEqual(result, numItems); } @@ -203,7 +207,8 @@ private static async Task CreateOptimisticDirectExecutionPi trace: NoOpTrace.Singleton, cancellationToken: default); - FeedRangeEpk firstRange = targetRanges[0]; // only one range is taken because ODE pipeline can only accept one range + // only one range is taken because ODE pipeline can only accept one range + FeedRangeEpk firstRange = targetRanges[0]; TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, @@ -220,7 +225,7 @@ private static async Task CreateOptimisticDirectExecutionPi return queryPipelineStage; } - private async Task CreateOptimisticPipelineAndDrainAsync(int numItems, bool isMultiPartition, string query) + private async Task CreateOptimisticPipelineAndDrainAsync(int numItems, bool isMultiPartition, string query, int expectedContinuationTokenCount) { DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: isMultiPartition); IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); @@ -247,6 +252,7 @@ private async Task CreateOptimisticPipelineAndDrainAsync(int numItems, bool continuationTokenCount++; } + Assert.AreEqual(expectedContinuationTokenCount, continuationTokenCount); return documents.Count; } @@ -257,15 +263,16 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa int numItems = 100; string query = "SELECT * FROM c"; List documents = new List(); - string errorMessage = $"Epk Range: Partition does not exist at the given range."; + string goneExceptionMessage = $"Epk Range: Partition does not exist at the given range."; CosmosException goneException = new CosmosException( - message: errorMessage, + message: goneExceptionMessage, statusCode: System.Net.HttpStatusCode.Gone, subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", requestCharge: default); int moveNextAsyncCounter = 0; + bool caughtGoneException = false; DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( numItems, multiPartition: isMultiPartition, @@ -278,23 +285,25 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { moveNextAsyncCounter++; - try - { - TryCatch tryGetPage = queryPipelineStage.Current; - tryGetPage.ThrowIfFailed(); + TryCatch tryGetPage = queryPipelineStage.Current; - documents.AddRange(tryGetPage.Result.Documents); - } - catch + if (tryGetPage.Failed == true) { - Assert.IsTrue(queryPipelineStage.Current.Failed); - Assert.AreEqual(queryPipelineStage.Current.InnerMostException.Message, errorMessage); - break; + string errorRecieved = tryGetPage.Exception.InnerException.Message; + if (errorRecieved.Equals(goneExceptionMessage)) + { + caughtGoneException = true; + break; + } } + + documents.AddRange(tryGetPage.Result.Documents); } // Once fallback plan is implemented, this test should be able to return all 100 documents - Assert.AreEqual(documents.Count, 10); + Assert.AreEqual(10, documents.Count); + Assert.IsTrue(caughtGoneException); + return true; } @@ -321,7 +330,8 @@ private static async Task CreateDocumentContainerAsync( DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); - int exponentPartitionKeyRanges = 2; // a value of 2 would lead to 4 partitions (2 * 2). 4 partitions are used because theyre easy to manage + demonstrates multi partition use case + // a value of 2 would lead to 4 partitions (2 * 2). 4 partitions are used because they're easy to manage + demonstrates multi partition use case + int exponentPartitionKeyRanges = 2; IReadOnlyList ranges; From ea0aac47fedbd1aee87d61e5cfeab864cd564ba7 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Sat, 12 Nov 2022 13:14:56 -0800 Subject: [PATCH 09/37] Added type Assert for ExecuteGoneExceptionOnODEPipeline() --- .../Query/OptimisticDirectExecutionQueryBaselineTests.cs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index c8afaf813b..9d34c456ea 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -29,6 +29,7 @@ using System.IO; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; + using Microsoft.IdentityModel.Tokens; [TestClass] public class OptimisticDirectExecutionQueryBaselineTests : BaselineTests @@ -157,7 +158,7 @@ public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() query: "SELECT * FROM c", expectedContinuationTokenCount: 10); - Assert.AreEqual(result, numItems); + Assert.AreEqual(numItems, result); } // test to check if pipeline handles a 410 exception properly and returns all the documents. @@ -182,7 +183,8 @@ public async Task TestPipelineForDistributedQueryAsync() query: "SELECT AVG(c) FROM c", expectedContinuationTokenCount: 0); - Assert.AreNotEqual(result, numItems); + //TODO: Add validation for actual value of average + Assert.AreEqual(1, result); } private static async Task TryMonadicCreate(int numItems, bool multiPartition, string query, FeedRangeEpk targetRange, CosmosElement continuationToken) @@ -290,6 +292,8 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa if (tryGetPage.Failed == true) { string errorRecieved = tryGetPage.Exception.InnerException.Message; + Assert.AreEqual(goneException.GetType(), tryGetPage.Exception.InnerException.GetType()); + if (errorRecieved.Equals(goneExceptionMessage)) { caughtGoneException = true; From b886f40eac36f6cabb1f87ecade7e7c7e5fb4386 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 16 Nov 2022 23:52:51 -0800 Subject: [PATCH 10/37] Replaced try-catch with if statement in MoveNextAsync() --- .../CosmosQueryExecutionContextFactory.cs | 8 +- ...misticDirectExecutionQueryPipelineStage.cs | 354 ++++++++++++------ .../Pagination/FlakyDocumentContainer.cs | 35 +- ...misticDirectExecutionQueryBaselineTests.cs | 75 +++- 4 files changed, 336 insertions(+), 136 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 4e6bacb145..0e4e78153c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -150,6 +150,7 @@ private static async Task> TryCreateCoreContextAsy return OptimisticDirectExecutionContext( documentContainer, + cosmosQueryContext, inputParameters, targetRange, cancellationToken); @@ -319,6 +320,7 @@ private static async Task> TryCreateFromPartitione tryCreatePipelineStage = CosmosQueryExecutionContextFactory.OptimisticDirectExecutionContext( documentContainer, + cosmosQueryContext, inputParameters, targetRange, cancellationToken); @@ -379,6 +381,7 @@ private static async Task> TryCreateFromPartitione private static TryCatch OptimisticDirectExecutionContext( DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, Documents.PartitionKeyRange targetRange, CancellationToken cancellationToken) @@ -386,11 +389,10 @@ private static TryCatch OptimisticDirectExecutionContext( // Return a OptimisticDirectExecution context return OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, - sqlQuerySpec: inputParameters.SqlQuerySpec, + cosmosQueryContext: cosmosQueryContext, + inputParameters: inputParameters, targetRange: new FeedRangeEpk(targetRange.ToRange()), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), - partitionKey: inputParameters.PartitionKey, - continuationToken: inputParameters.InitialUserContinuationToken, cancellationToken: cancellationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 9f86b0ad8b..b88b214b14 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -8,172 +8,306 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu using System.Collections.Generic; using System.Diagnostics; using System.Linq; + using System.Reflection; + using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.ChangeFeed; + using Microsoft.Azure.Cosmos.ChangeFeed.Pagination; + using Microsoft.Azure.Cosmos.Core.Collections; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Tracing; + using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory; using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage { - private readonly QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator; - - private OptimisticDirectExecutionQueryPipelineStage( - QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator) - { - this.queryPartitionRangePageAsyncEnumerator = queryPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(queryPartitionRangePageAsyncEnumerator)); - } - - public TryCatch Current { get; private set; } - - public ValueTask DisposeAsync() + //private readonly Func queryPipelineStageFactory; + private TryCatch inner; + //private TryCatch currentPipeline; + private QueryPartitionRangePageAsyncEnumerator partitionPageEnumerator; + private static DocumentContainer mDocumentContainer; + private static CosmosQueryContext mCosmosQueryContext; + private static InputParameters mInputParameters; + private static FeedRangeEpk mTargetRange; + + public OptimisticDirectExecutionQueryPipelineStage(TryCatch monadicCreate) { - return this.queryPartitionRangePageAsyncEnumerator.DisposeAsync(); + this.inner = monadicCreate; } - public void SetCancellationToken(CancellationToken cancellationToken) - { - this.queryPartitionRangePageAsyncEnumerator.SetCancellationToken(cancellationToken); - } + public TryCatch Current { get; private set; } + public ValueTask DisposeAsync() => default; + public async ValueTask MoveNextAsync(ITrace trace) { - if (trace == null) + bool result = await this.inner.Result.MoveNextAsync(trace); + bool isGoneException = this.inner.Result.Current.Failed + && this.inner.Result.Current.InnerMostException is CosmosException exception + && exception.StatusCode == System.Net.HttpStatusCode.Gone + && exception.SubStatusCode == 1002; + + if (isGoneException) { - throw new ArgumentNullException(nameof(trace)); + IQueryPipelineStage pipelineResult = CosmosQueryExecutionContextFactory.Create(mDocumentContainer, mCosmosQueryContext, mInputParameters, trace); + this.inner = TryCatch.FromResult(pipelineResult); + result = await this.inner.Result.MoveNextAsync(trace); + this.Current = this.inner.Result.Current; + return result; } - - if (!await this.queryPartitionRangePageAsyncEnumerator.MoveNextAsync(trace)) + else { - this.Current = default; - return false; + this.Current = this.inner.Result.Current; } - TryCatch partitionPage = this.queryPartitionRangePageAsyncEnumerator.Current; - if (partitionPage.Failed) + if (this.inner.Result.Current.Result?.State?.Value != null) { - this.Current = partitionPage; - return true; + AddContTokenToParams(this.inner.Result.Current.Result.State.Value); } - QueryPage backendQueryPage = partitionPage.Result; + return result; + } + + public void SetCancellationToken(CancellationToken cancellationToken) + { + this.partitionPageEnumerator = new QueryPartitionRangePageAsyncEnumerator( + mDocumentContainer, + mInputParameters.SqlQuerySpec, + new FeedRangeState(mTargetRange, default), + mInputParameters.PartitionKey, + new QueryPaginationOptions(pageSizeHint: mInputParameters.MaxItemCount), + cancellationToken); + + OptimisticDirectExecutionQueryPipelineImpl ode = new OptimisticDirectExecutionQueryPipelineImpl(this.partitionPageEnumerator); + ode.SetCancellationToken(cancellationToken); + } + + public static TryCatch MonadicCreate( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + InputParameters inputParameters, + FeedRangeEpk targetRange, + QueryPaginationOptions queryPaginationOptions, + CancellationToken cancellationToken) + { + mDocumentContainer = documentContainer; + mCosmosQueryContext = cosmosQueryContext; + mInputParameters = inputParameters; + mTargetRange = targetRange; + + TryCatch monadicCreate = OptimisticDirectExecutionQueryPipelineImpl.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: inputParameters.SqlQuerySpec, + targetRange: targetRange, + queryPaginationOptions: queryPaginationOptions, + partitionKey: inputParameters.PartitionKey, + continuationToken: inputParameters.InitialUserContinuationToken, + cancellationToken: cancellationToken); - QueryState queryState; - if (backendQueryPage.State == null) + if (monadicCreate.Failed) { - queryState = null; + return TryCatch.FromException( + new Exception( + message: "Failed when creating MonadicCreate", + innerException: monadicCreate.Exception)); } - else + + OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(monadicCreate); + return TryCatch.FromResult(odePipelineStageMonadicCreate); + } + + private static void AddContTokenToParams(CosmosElement continuationToken) + { + if (continuationToken is CosmosObject) { - QueryState backendQueryState = backendQueryPage.State; - ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( - token: (backendQueryState?.Value as CosmosString)?.Value, - range: ((FeedRangeEpk)this.queryPartitionRangePageAsyncEnumerator.FeedRangeState.FeedRange).Range); - - OptimisticDirectExecutionContinuationToken optimisticDirectExecutionContinuationToken = new OptimisticDirectExecutionContinuationToken(parallelContinuationToken); - CosmosElement cosmosElementContinuationToken = OptimisticDirectExecutionContinuationToken.ToCosmosElement(optimisticDirectExecutionContinuationToken); - queryState = new QueryState(cosmosElementContinuationToken); + ((CosmosObject)continuationToken).TryGetValue("OptimisticDirectExecutionToken", out CosmosElement parallelContinuationToken); + CosmosArray cosmosElementParallelContinuationToken = CosmosArray.Create(parallelContinuationToken); + continuationToken = cosmosElementParallelContinuationToken; } - QueryPage queryPage = new QueryPage( - backendQueryPage.Documents, - backendQueryPage.RequestCharge, - backendQueryPage.ActivityId, - backendQueryPage.ResponseLengthInBytes, - backendQueryPage.CosmosQueryExecutionInfo, - disallowContinuationTokenMessage: null, - backendQueryPage.AdditionalHeaders, - queryState); - - this.Current = TryCatch.FromResult(queryPage); - return true; + InputParameters inputParams = new InputParameters( + mInputParameters.SqlQuerySpec, + continuationToken, + mInputParameters.InitialFeedRange, + mInputParameters.MaxConcurrency, + mInputParameters.MaxItemCount, + mInputParameters.MaxBufferedItemCount, + mInputParameters.PartitionKey, + mInputParameters.Properties, + mInputParameters.PartitionedQueryExecutionInfo, + mInputParameters.ExecutionEnvironment, + mInputParameters.ReturnResultsInDeterministicOrder, + mInputParameters.ForcePassthrough, + mInputParameters.TestInjections); + + mInputParameters = inputParams; } - public static TryCatch MonadicCreate( - IDocumentContainer documentContainer, - SqlQuerySpec sqlQuerySpec, - FeedRangeEpk targetRange, - Cosmos.PartitionKey? partitionKey, - QueryPaginationOptions queryPaginationOptions, - CosmosElement continuationToken, - CancellationToken cancellationToken) + private class OptimisticDirectExecutionQueryPipelineImpl : IQueryPipelineStage { - if (targetRange == null) + private readonly QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator; + + internal OptimisticDirectExecutionQueryPipelineImpl( + QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator) { - throw new ArgumentNullException(nameof(targetRange)); + this.queryPartitionRangePageAsyncEnumerator = queryPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(queryPartitionRangePageAsyncEnumerator)); } - - TryCatch> monadicExtractState; - if (continuationToken == null) + + public TryCatch Current { get; private set; } + + public ValueTask DisposeAsync() { - FeedRangeState getState = new (targetRange, (QueryState)null); - monadicExtractState = TryCatch>.FromResult(getState); + return this.queryPartitionRangePageAsyncEnumerator.DisposeAsync(); } - else + + public void SetCancellationToken(CancellationToken cancellationToken) { - monadicExtractState = MonadicExtractState(continuationToken, targetRange); + this.queryPartitionRangePageAsyncEnumerator.SetCancellationToken(cancellationToken); } - - if (monadicExtractState.Failed) + + public async ValueTask MoveNextAsync(ITrace trace) { - return TryCatch.FromException(monadicExtractState.Exception); - } + if (trace == null) + { + throw new ArgumentNullException(nameof(trace)); + } - FeedRangeState feedRangeState = monadicExtractState.Result; + if (!await this.queryPartitionRangePageAsyncEnumerator.MoveNextAsync(trace)) + { + this.Current = default; + return false; + } - QueryPartitionRangePageAsyncEnumerator partitionPageEnumerator = new QueryPartitionRangePageAsyncEnumerator( - documentContainer, - sqlQuerySpec, - feedRangeState, - partitionKey, - queryPaginationOptions, - cancellationToken); + TryCatch partitionPage = this.queryPartitionRangePageAsyncEnumerator.Current; + if (partitionPage.Failed) + { + this.Current = partitionPage; + return false; + } - OptimisticDirectExecutionQueryPipelineStage stage = new OptimisticDirectExecutionQueryPipelineStage(partitionPageEnumerator); - return TryCatch.FromResult(stage); - } + QueryPage backendQueryPage = partitionPage.Result; - private static TryCatch> MonadicExtractState( - CosmosElement continuationToken, - FeedRangeEpk range) - { - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); + QueryState queryState; + if (backendQueryPage.State == null) + { + queryState = null; + } + else + { + QueryState backendQueryState = backendQueryPage.State; + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( + token: (backendQueryState?.Value as CosmosString)?.Value, + range: ((FeedRangeEpk)this.queryPartitionRangePageAsyncEnumerator.FeedRangeState.FeedRange).Range); + + OptimisticDirectExecutionContinuationToken optimisticDirectExecutionContinuationToken = new OptimisticDirectExecutionContinuationToken(parallelContinuationToken); + CosmosElement cosmosElementContinuationToken = OptimisticDirectExecutionContinuationToken.ToCosmosElement(optimisticDirectExecutionContinuationToken); + queryState = new QueryState(cosmosElementContinuationToken); + } + + QueryPage queryPage = new QueryPage( + backendQueryPage.Documents, + backendQueryPage.RequestCharge, + backendQueryPage.ActivityId, + backendQueryPage.ResponseLengthInBytes, + backendQueryPage.CosmosQueryExecutionInfo, + disallowContinuationTokenMessage: null, + backendQueryPage.AdditionalHeaders, + queryState); + + this.Current = TryCatch.FromResult(queryPage); + return true; } - TryCatch tryCreateContinuationToken = OptimisticDirectExecutionContinuationToken.TryCreateFromCosmosElement(continuationToken); - if (tryCreateContinuationToken.Failed) + public static TryCatch MonadicCreate( + IDocumentContainer documentContainer, + SqlQuerySpec sqlQuerySpec, + FeedRangeEpk targetRange, + Cosmos.PartitionKey? partitionKey, + QueryPaginationOptions queryPaginationOptions, + CosmosElement continuationToken, + CancellationToken cancellationToken) { - return TryCatch>.FromException(tryCreateContinuationToken.Exception); - } + if (targetRange == null) + { + throw new ArgumentNullException(nameof(targetRange)); + } - TryCatch> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping( - range, - tryCreateContinuationToken.Result); + TryCatch> monadicExtractState; + if (continuationToken == null) + { + FeedRangeState getState = new (targetRange, (QueryState)null); + monadicExtractState = TryCatch>.FromResult(getState); + } + else + { + monadicExtractState = MonadicExtractState(continuationToken, targetRange); + } - if (partitionMappingMonad.Failed) - { - return TryCatch>.FromException( - partitionMappingMonad.Exception); + if (monadicExtractState.Failed) + { + return TryCatch.FromException(monadicExtractState.Exception); + } + + FeedRangeState feedRangeState = monadicExtractState.Result; + + QueryPartitionRangePageAsyncEnumerator partitionPageEnumerator = new QueryPartitionRangePageAsyncEnumerator( + documentContainer, + sqlQuerySpec, + feedRangeState, + partitionKey, + queryPaginationOptions, + cancellationToken); + + OptimisticDirectExecutionQueryPipelineImpl stage = new OptimisticDirectExecutionQueryPipelineImpl(partitionPageEnumerator); + return TryCatch.FromResult(stage); } - - PartitionMapping partitionMapping = partitionMappingMonad.Result; - - KeyValuePair kvpRange = new KeyValuePair( - partitionMapping.TargetMapping.Keys.First(), - partitionMapping.TargetMapping.Values.First()); - FeedRangeState feedRangeState = new FeedRangeState(kvpRange.Key, kvpRange.Value?.Token != null ? new QueryState(CosmosString.Create(kvpRange.Value.Token.Token)) : null); - - return TryCatch>.FromResult(feedRangeState); + private static TryCatch> MonadicExtractState( + CosmosElement continuationToken, + FeedRangeEpk range) + { + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + + TryCatch tryCreateContinuationToken = OptimisticDirectExecutionContinuationToken.TryCreateFromCosmosElement(continuationToken); + if (tryCreateContinuationToken.Failed) + { + return TryCatch>.FromException(tryCreateContinuationToken.Exception); + } + + TryCatch> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping( + range, + tryCreateContinuationToken.Result); + + if (partitionMappingMonad.Failed) + { + return TryCatch>.FromException( + partitionMappingMonad.Exception); + } + + PartitionMapping partitionMapping = partitionMappingMonad.Result; + + KeyValuePair kvpRange = new KeyValuePair( + partitionMapping.TargetMapping.Keys.First(), + partitionMapping.TargetMapping.Values.First()); + + FeedRangeState feedRangeState = new FeedRangeState(kvpRange.Key, kvpRange.Value?.Token != null ? new QueryState(CosmosString.Create(kvpRange.Value.Token.Token)) : null); + + return TryCatch>.FromResult(feedRangeState); + } } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs index cc038ff4ed..3cca5432e7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs @@ -27,7 +27,9 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination /// internal sealed class FlakyDocumentContainer : IMonadicDocumentContainer { - private readonly FailureConfigs failureConfigs; + private FailureConfigs failureConfigs; + private bool isODETest; + private int iterationCount = 0; private readonly Random random; private static class Throttle @@ -64,10 +66,12 @@ private static class Throttle public FlakyDocumentContainer( IMonadicDocumentContainer documentContainer, - FailureConfigs failureConfigs) + FailureConfigs failureConfigs, + bool isODETest = false) { this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); this.failureConfigs = failureConfigs ?? throw new ArgumentNullException(nameof(failureConfigs)); + this.isODETest = isODETest; this.random = new Random(); } @@ -175,10 +179,25 @@ public async Task> MonadicQueryAsync( state: feedRangeState.State ?? StateForStartedButNoDocumentsReturned)); } - Exception failure = await this.ShouldReturnFailure(); - if (failure != null) + if (this.isODETest) { - return TryCatch.FromException(failure); + this.iterationCount++; + if (this.iterationCount == 2) + { + Exception failure = await this.ShouldReturnFailure(); + if (failure != null) + { + return TryCatch.FromException(failure); + } + } + } + else + { + Exception failure = await this.ShouldReturnFailure(); + if (failure != null) + { + return TryCatch.FromException(failure); + } } return await this.documentContainer.MonadicQueryAsync( @@ -190,8 +209,8 @@ public async Task> MonadicQueryAsync( } public async Task> MonadicChangeFeedAsync( - FeedRangeState feedRangeState, - ChangeFeedPaginationOptions changeFeedPaginationOptions, + FeedRangeState feedRangeState, + ChangeFeedPaginationOptions changeFeedPaginationOptions, ITrace trace, CancellationToken cancellationToken) { @@ -289,7 +308,7 @@ public sealed class FailureConfigs public delegate Task ShouldReturnFailure(); public FailureConfigs( - bool inject429s, + bool inject429s, bool injectEmptyPages, Exception throwException = null, Exception returnFailure = null) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 9d34c456ea..a6488abe34 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -191,34 +191,36 @@ private static async Task TryMonadicCreate(int numItems, bool multiPartiti { DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition); + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = await CreateInputParamsAndQueryContext(query, continuationToken); + TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( documentContainer: inMemoryCollection, - sqlQuerySpec: new SqlQuerySpec(query), + cosmosQueryContext: cosmosQueryContextCore, + inputParameters: inputParameters, targetRange: targetRange, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), - partitionKey: null, - cancellationToken: default, - continuationToken: continuationToken); + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), + cancellationToken: default); return monadicQueryPipelineStage.Succeeded; } private static async Task CreateOptimisticDirectExecutionPipelineStateAsync(DocumentContainer documentContainer, string query, CosmosElement continuationToken) { + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = await CreateInputParamsAndQueryContext(query, continuationToken); + List targetRanges = await documentContainer.GetFeedRangesAsync( trace: NoOpTrace.Singleton, cancellationToken: default); - // only one range is taken because ODE pipeline can only accept one range - FeedRangeEpk firstRange = targetRanges[0]; + // only the first range is taken because ODE pipeline can only accept one range + FeedRangeEpk firstRange = targetRanges.Where(r => r.Range.Min == string.Empty || targetRanges.Count == 1).First(); TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, - sqlQuerySpec: new SqlQuerySpec(query), + cosmosQueryContext: cosmosQueryContextCore, + inputParameters: inputParameters, targetRange: firstRange, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), - partitionKey: null, - continuationToken: continuationToken, + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), cancellationToken: default); Assert.IsTrue(monadicQueryPipelineStage.Succeeded); @@ -258,6 +260,48 @@ private async Task CreateOptimisticPipelineAndDrainAsync(int numItems, bool return documents.Count; } + private static async Task> CreateInputParamsAndQueryContext(string query, CosmosElement continuationToken) + { + // gets input parameters + QueryRequestOptions queryRequestOptions = new QueryRequestOptions + { + MaxConcurrency = 0, + MaxItemCount = 10, + MaxBufferedItemCount = 7000, + TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) + }; + + CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( + sqlQuerySpec: new SqlQuerySpec(query), + initialUserContinuationToken: continuationToken, + initialFeedRange: null, + maxConcurrency: queryRequestOptions.MaxConcurrency, + maxItemCount: queryRequestOptions.MaxItemCount, + maxBufferedItemCount: queryRequestOptions.MaxBufferedItemCount, + partitionKey: null, + properties: queryRequestOptions.Properties, + partitionedQueryExecutionInfo: null, + executionEnvironment: null, + returnResultsInDeterministicOrder: null, + forcePassthrough: true, + testInjections: queryRequestOptions.TestSettings); + + string databaseId = "db1234"; + string resourceLink = string.Format("dbs/{0}/colls", databaseId); + CosmosQueryContextCore cosmosQueryContextCore = new CosmosQueryContextCore( + client: new TestCosmosQueryClient(), + resourceTypeEnum: Documents.ResourceType.Document, + operationType: Documents.OperationType.Query, + resourceType: typeof(QueryResponseCore), + resourceLink: resourceLink, + isContinuationExpected: false, + allowNonValueAggregateQuery: true, + useSystemPrefix: false, + correlatedActivityId: Guid.NewGuid()); + + return Tuple.Create(inputParameters, cosmosQueryContextCore); + } + // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing // TODO: With the addition of the merge/split support, this queryPipelineStage should be able to return all documents regardless of a gone exception happening private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) @@ -281,7 +325,7 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa failureConfigs: new FlakyDocumentContainer.FailureConfigs( inject429s: false, injectEmptyPages: false, - shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter != 1 ? null : goneException))); + shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter == 1 ? goneException : null))); IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) @@ -305,8 +349,9 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa } // Once fallback plan is implemented, this test should be able to return all 100 documents - Assert.AreEqual(10, documents.Count); - Assert.IsTrue(caughtGoneException); + Assert.AreEqual(100, documents.Count); + // cannot check for assert because the gone exception never reaches till here + //Assert.IsTrue(caughtGoneException); return true; } @@ -329,7 +374,7 @@ private static async Task CreateDocumentContainerAsync( IMonadicDocumentContainer monadicDocumentContainer = new InMemoryContainer(partitionKeyDefinition); if (failureConfigs != null) { - monadicDocumentContainer = new FlakyDocumentContainer(monadicDocumentContainer, failureConfigs); + monadicDocumentContainer = new FlakyDocumentContainer(monadicDocumentContainer, failureConfigs, isODETest: true); } DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); From 04389e0d2a38d2eeb0d522ea595921205605aa5f Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Tue, 29 Nov 2022 16:18:03 -0800 Subject: [PATCH 11/37] Added delegate to access TryCreateCoreContextAsync() --- .../CosmosQueryExecutionContextFactory.cs | 21 +- ...misticDirectExecutionQueryPipelineStage.cs | 118 +++-- ...misticDirectExecutionQueryBaselineTests.cs | 409 +++++++++--------- 3 files changed, 274 insertions(+), 274 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 0e4e78153c..6b0121dac8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -389,10 +389,22 @@ private static TryCatch OptimisticDirectExecutionContext( // Return a OptimisticDirectExecution context return OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( documentContainer: documentContainer, - cosmosQueryContext: cosmosQueryContext, inputParameters: inputParameters, targetRange: new FeedRangeEpk(targetRange.ToRange()), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), + queryPipelineStage: (InputParameters updatedInputParameters) => + { + // Query Iterator requires that the creation of the query context is deferred until fallback pipeline is called + Task> tryCreateContext = + CosmosQueryExecutionContextFactory.TryCreateCoreContextAsync( + documentContainer, + cosmosQueryContext, + updatedInputParameters, + NoOpTrace.Singleton, + default); + + return tryCreateContext; + }, cancellationToken: cancellationToken); } @@ -615,7 +627,7 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP ContainerQueryProperties containerQueryProperties, ITrace trace) { - if (inputParameters.TestInjections == null || !inputParameters.TestInjections.EnableOptimisticDirectExecution) return null; + if (inputParameters.TestInjections == null || !inputParameters.TestInjections.EnableOptimisticDirectExecution || inputParameters.IsOdeFallBackPlan) return null; // case 1: Is query going to a single partition bool hasPartitionKey = inputParameters.PartitionKey.HasValue @@ -687,7 +699,8 @@ public InputParameters( ExecutionEnvironment? executionEnvironment, bool? returnResultsInDeterministicOrder, bool forcePassthrough, - TestInjections testInjections) + TestInjections testInjections, + bool isOdeFallBackPlan = false) { this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); this.InitialUserContinuationToken = initialUserContinuationToken; @@ -721,6 +734,7 @@ public InputParameters( this.ReturnResultsInDeterministicOrder = returnResultsInDeterministicOrder.GetValueOrDefault(InputParameters.DefaultReturnResultsInDeterministicOrder); this.ForcePassthrough = forcePassthrough; this.TestInjections = testInjections; + this.IsOdeFallBackPlan = isOdeFallBackPlan; } public SqlQuerySpec SqlQuerySpec { get; } @@ -736,6 +750,7 @@ public InputParameters( public bool ReturnResultsInDeterministicOrder { get; } public TestInjections TestInjections { get; } public bool ForcePassthrough { get; } + public bool IsOdeFallBackPlan { get; set; } } internal sealed class AggregateProjectionDetector diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index b88b214b14..336c5c167e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -8,17 +8,10 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu using System.Collections.Generic; using System.Diagnostics; using System.Linq; - using System.Reflection; - using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.ChangeFeed; - using Microsoft.Azure.Cosmos.ChangeFeed.Pagination; - using Microsoft.Azure.Cosmos.Core.Collections; using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Diagnostics; using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; @@ -32,21 +25,17 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage { - //private readonly Func queryPipelineStageFactory; + private readonly string optimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; private TryCatch inner; - //private TryCatch currentPipeline; - private QueryPartitionRangePageAsyncEnumerator partitionPageEnumerator; - private static DocumentContainer mDocumentContainer; - private static CosmosQueryContext mCosmosQueryContext; - private static InputParameters mInputParameters; - private static FeedRangeEpk mTargetRange; + private Func>> queryPipelineStageFactory; + private InputParameters inputParameters; - public OptimisticDirectExecutionQueryPipelineStage(TryCatch monadicCreate) + private OptimisticDirectExecutionQueryPipelineStage(TryCatch monadicCreate) { this.inner = monadicCreate; } - public TryCatch Current { get; private set; } + public TryCatch Current => this.inner.Result.Current; public ValueTask DisposeAsync() => default; @@ -57,56 +46,45 @@ public async ValueTask MoveNextAsync(ITrace trace) && this.inner.Result.Current.InnerMostException is CosmosException exception && exception.StatusCode == System.Net.HttpStatusCode.Gone && exception.SubStatusCode == 1002; - - if (isGoneException) + + if (result) { - IQueryPipelineStage pipelineResult = CosmosQueryExecutionContextFactory.Create(mDocumentContainer, mCosmosQueryContext, mInputParameters, trace); - this.inner = TryCatch.FromResult(pipelineResult); - result = await this.inner.Result.MoveNextAsync(trace); - this.Current = this.inner.Result.Current; - return result; + if (this.Current.Result?.State?.Value != null) + { + CosmosElement continuationToken = this.inner.Result.Current.Result.State.Value; + this.SaveContinuation(continuationToken); + return result; + } } else { - this.Current = this.inner.Result.Current; - } + if (isGoneException) + { + Debug.Assert(result != this.Current.Failed); - if (this.inner.Result.Current.Result?.State?.Value != null) - { - AddContTokenToParams(this.inner.Result.Current.Result.State.Value); + this.inputParameters.IsOdeFallBackPlan = true; + this.inner = await this.queryPipelineStageFactory.Invoke(this.inputParameters); + return await this.inner.Result.MoveNextAsync(trace); + } } - return result; + return result; } public void SetCancellationToken(CancellationToken cancellationToken) { - this.partitionPageEnumerator = new QueryPartitionRangePageAsyncEnumerator( - mDocumentContainer, - mInputParameters.SqlQuerySpec, - new FeedRangeState(mTargetRange, default), - mInputParameters.PartitionKey, - new QueryPaginationOptions(pageSizeHint: mInputParameters.MaxItemCount), - cancellationToken); - - OptimisticDirectExecutionQueryPipelineImpl ode = new OptimisticDirectExecutionQueryPipelineImpl(this.partitionPageEnumerator); - ode.SetCancellationToken(cancellationToken); + this.inner.Result.SetCancellationToken(cancellationToken); } public static TryCatch MonadicCreate( DocumentContainer documentContainer, - CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, FeedRangeEpk targetRange, QueryPaginationOptions queryPaginationOptions, + Func>> queryPipelineStage, CancellationToken cancellationToken) { - mDocumentContainer = documentContainer; - mCosmosQueryContext = cosmosQueryContext; - mInputParameters = inputParameters; - mTargetRange = targetRange; - - TryCatch monadicCreate = OptimisticDirectExecutionQueryPipelineImpl.MonadicCreate( + TryCatch pipelineStage = OptimisticDirectExecutionQueryPipelineImpl.MonadicCreate( documentContainer: documentContainer, sqlQuerySpec: inputParameters.SqlQuerySpec, targetRange: targetRange, @@ -115,43 +93,45 @@ public static TryCatch MonadicCreate( continuationToken: inputParameters.InitialUserContinuationToken, cancellationToken: cancellationToken); - if (monadicCreate.Failed) + if (pipelineStage.Failed) { - return TryCatch.FromException( - new Exception( - message: "Failed when creating MonadicCreate", - innerException: monadicCreate.Exception)); + return TryCatch.FromException(pipelineStage.Exception); } - OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(monadicCreate); + OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage) + { + queryPipelineStageFactory = queryPipelineStage, + inputParameters = inputParameters + }; + return TryCatch.FromResult(odePipelineStageMonadicCreate); } - private static void AddContTokenToParams(CosmosElement continuationToken) + private void SaveContinuation(CosmosElement continuationToken) { if (continuationToken is CosmosObject) { - ((CosmosObject)continuationToken).TryGetValue("OptimisticDirectExecutionToken", out CosmosElement parallelContinuationToken); + ((CosmosObject)continuationToken).TryGetValue(this.optimisticDirectExecutionToken, out CosmosElement parallelContinuationToken); CosmosArray cosmosElementParallelContinuationToken = CosmosArray.Create(parallelContinuationToken); continuationToken = cosmosElementParallelContinuationToken; } - + InputParameters inputParams = new InputParameters( - mInputParameters.SqlQuerySpec, + this.inputParameters.SqlQuerySpec, continuationToken, - mInputParameters.InitialFeedRange, - mInputParameters.MaxConcurrency, - mInputParameters.MaxItemCount, - mInputParameters.MaxBufferedItemCount, - mInputParameters.PartitionKey, - mInputParameters.Properties, - mInputParameters.PartitionedQueryExecutionInfo, - mInputParameters.ExecutionEnvironment, - mInputParameters.ReturnResultsInDeterministicOrder, - mInputParameters.ForcePassthrough, - mInputParameters.TestInjections); - - mInputParameters = inputParams; + this.inputParameters.InitialFeedRange, + this.inputParameters.MaxConcurrency, + this.inputParameters.MaxItemCount, + this.inputParameters.MaxBufferedItemCount, + this.inputParameters.PartitionKey, + this.inputParameters.Properties, + this.inputParameters.PartitionedQueryExecutionInfo, + this.inputParameters.ExecutionEnvironment, + this.inputParameters.ReturnResultsInDeterministicOrder, + this.inputParameters.ForcePassthrough, + this.inputParameters.TestInjections); + + this.inputParameters = inputParams; } private class OptimisticDirectExecutionQueryPipelineImpl : IQueryPipelineStage diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index a6488abe34..060b8da50e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -27,9 +27,6 @@ using System.Linq; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; using System.IO; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; - using Microsoft.IdentityModel.Tokens; [TestClass] public class OptimisticDirectExecutionQueryBaselineTests : BaselineTests @@ -38,6 +35,10 @@ public class OptimisticDirectExecutionQueryBaselineTests : BaselineTests testVariations = new List { CreateInput( @@ -60,6 +61,24 @@ public void PositiveOptimisticDirectExecutionOutput() expectedOptimisticDirectExecution: true, partitionKeyPath: @"/pk", partitionKeyValue: @"value"), + + CreateInput( + description: @"Single Partition Key and Value Field", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: "a", + continuationToken: null), + + CreateInput( + description: @"Single Partition Key and continuation token", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: "a", + continuationToken: cosmosElementContinuationToken), + + //TODO: Add non ODE continuation token case }; this.ExecuteTestSuite(testVariations); } @@ -97,51 +116,50 @@ public void NegativeOptimisticDirectExecutionOutput() // This test confirms that TestInjection.EnableOptimisticDirectExection is set to false from default. // Check test "TestPipelineForDistributedQueryAsync" to understand why this is done [TestMethod] - public async Task TestDefaultTestInjectionSettings() + public async Task TestDefaultTestInjectionSettingsAsync() { TestInjections testInjection = new TestInjections(simulate429s: false, simulateEmptyPages: false); Assert.AreEqual(testInjection.EnableOptimisticDirectExecution, false); } - [TestMethod] - [Owner("akotalwar")] - public async Task TestMonadicCreateODEPipeline() - { - int numItems = 10; - bool multiPartition = false; - string query = "SELECT * FROM c"; - - // null continuation token - Assert.IsTrue(await TryMonadicCreate(numItems, multiPartition, query, targetRange: FeedRangeEpk.FullRange, continuationToken: null)); - - CosmosElement cosmosElementContinuationToken = CosmosElement.Parse( - "{\"OptimisticDirectExecutionToken\":{\"token\":\"{\\\"resourceId\\\":\\\"AQAAAMmFOw8LAAAAAAAAAA==\\\",\\\"skipCount\\\":1}\"," + - "\"range\":{\"min\":\"\",\"max\":\"FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF\"}}}"); - Range range = new Documents.Routing.Range("", "FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF", isMinInclusive: true, isMaxInclusive: false); - - // single continuation token - Assert.IsTrue(await TryMonadicCreate(numItems, multiPartition, query, targetRange: new FeedRangeEpk(range), continuationToken: cosmosElementContinuationToken)); - - //TODO: Add non ODE continuation token case - } - - // test checks that the pipeline can take a query to the backend and returns its associated document(s). + // test checks that the pipeline can take a query to the backend and returns its associated document(s). [TestMethod] public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() { int numItems = 100; - string query = "SELECT VALUE COUNT(1) FROM c"; + OptimisticDirectExecutionTestInput input = CreateInput( + description: @"Single Partition Key and Value Field", + query: "SELECT VALUE COUNT(1) FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: "a"); + + QueryRequestOptions queryRequestOptions = new QueryRequestOptions + { + MaxConcurrency = 0, + MaxItemCount = 10, + MaxBufferedItemCount = 7000, + TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) + }; + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); - IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + IQueryPipelineStage queryPipelineStage = await CreateAndGetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); int documentCountInSinglePartition = 0; while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { + Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); + TryCatch tryGetPage = queryPipelineStage.Current; tryGetPage.ThrowIfFailed(); documentCountInSinglePartition += Int32.Parse(tryGetPage.Result.Documents[0].ToString()); + + if (tryGetPage.Result.State == null) + { + break; + } } Assert.AreEqual(100, documentCountInSinglePartition); @@ -152,10 +170,17 @@ public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() { int numItems = 100; + OptimisticDirectExecutionTestInput input = CreateInput( + description: @"Single Partition Key and Value Field", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: "a"); + int result = await this.CreateOptimisticPipelineAndDrainAsync( - numItems: numItems, - isMultiPartition: false, - query: "SELECT * FROM c", + input, + numItems: numItems, + isMultiPartition: false, expectedContinuationTokenCount: 10); Assert.AreEqual(numItems, result); @@ -177,137 +202,37 @@ public async Task TestPipelineForGoneExceptionOnSingleAndMultiplePartitionAsync( public async Task TestPipelineForDistributedQueryAsync() { int numItems = 100; + OptimisticDirectExecutionTestInput input = CreateInput( + description: @"Single Partition Key and Value Field", + query: "SELECT AVG(c) FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: "a"); + int result = await this.CreateOptimisticPipelineAndDrainAsync( - numItems: numItems, - isMultiPartition: false, - query: "SELECT AVG(c) FROM c", + input, + numItems: numItems, + isMultiPartition: false, expectedContinuationTokenCount: 0); //TODO: Add validation for actual value of average Assert.AreEqual(1, result); } - private static async Task TryMonadicCreate(int numItems, bool multiPartition, string query, FeedRangeEpk targetRange, CosmosElement continuationToken) - { - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition); - - (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = await CreateInputParamsAndQueryContext(query, continuationToken); - - TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( - documentContainer: inMemoryCollection, - cosmosQueryContext: cosmosQueryContextCore, - inputParameters: inputParameters, - targetRange: targetRange, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), - cancellationToken: default); - - return monadicQueryPipelineStage.Succeeded; - } - - private static async Task CreateOptimisticDirectExecutionPipelineStateAsync(DocumentContainer documentContainer, string query, CosmosElement continuationToken) - { - (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = await CreateInputParamsAndQueryContext(query, continuationToken); - - List targetRanges = await documentContainer.GetFeedRangesAsync( - trace: NoOpTrace.Singleton, - cancellationToken: default); - - // only the first range is taken because ODE pipeline can only accept one range - FeedRangeEpk firstRange = targetRanges.Where(r => r.Range.Min == string.Empty || targetRanges.Count == 1).First(); - - TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( - documentContainer: documentContainer, - cosmosQueryContext: cosmosQueryContextCore, - inputParameters: inputParameters, - targetRange: firstRange, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), - cancellationToken: default); - - Assert.IsTrue(monadicQueryPipelineStage.Succeeded); - IQueryPipelineStage queryPipelineStage = monadicQueryPipelineStage.Result; - - return queryPipelineStage; - } - - private async Task CreateOptimisticPipelineAndDrainAsync(int numItems, bool isMultiPartition, string query, int expectedContinuationTokenCount) - { - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: isMultiPartition); - IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); - - List documents = new List(); - int continuationTokenCount = 0; - - while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) - { - TryCatch tryGetPage = queryPipelineStage.Current; - tryGetPage.ThrowIfFailed(); - - documents.AddRange(tryGetPage.Result.Documents); - - if (tryGetPage.Result.State == null) - { - break; - } - else - { - queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: tryGetPage.Result.State.Value); - } - - continuationTokenCount++; - } - - Assert.AreEqual(expectedContinuationTokenCount, continuationTokenCount); - return documents.Count; - } - - private static async Task> CreateInputParamsAndQueryContext(string query, CosmosElement continuationToken) - { - // gets input parameters - QueryRequestOptions queryRequestOptions = new QueryRequestOptions - { - MaxConcurrency = 0, - MaxItemCount = 10, - MaxBufferedItemCount = 7000, - TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) - }; - - CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( - sqlQuerySpec: new SqlQuerySpec(query), - initialUserContinuationToken: continuationToken, - initialFeedRange: null, - maxConcurrency: queryRequestOptions.MaxConcurrency, - maxItemCount: queryRequestOptions.MaxItemCount, - maxBufferedItemCount: queryRequestOptions.MaxBufferedItemCount, - partitionKey: null, - properties: queryRequestOptions.Properties, - partitionedQueryExecutionInfo: null, - executionEnvironment: null, - returnResultsInDeterministicOrder: null, - forcePassthrough: true, - testInjections: queryRequestOptions.TestSettings); - - string databaseId = "db1234"; - string resourceLink = string.Format("dbs/{0}/colls", databaseId); - CosmosQueryContextCore cosmosQueryContextCore = new CosmosQueryContextCore( - client: new TestCosmosQueryClient(), - resourceTypeEnum: Documents.ResourceType.Document, - operationType: Documents.OperationType.Query, - resourceType: typeof(QueryResponseCore), - resourceLink: resourceLink, - isContinuationExpected: false, - allowNonValueAggregateQuery: true, - useSystemPrefix: false, - correlatedActivityId: Guid.NewGuid()); - - return Tuple.Create(inputParameters, cosmosQueryContextCore); - } - // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing // TODO: With the addition of the merge/split support, this queryPipelineStage should be able to return all documents regardless of a gone exception happening private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) { int numItems = 100; string query = "SELECT * FROM c"; + + OptimisticDirectExecutionTestInput input = CreateInput( + description: @"Single Partition Key and Value Field", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: "a"); + List documents = new List(); string goneExceptionMessage = $"Epk Range: Partition does not exist at the given range."; CosmosException goneException = new CosmosException( @@ -317,6 +242,14 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", requestCharge: default); + QueryRequestOptions queryRequestOptions = new QueryRequestOptions + { + MaxConcurrency = 0, + MaxItemCount = 10, + MaxBufferedItemCount = 7000, + TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) + }; + int moveNextAsyncCounter = 0; bool caughtGoneException = false; DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( @@ -327,22 +260,26 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa injectEmptyPages: false, shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter == 1 ? goneException : null))); - IQueryPipelineStage queryPipelineStage = await CreateOptimisticDirectExecutionPipelineStateAsync(inMemoryCollection, query, continuationToken: null); + IQueryPipelineStage queryPipelineStage = await CreateAndGetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { moveNextAsyncCounter++; + if (moveNextAsyncCounter == 1) + { + Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); + } + else + { + Assert.AreNotEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); + } + TryCatch tryGetPage = queryPipelineStage.Current; - if (tryGetPage.Failed == true) + if (tryGetPage.Failed) { - string errorRecieved = tryGetPage.Exception.InnerException.Message; - Assert.AreEqual(goneException.GetType(), tryGetPage.Exception.InnerException.GetType()); - - if (errorRecieved.Equals(goneExceptionMessage)) - { - caughtGoneException = true; - break; - } + // failure should never come till here. Should be handled before + Assert.Fail(); } documents.AddRange(tryGetPage.Result.Documents); @@ -350,12 +287,59 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa // Once fallback plan is implemented, this test should be able to return all 100 documents Assert.AreEqual(100, documents.Count); - // cannot check for assert because the gone exception never reaches till here - //Assert.IsTrue(caughtGoneException); return true; } + private async Task CreateOptimisticPipelineAndDrainAsync(OptimisticDirectExecutionTestInput input, int numItems, bool isMultiPartition, int expectedContinuationTokenCount) + { + QueryRequestOptions queryRequestOptions = new QueryRequestOptions + { + MaxConcurrency = 0, + MaxItemCount = 10, + MaxBufferedItemCount = 7000, + TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) + }; + + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: isMultiPartition); + IQueryPipelineStage queryPipelineStage = await CreateAndGetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); + + List documents = new List(); + int continuationTokenCount = 0; + + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) + { + Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); + + TryCatch tryGetPage = queryPipelineStage.Current; + tryGetPage.ThrowIfFailed(); + + documents.AddRange(tryGetPage.Result.Documents); + + if (tryGetPage.Result.State == null) + { + break; + } + else + { + input = CreateInput( + description: input.Description, + query: input.Query, + expectedOptimisticDirectExecution: input.ExpectedOptimisticDirectExecution, + partitionKeyPath: @"/pk", + partitionKeyValue: input.PartitionKeyValue, + continuationToken: tryGetPage.Result.State.Value); + + queryPipelineStage = await CreateAndGetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); + } + + continuationTokenCount++; + } + + Assert.AreEqual(expectedContinuationTokenCount, continuationTokenCount); + return documents.Count; + } + private static async Task CreateDocumentContainerAsync( int numItems, bool multiPartition, @@ -412,7 +396,7 @@ private static async Task CreateDocumentContainerAsync( for (int i = 0; i < numItems; i++) { // Insert an item - CosmosObject item = CosmosObject.Parse($"{{\"pk\" : {i} }}"); + CosmosObject item = CosmosObject.Parse($"{{\"pk\" : \"a\" }}"); TryCatch monadicCreateRecord = await documentContainer.MonadicCreateItemAsync(item, cancellationToken: default); Assert.IsTrue(monadicCreateRecord.Succeeded); } @@ -461,36 +445,64 @@ private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(st return tryGetQueryPlan.Result; } + private static async Task CreateAndGetOdePipelineAsync(OptimisticDirectExecutionTestInput input, DocumentContainer documentContainer, QueryRequestOptions queryRequestOptions) + { + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input, queryRequestOptions); + + IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( + documentContainer, + cosmosQueryContextCore, + inputParameters, + NoOpTrace.Singleton); + + Assert.IsNotNull(queryPipelineStage); + + return queryPipelineStage; + } + public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirectExecutionTestInput input) { // gets DocumentContainer IMonadicDocumentContainer monadicDocumentContainer = new InMemoryContainer(input.PartitionKeyDefinition); DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); - SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(input.Query); - - // gets query context - string databaseId = "db1234"; - string resourceLink = $"dbs/{databaseId}/colls"; - CosmosQueryContextCore cosmosQueryContextCore = new CosmosQueryContextCore( - client: new TestCosmosQueryClient(), - resourceTypeEnum: Documents.ResourceType.Document, - operationType: Documents.OperationType.Query, - resourceType: typeof(QueryResponseCore), - resourceLink: resourceLink, - isContinuationExpected: true, - allowNonValueAggregateQuery: true, - useSystemPrefix: false, - correlatedActivityId: Guid.NewGuid()); - - // gets input parameters QueryRequestOptions queryRequestOptions = new QueryRequestOptions { + MaxConcurrency = 0, + MaxItemCount = 10, + MaxBufferedItemCount = 7000, TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) }; + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input, queryRequestOptions); + + IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( + documentContainer, + cosmosQueryContextCore, + inputParameters, + NoOpTrace.Singleton); + + bool result = queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton).Result; + + if (input.ExpectedOptimisticDirectExecution) + { + Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); + } + else + { + Assert.AreNotEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); + } + + Assert.IsNotNull(queryPipelineStage); + Assert.IsTrue(result); + + return new OptimisticDirectExecutionTestOutput(input.ExpectedOptimisticDirectExecution); + } + + private static Tuple CreateInputParamsAndQueryContext(OptimisticDirectExecutionTestInput input, QueryRequestOptions queryRequestOptions) + { CosmosSerializerCore serializerCore = new(); - using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, Documents.ResourceType.Document)); + using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(new SqlQuerySpec(input.Query), Documents.ResourceType.Document)); string sqlQuerySpecJsonString = streamReader.ReadToEnd(); PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = GetPartitionedQueryExecutionInfo(sqlQuerySpecJsonString, input.PartitionKeyDefinition); @@ -500,7 +512,7 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect } CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( - sqlQuerySpec: sqlQuerySpec, + sqlQuerySpec: new SqlQuerySpec(input.Query), initialUserContinuationToken: input.ContinuationToken, initialFeedRange: null, maxConcurrency: queryRequestOptions.MaxConcurrency, @@ -514,27 +526,20 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect forcePassthrough: true, testInjections: queryRequestOptions.TestSettings); - IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( - documentContainer, - cosmosQueryContextCore, - inputParameters, - NoOpTrace.Singleton); - - bool result = queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton).Result; - - if (input.ExpectedOptimisticDirectExecution) - { - Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); - } - else - { - Assert.AreNotEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); - } - - Assert.IsNotNull(queryPipelineStage); - Assert.IsTrue(result); + string databaseId = "db1234"; + string resourceLink = $"dbs/{databaseId}/colls"; + CosmosQueryContextCore cosmosQueryContextCore = new CosmosQueryContextCore( + client: new TestCosmosQueryClient(), + resourceTypeEnum: Documents.ResourceType.Document, + operationType: Documents.OperationType.Query, + resourceType: typeof(QueryResponseCore), + resourceLink: resourceLink, + isContinuationExpected: true, + allowNonValueAggregateQuery: true, + useSystemPrefix: false, + correlatedActivityId: Guid.NewGuid()); - return new OptimisticDirectExecutionTestOutput(input.ExpectedOptimisticDirectExecution); + return Tuple.Create(inputParameters, cosmosQueryContextCore); } } From 4af203617fbdc6a4e8c3575be0ff2b04d2937938 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 30 Nov 2022 10:36:01 -0800 Subject: [PATCH 12/37] Added check to confirm Ode pipeline is not called in fallback plan --- ...ptimisticDirectExecutionQueryPipelineStage.cs | 16 +++++++++++++--- ...ptimisticDirectExecutionQueryBaselineTests.cs | 2 +- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 336c5c167e..508f95a63e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -51,7 +51,7 @@ public async ValueTask MoveNextAsync(ITrace trace) { if (this.Current.Result?.State?.Value != null) { - CosmosElement continuationToken = this.inner.Result.Current.Result.State.Value; + CosmosElement continuationToken = this.Current.Result.State.Value; this.SaveContinuation(continuationToken); return result; } @@ -61,10 +61,20 @@ public async ValueTask MoveNextAsync(ITrace trace) if (isGoneException) { Debug.Assert(result != this.Current.Failed); - this.inputParameters.IsOdeFallBackPlan = true; this.inner = await this.queryPipelineStageFactory.Invoke(this.inputParameters); - return await this.inner.Result.MoveNextAsync(trace); + // TODO: Failure check for this.inner + bool fallbackPipelineResult = await this.inner.Result.MoveNextAsync(trace); + + if (this.Current.Result?.State?.Value != null) + { + if (this.Current.Result.State.Value is CosmosObject) + { + // Fallback plan returned a Ode pipeline which is wrong + return false; + } + } + return fallbackPipelineResult; } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 060b8da50e..f1f064bd6f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -447,7 +447,7 @@ private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(st private static async Task CreateAndGetOdePipelineAsync(OptimisticDirectExecutionTestInput input, DocumentContainer documentContainer, QueryRequestOptions queryRequestOptions) { - (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input, queryRequestOptions); + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input, queryRequestOptions); IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( documentContainer, From 1257ad38583bb89f0c3bdf423b4f27079a541568 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 30 Nov 2022 14:46:27 -0800 Subject: [PATCH 13/37] Updated method name from OptimisticDirectExecutionContext() to TryCreateOptimisticDirectExecutionContext() --- .../CosmosQueryExecutionContextFactory.cs | 8 +- ...misticDirectExecutionQueryPipelineStage.cs | 73 +++++++++---------- 2 files changed, 38 insertions(+), 43 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 6b0121dac8..d1987fcd00 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -148,7 +148,7 @@ private static async Task> TryCreateCoreContextAsy // Test code added to confirm the correct pipeline is being utilized SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); - return OptimisticDirectExecutionContext( + return TryCreateOptimisticDirectExecutionContext( documentContainer, cosmosQueryContext, inputParameters, @@ -318,7 +318,7 @@ private static async Task> TryCreateFromPartitione { SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); - tryCreatePipelineStage = CosmosQueryExecutionContextFactory.OptimisticDirectExecutionContext( + tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreateOptimisticDirectExecutionContext( documentContainer, cosmosQueryContext, inputParameters, @@ -379,7 +379,7 @@ private static async Task> TryCreateFromPartitione return tryCreatePipelineStage; } - private static TryCatch OptimisticDirectExecutionContext( + private static TryCatch TryCreateOptimisticDirectExecutionContext( DocumentContainer documentContainer, CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, @@ -392,7 +392,7 @@ private static TryCatch OptimisticDirectExecutionContext( inputParameters: inputParameters, targetRange: new FeedRangeEpk(targetRange.ToRange()), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), - queryPipelineStage: (InputParameters updatedInputParameters) => + queryPipelineStage: (updatedInputParameters) => { // Query Iterator requires that the creation of the query context is deferred until fallback pipeline is called Task> tryCreateContext = diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 508f95a63e..8a27de8ac1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -25,65 +25,59 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage { - private readonly string optimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; - private TryCatch inner; - private Func>> queryPipelineStageFactory; + private const string optimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; + private readonly Func>> queryPipelineStageFactory; + private TryCatch innerQueryPipelineStage; private InputParameters inputParameters; - private OptimisticDirectExecutionQueryPipelineStage(TryCatch monadicCreate) + private OptimisticDirectExecutionQueryPipelineStage(TryCatch queryPipelineStage, Func>> queryPipelineStageFactory, InputParameters inputParameters) { - this.inner = monadicCreate; + this.innerQueryPipelineStage = queryPipelineStage; + this.queryPipelineStageFactory = queryPipelineStageFactory; + this.inputParameters = inputParameters; } - public TryCatch Current => this.inner.Result.Current; + public TryCatch Current => this.innerQueryPipelineStage.Result.Current; public ValueTask DisposeAsync() => default; public async ValueTask MoveNextAsync(ITrace trace) { - bool result = await this.inner.Result.MoveNextAsync(trace); - bool isGoneException = this.inner.Result.Current.Failed - && this.inner.Result.Current.InnerMostException is CosmosException exception + bool success = await this.innerQueryPipelineStage.Result.MoveNextAsync(trace); + bool isGoneException = this.Current.Failed + && this.Current.InnerMostException is CosmosException exception && exception.StatusCode == System.Net.HttpStatusCode.Gone && exception.SubStatusCode == 1002; - if (result) + if (success) { - if (this.Current.Result?.State?.Value != null) - { - CosmosElement continuationToken = this.Current.Result.State.Value; - this.SaveContinuation(continuationToken); - return result; - } + this.SaveContinuation(this.Current.Result?.State?.Value); } - else + else if (isGoneException) { - if (isGoneException) - { - Debug.Assert(result != this.Current.Failed); - this.inputParameters.IsOdeFallBackPlan = true; - this.inner = await this.queryPipelineStageFactory.Invoke(this.inputParameters); - // TODO: Failure check for this.inner - bool fallbackPipelineResult = await this.inner.Result.MoveNextAsync(trace); + this.inputParameters.IsOdeFallBackPlan = true; + this.innerQueryPipelineStage = await this.queryPipelineStageFactory(this.inputParameters); + + // TODO: Failure check for this.inner + bool fallbackPipelineSuccess = await this.innerQueryPipelineStage.Result.MoveNextAsync(trace); - if (this.Current.Result?.State?.Value != null) + if (this.Current.Result?.State?.Value != null) + { + if (this.Current.Result.State.Value is CosmosObject) { - if (this.Current.Result.State.Value is CosmosObject) - { - // Fallback plan returned a Ode pipeline which is wrong - return false; - } + // Fallback plan returned a Ode pipeline which is wrong + return false; } - return fallbackPipelineResult; } + return fallbackPipelineSuccess; } - return result; + return success; } public void SetCancellationToken(CancellationToken cancellationToken) { - this.inner.Result.SetCancellationToken(cancellationToken); + this.innerQueryPipelineStage.Result.SetCancellationToken(cancellationToken); } public static TryCatch MonadicCreate( @@ -108,20 +102,21 @@ public static TryCatch MonadicCreate( return TryCatch.FromException(pipelineStage.Exception); } - OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage) - { - queryPipelineStageFactory = queryPipelineStage, - inputParameters = inputParameters - }; + OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, queryPipelineStage, inputParameters); return TryCatch.FromResult(odePipelineStageMonadicCreate); } private void SaveContinuation(CosmosElement continuationToken) { + if (continuationToken == null) + { + return; + } + if (continuationToken is CosmosObject) { - ((CosmosObject)continuationToken).TryGetValue(this.optimisticDirectExecutionToken, out CosmosElement parallelContinuationToken); + ((CosmosObject)continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement parallelContinuationToken); CosmosArray cosmosElementParallelContinuationToken = CosmosArray.Create(parallelContinuationToken); continuationToken = cosmosElementParallelContinuationToken; } From bd4c195fd5c05ef3b189f2688b81d52d2b397b0f Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Thu, 1 Dec 2022 12:46:28 -0800 Subject: [PATCH 14/37] Using delegate instead of Func<>. --- .../CosmosQueryExecutionContextFactory.cs | 20 +++++++- ...misticDirectExecutionQueryPipelineStage.cs | 50 ++++++++----------- 2 files changed, 41 insertions(+), 29 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index d1987fcd00..811280c7fc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -392,8 +392,26 @@ private static TryCatch TryCreateOptimisticDirectExecutionC inputParameters: inputParameters, targetRange: new FeedRangeEpk(targetRange.ToRange()), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), - queryPipelineStage: (updatedInputParameters) => + queryPipelineStage: (CosmosElement continuationToken) => { + InputParameters updatedInputParameters = new InputParameters( + inputParameters.SqlQuerySpec, + continuationToken, + inputParameters.InitialFeedRange, + inputParameters.MaxConcurrency, + inputParameters.MaxItemCount, + inputParameters.MaxBufferedItemCount, + inputParameters.PartitionKey, + inputParameters.Properties, + inputParameters.PartitionedQueryExecutionInfo, + inputParameters.ExecutionEnvironment, + inputParameters.ReturnResultsInDeterministicOrder, + inputParameters.ForcePassthrough, + inputParameters.TestInjections) + { + IsOdeFallBackPlan = true + }; + // Query Iterator requires that the creation of the query context is deferred until fallback pipeline is called Task> tryCreateContext = CosmosQueryExecutionContextFactory.TryCreateCoreContextAsync( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 8a27de8ac1..83e28bcd21 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -20,43 +20,52 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Tracing; + using Microsoft.Azure.Documents; using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory; using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage { private const string optimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; - private readonly Func>> queryPipelineStageFactory; + public delegate Task> FallbackQueryPipelineStageFactory(CosmosElement continuationToken); + private readonly FallbackQueryPipelineStageFactory queryPipelineStageFactory; private TryCatch innerQueryPipelineStage; - private InputParameters inputParameters; + private CosmosElement continuationToken; - private OptimisticDirectExecutionQueryPipelineStage(TryCatch queryPipelineStage, Func>> queryPipelineStageFactory, InputParameters inputParameters) + private OptimisticDirectExecutionQueryPipelineStage(TryCatch queryPipelineStage, FallbackQueryPipelineStageFactory queryPipelineStageFactory, CosmosElement continuationToken) { this.innerQueryPipelineStage = queryPipelineStage; this.queryPipelineStageFactory = queryPipelineStageFactory; - this.inputParameters = inputParameters; + this.continuationToken = continuationToken; } public TryCatch Current => this.innerQueryPipelineStage.Result.Current; - public ValueTask DisposeAsync() => default; + public ValueTask DisposeAsync() + { + return this.innerQueryPipelineStage.Result.DisposeAsync(); + } public async ValueTask MoveNextAsync(ITrace trace) { + if (this.innerQueryPipelineStage.Result == null) + { + return false; + } + bool success = await this.innerQueryPipelineStage.Result.MoveNextAsync(trace); bool isGoneException = this.Current.Failed && this.Current.InnerMostException is CosmosException exception && exception.StatusCode == System.Net.HttpStatusCode.Gone - && exception.SubStatusCode == 1002; + && exception.SubStatusCode == (int)SubStatusCodes.PartitionKeyRangeGone; if (success) { - this.SaveContinuation(this.Current.Result?.State?.Value); + this.SaveContinuation(this.Current.Result.State?.Value); } else if (isGoneException) { - this.inputParameters.IsOdeFallBackPlan = true; - this.innerQueryPipelineStage = await this.queryPipelineStageFactory(this.inputParameters); + this.innerQueryPipelineStage = await this.queryPipelineStageFactory(this.continuationToken); // TODO: Failure check for this.inner bool fallbackPipelineSuccess = await this.innerQueryPipelineStage.Result.MoveNextAsync(trace); @@ -85,7 +94,7 @@ public static TryCatch MonadicCreate( InputParameters inputParameters, FeedRangeEpk targetRange, QueryPaginationOptions queryPaginationOptions, - Func>> queryPipelineStage, + FallbackQueryPipelineStageFactory queryPipelineStage, CancellationToken cancellationToken) { TryCatch pipelineStage = OptimisticDirectExecutionQueryPipelineImpl.MonadicCreate( @@ -102,7 +111,7 @@ public static TryCatch MonadicCreate( return TryCatch.FromException(pipelineStage.Exception); } - OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, queryPipelineStage, inputParameters); + OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, queryPipelineStage, inputParameters.InitialUserContinuationToken); return TryCatch.FromResult(odePipelineStageMonadicCreate); } @@ -120,23 +129,8 @@ private void SaveContinuation(CosmosElement continuationToken) CosmosArray cosmosElementParallelContinuationToken = CosmosArray.Create(parallelContinuationToken); continuationToken = cosmosElementParallelContinuationToken; } - - InputParameters inputParams = new InputParameters( - this.inputParameters.SqlQuerySpec, - continuationToken, - this.inputParameters.InitialFeedRange, - this.inputParameters.MaxConcurrency, - this.inputParameters.MaxItemCount, - this.inputParameters.MaxBufferedItemCount, - this.inputParameters.PartitionKey, - this.inputParameters.Properties, - this.inputParameters.PartitionedQueryExecutionInfo, - this.inputParameters.ExecutionEnvironment, - this.inputParameters.ReturnResultsInDeterministicOrder, - this.inputParameters.ForcePassthrough, - this.inputParameters.TestInjections); - - this.inputParameters = inputParams; + + this.continuationToken = continuationToken; } private class OptimisticDirectExecutionQueryPipelineImpl : IQueryPipelineStage From 661b4f83059ff211bd2b3a369d4f328020fe71d5 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Fri, 2 Dec 2022 12:24:04 -0800 Subject: [PATCH 15/37] Ode fallback plan always calls Specialized pipeline --- .../CosmosQueryExecutionContextFactory.cs | 165 +++++++++++++----- ...misticDirectExecutionQueryPipelineStage.cs | 53 +++--- 2 files changed, 144 insertions(+), 74 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 811280c7fc..ef0651a41c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -27,6 +27,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.SqlObjects; using Microsoft.Azure.Cosmos.SqlObjects.Visitors; using Microsoft.Azure.Cosmos.Tracing; + using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory; internal static class CosmosQueryExecutionContextFactory { @@ -151,6 +152,7 @@ private static async Task> TryCreateCoreContextAsy return TryCreateOptimisticDirectExecutionContext( documentContainer, cosmosQueryContext, + containerQueryProperties, inputParameters, targetRange, cancellationToken); @@ -321,6 +323,7 @@ private static async Task> TryCreateFromPartitione tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreateOptimisticDirectExecutionContext( documentContainer, cosmosQueryContext, + containerQueryProperties, inputParameters, targetRange, cancellationToken); @@ -340,48 +343,60 @@ private static async Task> TryCreateFromPartitione } else { - SetTestInjectionPipelineType(inputParameters, Specialized); + tryCreatePipelineStage = CallSpecializedPipeline(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken); + } - if (!string.IsNullOrEmpty(partitionedQueryExecutionInfo.QueryInfo.RewrittenQuery)) - { - // We need pass down the rewritten query. - SqlQuerySpec rewrittenQuerySpec = new SqlQuerySpec() - { - QueryText = partitionedQueryExecutionInfo.QueryInfo.RewrittenQuery, - Parameters = inputParameters.SqlQuerySpec.Parameters - }; + return tryCreatePipelineStage; + } - inputParameters = new InputParameters( - rewrittenQuerySpec, - inputParameters.InitialUserContinuationToken, - inputParameters.InitialFeedRange, - inputParameters.MaxConcurrency, - inputParameters.MaxItemCount, - inputParameters.MaxBufferedItemCount, - inputParameters.PartitionKey, - inputParameters.Properties, - inputParameters.PartitionedQueryExecutionInfo, - inputParameters.ExecutionEnvironment, - inputParameters.ReturnResultsInDeterministicOrder, - inputParameters.ForcePassthrough, - inputParameters.TestInjections); - } + private static TryCatch CallSpecializedPipeline( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + InputParameters inputParameters, + List targetRanges, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, + CancellationToken cancellationToken) + { + SetTestInjectionPipelineType(inputParameters, Specialized); - tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreateSpecializedDocumentQueryExecutionContext( - documentContainer, - cosmosQueryContext, - inputParameters, - partitionedQueryExecutionInfo, - targetRanges, - cancellationToken); + if (!string.IsNullOrEmpty(partitionedQueryExecutionInfo.QueryInfo.RewrittenQuery)) + { + // We need pass down the rewritten query. + SqlQuerySpec rewrittenQuerySpec = new SqlQuerySpec() + { + QueryText = partitionedQueryExecutionInfo.QueryInfo.RewrittenQuery, + Parameters = inputParameters.SqlQuerySpec.Parameters + }; + + inputParameters = new InputParameters( + rewrittenQuerySpec, + inputParameters.InitialUserContinuationToken, + inputParameters.InitialFeedRange, + inputParameters.MaxConcurrency, + inputParameters.MaxItemCount, + inputParameters.MaxBufferedItemCount, + inputParameters.PartitionKey, + inputParameters.Properties, + inputParameters.PartitionedQueryExecutionInfo, + inputParameters.ExecutionEnvironment, + inputParameters.ReturnResultsInDeterministicOrder, + inputParameters.ForcePassthrough, + inputParameters.TestInjections); } - return tryCreatePipelineStage; + return CosmosQueryExecutionContextFactory.TryCreateSpecializedDocumentQueryExecutionContext( + documentContainer, + cosmosQueryContext, + inputParameters, + partitionedQueryExecutionInfo, + targetRanges, + cancellationToken); } - + private static TryCatch TryCreateOptimisticDirectExecutionContext( DocumentContainer documentContainer, CosmosQueryContext cosmosQueryContext, + ContainerQueryProperties containerQueryProperties, InputParameters inputParameters, Documents.PartitionKeyRange targetRange, CancellationToken cancellationToken) @@ -407,16 +422,14 @@ private static TryCatch TryCreateOptimisticDirectExecutionC inputParameters.ExecutionEnvironment, inputParameters.ReturnResultsInDeterministicOrder, inputParameters.ForcePassthrough, - inputParameters.TestInjections) - { - IsOdeFallBackPlan = true - }; + inputParameters.TestInjections); - // Query Iterator requires that the creation of the query context is deferred until fallback pipeline is called + // In fallback scenario, the Specialized pipeline is always invoked Task> tryCreateContext = - CosmosQueryExecutionContextFactory.TryCreateCoreContextAsync( + CosmosQueryExecutionContextFactory.TryCreateOdeFallbackPipelineAsync( documentContainer, cosmosQueryContext, + containerQueryProperties, updatedInputParameters, NoOpTrace.Singleton, default); @@ -425,7 +438,72 @@ private static TryCatch TryCreateOptimisticDirectExecutionC }, cancellationToken: cancellationToken); } - + + public static async Task> TryCreateOdeFallbackPipelineAsync( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + ContainerQueryProperties containerQueryProperties, + InputParameters inputParameters, + ITrace trace, + CancellationToken cancellationToken) + { + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; + Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); + + if (partitionKeyDefinition != null) + { + partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync( + cosmosQueryContext.QueryClient, + inputParameters.SqlQuerySpec, + cosmosQueryContext.ResourceTypeEnum, + partitionKeyDefinition, + inputParameters.PartitionKey != null, + containerQueryProperties.GeospatialType, + cosmosQueryContext.UseSystemPrefix, + trace, + cancellationToken); + } + else + { + partitionedQueryExecutionInfo = new PartitionedQueryExecutionInfo() + { + QueryInfo = new QueryInfo() + { + Aggregates = null, + DistinctType = DistinctQueryType.None, + GroupByAliases = null, + GroupByAliasToAggregateType = null, + GroupByExpressions = null, + HasSelectValue = false, + Limit = null, + Offset = null, + OrderBy = null, + OrderByExpressions = null, + RewrittenQuery = null, + Top = null, + }, + QueryRanges = new List>(), + }; + } + + List targetRanges = await CosmosQueryExecutionContextFactory.GetTargetPartitionKeyRangesAsync( + cosmosQueryContext.QueryClient, + cosmosQueryContext.ResourceLink, + partitionedQueryExecutionInfo, + containerQueryProperties, + inputParameters.Properties, + inputParameters.InitialFeedRange, + trace); + + return CallSpecializedPipeline( + documentContainer, + cosmosQueryContext, + inputParameters, + targetRanges, + partitionedQueryExecutionInfo, + cancellationToken); + } + private static TryCatch TryCreatePassthroughQueryExecutionContext( DocumentContainer documentContainer, InputParameters inputParameters, @@ -645,7 +723,7 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP ContainerQueryProperties containerQueryProperties, ITrace trace) { - if (inputParameters.TestInjections == null || !inputParameters.TestInjections.EnableOptimisticDirectExecution || inputParameters.IsOdeFallBackPlan) return null; + if (inputParameters.TestInjections == null || !inputParameters.TestInjections.EnableOptimisticDirectExecution) return null; // case 1: Is query going to a single partition bool hasPartitionKey = inputParameters.PartitionKey.HasValue @@ -717,8 +795,7 @@ public InputParameters( ExecutionEnvironment? executionEnvironment, bool? returnResultsInDeterministicOrder, bool forcePassthrough, - TestInjections testInjections, - bool isOdeFallBackPlan = false) + TestInjections testInjections) { this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); this.InitialUserContinuationToken = initialUserContinuationToken; @@ -752,7 +829,6 @@ public InputParameters( this.ReturnResultsInDeterministicOrder = returnResultsInDeterministicOrder.GetValueOrDefault(InputParameters.DefaultReturnResultsInDeterministicOrder); this.ForcePassthrough = forcePassthrough; this.TestInjections = testInjections; - this.IsOdeFallBackPlan = isOdeFallBackPlan; } public SqlQuerySpec SqlQuerySpec { get; } @@ -768,7 +844,6 @@ public InputParameters( public bool ReturnResultsInDeterministicOrder { get; } public TestInjections TestInjections { get; } public bool ForcePassthrough { get; } - public bool IsOdeFallBackPlan { get; set; } } internal sealed class AggregateProjectionDetector diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 83e28bcd21..f85fb7735e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -27,33 +27,36 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage { private const string optimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; + private const string optimisticDirectExecution = "OptimisticDirectExecution"; + private const string specializedDocumentQueryExecution = "SpecializedDocumentQueryExecution"; public delegate Task> FallbackQueryPipelineStageFactory(CosmosElement continuationToken); private readonly FallbackQueryPipelineStageFactory queryPipelineStageFactory; - private TryCatch innerQueryPipelineStage; + private TryCatch inner; private CosmosElement continuationToken; + private string queryState; private OptimisticDirectExecutionQueryPipelineStage(TryCatch queryPipelineStage, FallbackQueryPipelineStageFactory queryPipelineStageFactory, CosmosElement continuationToken) { - this.innerQueryPipelineStage = queryPipelineStage; + this.inner = queryPipelineStage; this.queryPipelineStageFactory = queryPipelineStageFactory; this.continuationToken = continuationToken; } - public TryCatch Current => this.innerQueryPipelineStage.Result.Current; + public TryCatch Current => this.inner.Result.Current; public ValueTask DisposeAsync() { - return this.innerQueryPipelineStage.Result.DisposeAsync(); + return this.inner.Result.DisposeAsync(); } public async ValueTask MoveNextAsync(ITrace trace) { - if (this.innerQueryPipelineStage.Result == null) + if (this.inner.Result == null) { return false; } - bool success = await this.innerQueryPipelineStage.Result.MoveNextAsync(trace); + bool success = await this.inner.Result.MoveNextAsync(trace); bool isGoneException = this.Current.Failed && this.Current.InnerMostException is CosmosException exception && exception.StatusCode == System.Net.HttpStatusCode.Gone @@ -61,23 +64,20 @@ public async ValueTask MoveNextAsync(ITrace trace) if (success) { + this.queryState = optimisticDirectExecution; this.SaveContinuation(this.Current.Result.State?.Value); } else if (isGoneException) { - this.innerQueryPipelineStage = await this.queryPipelineStageFactory(this.continuationToken); + // Only Ode pipeline GoneException should be handled + Debug.Assert(this.queryState == optimisticDirectExecution); + this.inner = await this.queryPipelineStageFactory(this.UnwrapContinuationToken(this.continuationToken)); + this.queryState = specializedDocumentQueryExecution; + // TODO: Failure check for this.inner - bool fallbackPipelineSuccess = await this.innerQueryPipelineStage.Result.MoveNextAsync(trace); + bool fallbackPipelineSuccess = await this.inner.Result.MoveNextAsync(trace); - if (this.Current.Result?.State?.Value != null) - { - if (this.Current.Result.State.Value is CosmosObject) - { - // Fallback plan returned a Ode pipeline which is wrong - return false; - } - } return fallbackPipelineSuccess; } @@ -86,7 +86,14 @@ public async ValueTask MoveNextAsync(ITrace trace) public void SetCancellationToken(CancellationToken cancellationToken) { - this.innerQueryPipelineStage.Result.SetCancellationToken(cancellationToken); + this.inner.Result.SetCancellationToken(cancellationToken); + } + + public CosmosElement UnwrapContinuationToken(CosmosElement continuationToken) + { + ((CosmosObject)continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken); + CosmosArray cosmosElementFallbackToken = CosmosArray.Create(specializedContinuationToken); + return cosmosElementFallbackToken; } public static TryCatch MonadicCreate( @@ -118,18 +125,6 @@ public static TryCatch MonadicCreate( private void SaveContinuation(CosmosElement continuationToken) { - if (continuationToken == null) - { - return; - } - - if (continuationToken is CosmosObject) - { - ((CosmosObject)continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement parallelContinuationToken); - CosmosArray cosmosElementParallelContinuationToken = CosmosArray.Create(parallelContinuationToken); - continuationToken = cosmosElementParallelContinuationToken; - } - this.continuationToken = continuationToken; } From 039fbe8d2072ddee59863aa62c040ca0999689a0 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Tue, 6 Dec 2022 18:34:56 -0800 Subject: [PATCH 16/37] Using ServiceInterop/Gateway to get QueryPlan for Specialized Pipeline --- .../Query/Core/Monads/TryCatch{TResult}.cs | 16 ++ .../CosmosQueryExecutionContextFactory.cs | 129 +++++------- ...misticDirectExecutionQueryPipelineStage.cs | 75 ++++--- ...ositiveOptimisticDirectExecutionOutput.xml | 38 +++- .../Pagination/FlakyDocumentContainer.cs | 2 +- ...misticDirectExecutionQueryBaselineTests.cs | 195 +++++++++--------- 6 files changed, 241 insertions(+), 214 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs index 8d9c43b036..0fb56995e2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs @@ -126,6 +126,22 @@ public async Task> TryAsync( return matchResult; } + public TryCatch Try( + Func> onSuccess) + { + TryCatch matchResult; + if (this.Succeeded) + { + matchResult = onSuccess(this.either.FromRight(default)); + } + else + { + matchResult = TryCatch.FromException(this.either.FromLeft(default)); + } + + return matchResult; + } + public TryCatch Catch( Action onError) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index ef0651a41c..7028c966bf 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -27,7 +27,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.SqlObjects; using Microsoft.Azure.Cosmos.SqlObjects.Visitors; using Microsoft.Azure.Cosmos.Tracing; - using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory; internal static class CosmosQueryExecutionContextFactory { @@ -138,12 +137,12 @@ private static async Task> TryCreateCoreContextAsy cosmosQueryContext.ContainerResourceId = containerQueryProperties.ResourceId; Documents.PartitionKeyRange targetRange = await GetTargetRangeOptimisticDirectExecutionAsync( - inputParameters, - queryPlanFromContinuationToken, - cosmosQueryContext, - containerQueryProperties, + inputParameters, + queryPlanFromContinuationToken, + cosmosQueryContext, + containerQueryProperties, trace); - + if (targetRange != null) { // Test code added to confirm the correct pipeline is being utilized @@ -229,33 +228,12 @@ private static async Task> TryCreateCoreContextAsy } } - if (cosmosQueryContext.QueryClient.ByPassQueryParsing()) - { - // For non-Windows platforms(like Linux and OSX) in .NET Core SDK, we cannot use ServiceInterop, so need to bypass in that case. - // We are also now bypassing this for 32 bit host process running even on Windows as there are many 32 bit apps that will not work without this - partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanThroughGatewayAsync( - cosmosQueryContext, - inputParameters.SqlQuerySpec, - cosmosQueryContext.ResourceLink, - inputParameters.PartitionKey, - createQueryPipelineTrace, - cancellationToken); - } - else - { - Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); - - partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync( - cosmosQueryContext.QueryClient, - inputParameters.SqlQuerySpec, - cosmosQueryContext.ResourceTypeEnum, - partitionKeyDefinition, - inputParameters.PartitionKey != null, - containerQueryProperties.GeospatialType, - cosmosQueryContext.UseSystemPrefix, - createQueryPipelineTrace, - cancellationToken); - } + partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoFromGatewayOrServiceInteropAsync( + cosmosQueryContext, + inputParameters, + containerQueryProperties, + createQueryPipelineTrace, + cancellationToken); } return await TryCreateFromPartitionedQueryExecutionInfoAsync( @@ -407,7 +385,7 @@ private static TryCatch TryCreateOptimisticDirectExecutionC inputParameters: inputParameters, targetRange: new FeedRangeEpk(targetRange.ToRange()), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), - queryPipelineStage: (CosmosElement continuationToken) => + fallbackQueryPipelineStageFactory: (continuationToken) => { InputParameters updatedInputParameters = new InputParameters( inputParameters.SqlQuerySpec, @@ -447,44 +425,12 @@ public static async Task> TryCreateOdeFallbackPipe ITrace trace, CancellationToken cancellationToken) { - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; - Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); - - if (partitionKeyDefinition != null) - { - partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync( - cosmosQueryContext.QueryClient, - inputParameters.SqlQuerySpec, - cosmosQueryContext.ResourceTypeEnum, - partitionKeyDefinition, - inputParameters.PartitionKey != null, - containerQueryProperties.GeospatialType, - cosmosQueryContext.UseSystemPrefix, - trace, - cancellationToken); - } - else - { - partitionedQueryExecutionInfo = new PartitionedQueryExecutionInfo() - { - QueryInfo = new QueryInfo() - { - Aggregates = null, - DistinctType = DistinctQueryType.None, - GroupByAliases = null, - GroupByAliasToAggregateType = null, - GroupByExpressions = null, - HasSelectValue = false, - Limit = null, - Offset = null, - OrderBy = null, - OrderByExpressions = null, - RewrittenQuery = null, - Top = null, - }, - QueryRanges = new List>(), - }; - } + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoFromGatewayOrServiceInteropAsync( + cosmosQueryContext, + inputParameters, + containerQueryProperties, + trace, + cancellationToken); List targetRanges = await CosmosQueryExecutionContextFactory.GetTargetPartitionKeyRangesAsync( cosmosQueryContext.QueryClient, @@ -588,6 +534,45 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx requestCancellationToken: cancellationToken); } + private static async Task GetPartitionedQueryExecutionInfoFromGatewayOrServiceInteropAsync( + CosmosQueryContext cosmosQueryContext, + InputParameters inputParameters, + ContainerQueryProperties containerQueryProperties, + ITrace trace, + CancellationToken cancellationToken) + { + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; + if (cosmosQueryContext.QueryClient.ByPassQueryParsing()) + { + // For non-Windows platforms(like Linux and OSX) in .NET Core SDK, we cannot use ServiceInterop, so need to bypass in that case. + // We are also now bypassing this for 32 bit host process running even on Windows as there are many 32 bit apps that will not work without this + partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanThroughGatewayAsync( + cosmosQueryContext, + inputParameters.SqlQuerySpec, + cosmosQueryContext.ResourceLink, + inputParameters.PartitionKey, + trace, + cancellationToken); + } + else + { + Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); + + partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync( + cosmosQueryContext.QueryClient, + inputParameters.SqlQuerySpec, + cosmosQueryContext.ResourceTypeEnum, + partitionKeyDefinition, + inputParameters.PartitionKey != null, + containerQueryProperties.GeospatialType, + cosmosQueryContext.UseSystemPrefix, + trace, + cancellationToken); + } + + return partitionedQueryExecutionInfo; + } + /// /// Gets the list of partition key ranges. /// 1. Check partition key range id diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index f85fb7735e..6188e214bb 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -6,7 +6,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu { using System; using System.Collections.Generic; - using System.Diagnostics; using System.Linq; using System.Threading; using System.Threading.Tasks; @@ -18,41 +17,44 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Tracing; using Microsoft.Azure.Documents; - using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory; - using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage { + private enum ExecutionState + { + OptimisticDirectExecution, + SpecializedDocumentQueryExecution, + } + private const string optimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; - private const string optimisticDirectExecution = "OptimisticDirectExecution"; - private const string specializedDocumentQueryExecution = "SpecializedDocumentQueryExecution"; public delegate Task> FallbackQueryPipelineStageFactory(CosmosElement continuationToken); private readonly FallbackQueryPipelineStageFactory queryPipelineStageFactory; private TryCatch inner; private CosmosElement continuationToken; - private string queryState; - - private OptimisticDirectExecutionQueryPipelineStage(TryCatch queryPipelineStage, FallbackQueryPipelineStageFactory queryPipelineStageFactory, CosmosElement continuationToken) + private ExecutionState executionState; + + private OptimisticDirectExecutionQueryPipelineStage(TryCatch inner, FallbackQueryPipelineStageFactory queryPipelineStageFactory, CosmosElement continuationToken) { - this.inner = queryPipelineStage; + this.inner = inner; this.queryPipelineStageFactory = queryPipelineStageFactory; this.continuationToken = continuationToken; + this.executionState = ExecutionState.OptimisticDirectExecution; } - public TryCatch Current => this.inner.Result.Current; + public TryCatch Current => this.inner.Try(pipelineStage => pipelineStage.Current); public ValueTask DisposeAsync() { - return this.inner.Result.DisposeAsync(); + return this.inner.Try(pipelineStage => pipelineStage.DisposeAsync()).Result; } - + public async ValueTask MoveNextAsync(ITrace trace) { - if (this.inner.Result == null) + if (this.inner.Failed) { + this.inner = TryCatch.FromException(this.inner.Exception); return false; } @@ -64,44 +66,49 @@ public async ValueTask MoveNextAsync(ITrace trace) if (success) { - this.queryState = optimisticDirectExecution; this.SaveContinuation(this.Current.Result.State?.Value); } - else if (isGoneException) + else if (isGoneException && this.executionState == ExecutionState.OptimisticDirectExecution) { - // Only Ode pipeline GoneException should be handled - Debug.Assert(this.queryState == optimisticDirectExecution); + this.inner = await this.queryPipelineStageFactory(TryUnwrapContinuationToken(this.continuationToken).Result); + this.executionState = ExecutionState.SpecializedDocumentQueryExecution; - this.inner = await this.queryPipelineStageFactory(this.UnwrapContinuationToken(this.continuationToken)); - this.queryState = specializedDocumentQueryExecution; - - // TODO: Failure check for this.inner - bool fallbackPipelineSuccess = await this.inner.Result.MoveNextAsync(trace); + if (this.inner.Failed) + { + this.inner = TryCatch.FromException(this.inner.Exception); + return false; + } - return fallbackPipelineSuccess; + success = await this.inner.Result.MoveNextAsync(trace); } - return success; + return success; } public void SetCancellationToken(CancellationToken cancellationToken) { - this.inner.Result.SetCancellationToken(cancellationToken); + this.inner.Try((pipelineStage) => pipelineStage = this.inner.Result).Result.SetCancellationToken(cancellationToken); } - public CosmosElement UnwrapContinuationToken(CosmosElement continuationToken) + public static TryCatch TryUnwrapContinuationToken(CosmosElement continuationToken) { - ((CosmosObject)continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken); + if (!((CosmosObject)continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken)) + { + return TryCatch.FromException( + new FormatException( + $"Unable to convert CosmosObject '{optimisticDirectExecutionToken}' to CosmosElement.")); + } + CosmosArray cosmosElementFallbackToken = CosmosArray.Create(specializedContinuationToken); - return cosmosElementFallbackToken; + return TryCatch.FromResult(cosmosElementFallbackToken); } public static TryCatch MonadicCreate( DocumentContainer documentContainer, - InputParameters inputParameters, + CosmosQueryExecutionContextFactory.InputParameters inputParameters, FeedRangeEpk targetRange, QueryPaginationOptions queryPaginationOptions, - FallbackQueryPipelineStageFactory queryPipelineStage, + FallbackQueryPipelineStageFactory fallbackQueryPipelineStageFactory, CancellationToken cancellationToken) { TryCatch pipelineStage = OptimisticDirectExecutionQueryPipelineImpl.MonadicCreate( @@ -118,7 +125,7 @@ public static TryCatch MonadicCreate( return TryCatch.FromException(pipelineStage.Exception); } - OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, queryPipelineStage, inputParameters.InitialUserContinuationToken); + OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, fallbackQueryPipelineStageFactory, inputParameters.InitialUserContinuationToken); return TryCatch.FromResult(odePipelineStageMonadicCreate); } @@ -262,7 +269,7 @@ private static TryCatch> MonadicExtractState( return TryCatch>.FromException(tryCreateContinuationToken.Exception); } - TryCatch> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping( + TryCatch> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping( range, tryCreateContinuationToken.Result); @@ -272,7 +279,7 @@ private static TryCatch> MonadicExtractState( partitionMappingMonad.Exception); } - PartitionMapping partitionMapping = partitionMappingMonad.Result; + PartitionMapper.PartitionMapping partitionMapping = partitionMappingMonad.Result; KeyValuePair kvpRange = new KeyValuePair( partitionMapping.TargetMapping.Keys.First(), diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.PositiveOptimisticDirectExecutionOutput.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.PositiveOptimisticDirectExecutionOutput.xml index 950947ff5a..cedd618124 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.PositiveOptimisticDirectExecutionOutput.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.PositiveOptimisticDirectExecutionOutput.xml @@ -1,8 +1,8 @@  - Partition Key + Value and Distinct - SELECT DISTINCT c.key FROM c + Single Partition Key and Distinct + SELECT DISTINCT c.age FROM c /pk @@ -14,8 +14,8 @@ - Partition Key + Value and Min Aggregate - SELECT VALUE MIN(c.key) FROM c + Single Partition Key and Min Aggregate + SELECT VALUE MIN(c.age) FROM c /pk @@ -27,8 +27,34 @@ - Partition Key + Value Fields - SELECT c.key FROM c + Single Partition Key and Value Field + SELECT c.age FROM c + + /pk + + Hash + + + true + + + + + Single Partition Key and Value Field + SELECT * FROM c + + /pk + + Hash + + + true + + + + + Single Partition Key and continuation token + SELECT * FROM c /pk diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs index 3cca5432e7..f7f7fbec1c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs @@ -182,7 +182,7 @@ public async Task> MonadicQueryAsync( if (this.isODETest) { this.iterationCount++; - if (this.iterationCount == 2) + if (this.iterationCount == 2 || this.iterationCount == 5) { Exception failure = await this.ShouldReturnFailure(); if (failure != null) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index f1f064bd6f..b0b9cb7d23 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -3,30 +3,28 @@ using System; using System.Collections.Generic; using System.Collections.ObjectModel; - using System.Xml; - using Microsoft.Azure.Documents; - using Microsoft.Azure.Documents.Routing; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Test.BaselineTest; - using Microsoft.VisualStudio.TestTools.UnitTesting; - using Newtonsoft.Json; - using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using System.IO; + using System.Linq; + using System.Threading; using System.Threading.Tasks; + using System.Xml; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Tests.Pagination; - using Microsoft.Azure.Cosmos.Tracing; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using Moq; using Microsoft.Azure.Cosmos.Query; + using Microsoft.Azure.Cosmos.Query.Core; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Routing; - using System.Threading; - using System.Linq; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; - using System.IO; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Cosmos.Test.BaselineTest; + using Microsoft.Azure.Cosmos.Tests.Pagination; + using Microsoft.Azure.Cosmos.Tracing; + using Microsoft.Azure.Documents; + using Microsoft.Azure.Documents.Routing; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Newtonsoft.Json; [TestClass] public class OptimisticDirectExecutionQueryBaselineTests : BaselineTests @@ -131,20 +129,14 @@ public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() OptimisticDirectExecutionTestInput input = CreateInput( description: @"Single Partition Key and Value Field", query: "SELECT VALUE COUNT(1) FROM c", - expectedOptimisticDirectExecution: true, + expectedOptimisticDirectExecution: true, partitionKeyPath: @"/pk", partitionKeyValue: "a"); - QueryRequestOptions queryRequestOptions = new QueryRequestOptions - { - MaxConcurrency = 0, - MaxItemCount = 10, - MaxBufferedItemCount = 7000, - TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) - }; + QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: false); - IQueryPipelineStage queryPipelineStage = await CreateAndGetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); + IQueryPipelineStage queryPipelineStage = await GetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); int documentCountInSinglePartition = 0; while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) @@ -177,7 +169,7 @@ public async Task TestPipelineForContinuationTokenOnSinglePartitionAsync() partitionKeyPath: @"/pk", partitionKeyValue: "a"); - int result = await this.CreateOptimisticPipelineAndDrainAsync( + int result = await this.GetPipelineAndDrainAsync( input, numItems: numItems, isMultiPartition: false, @@ -209,7 +201,7 @@ public async Task TestPipelineForDistributedQueryAsync() partitionKeyPath: @"/pk", partitionKeyValue: "a"); - int result = await this.CreateOptimisticPipelineAndDrainAsync( + int result = await this.GetPipelineAndDrainAsync( input, numItems: numItems, isMultiPartition: false, @@ -242,25 +234,18 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", requestCharge: default); - QueryRequestOptions queryRequestOptions = new QueryRequestOptions - { - MaxConcurrency = 0, - MaxItemCount = 10, - MaxBufferedItemCount = 7000, - TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) - }; + QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); int moveNextAsyncCounter = 0; - bool caughtGoneException = false; DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( - numItems, - multiPartition: isMultiPartition, - failureConfigs: new FlakyDocumentContainer.FailureConfigs( - inject429s: false, - injectEmptyPages: false, - shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter == 1 ? goneException : null))); + numItems, + multiPartition: isMultiPartition, + failureConfigs: new FlakyDocumentContainer.FailureConfigs( + inject429s: false, + injectEmptyPages: false, + shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter == 1 || moveNextAsyncCounter == 4 ? goneException : null))); - IQueryPipelineStage queryPipelineStage = await CreateAndGetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); + IQueryPipelineStage queryPipelineStage = await GetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { @@ -269,7 +254,7 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa { Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); } - else + else { Assert.AreNotEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); } @@ -285,24 +270,15 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa documents.AddRange(tryGetPage.Result.Documents); } - // Once fallback plan is implemented, this test should be able to return all 100 documents Assert.AreEqual(100, documents.Count); - return true; } - private async Task CreateOptimisticPipelineAndDrainAsync(OptimisticDirectExecutionTestInput input, int numItems, bool isMultiPartition, int expectedContinuationTokenCount) + private async Task GetPipelineAndDrainAsync(OptimisticDirectExecutionTestInput input, int numItems, bool isMultiPartition, int expectedContinuationTokenCount) { - QueryRequestOptions queryRequestOptions = new QueryRequestOptions - { - MaxConcurrency = 0, - MaxItemCount = 10, - MaxBufferedItemCount = 7000, - TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) - }; - + QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync(numItems, multiPartition: isMultiPartition); - IQueryPipelineStage queryPipelineStage = await CreateAndGetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); + IQueryPipelineStage queryPipelineStage = await GetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); List documents = new List(); int continuationTokenCount = 0; @@ -330,7 +306,7 @@ private async Task CreateOptimisticPipelineAndDrainAsync(OptimisticDirectEx partitionKeyValue: input.PartitionKeyValue, continuationToken: tryGetPage.Result.State.Value); - queryPipelineStage = await CreateAndGetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); + queryPipelineStage = await GetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); } continuationTokenCount++; @@ -340,6 +316,36 @@ private async Task CreateOptimisticPipelineAndDrainAsync(OptimisticDirectEx return documents.Count; } + private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(string querySpecJsonString, PartitionKeyDefinition pkDefinition) + { + TryCatch tryGetQueryPlan = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfo( + querySpecJsonString: querySpecJsonString, + partitionKeyDefinition: pkDefinition, + requireFormattableOrderByQuery: true, + isContinuationExpected: true, + allowNonValueAggregateQuery: true, + hasLogicalPartitionKey: false, + allowDCount: true, + useSystemPrefix: false, + geospatialType: Cosmos.GeospatialType.Geography); + + return tryGetQueryPlan.Result; + } + + private static async Task GetOdePipelineAsync(OptimisticDirectExecutionTestInput input, DocumentContainer documentContainer, QueryRequestOptions queryRequestOptions) + { + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input, queryRequestOptions); + + IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( + documentContainer, + cosmosQueryContextCore, + inputParameters, + NoOpTrace.Singleton); + + Assert.IsNotNull(queryPipelineStage); + return queryPipelineStage; + } + private static async Task CreateDocumentContainerAsync( int numItems, bool multiPartition, @@ -364,7 +370,7 @@ private static async Task CreateDocumentContainerAsync( DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); // a value of 2 would lead to 4 partitions (2 * 2). 4 partitions are used because they're easy to manage + demonstrates multi partition use case - int exponentPartitionKeyRanges = 2; + int exponentPartitionKeyRanges = 2; IReadOnlyList ranges; @@ -388,7 +394,7 @@ private static async Task CreateDocumentContainerAsync( ranges = await documentContainer.GetFeedRangesAsync( trace: NoOpTrace.Singleton, cancellationToken: default); - + int rangeCount = multiPartition ? 4 : 1; Assert.AreEqual(rangeCount, ranges.Count); @@ -429,50 +435,13 @@ private static OptimisticDirectExecutionTestInput CreateInput( return new OptimisticDirectExecutionTestInput(description, query, new SqlQuerySpec(query), expectedOptimisticDirectExecution, partitionKeyPath, partitionKeyValue, continuationToken); } - private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(string querySpecJsonString, PartitionKeyDefinition pkDefinition) - { - TryCatch tryGetQueryPlan = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfo( - querySpecJsonString: querySpecJsonString, - partitionKeyDefinition: pkDefinition, - requireFormattableOrderByQuery: true, - isContinuationExpected: true, - allowNonValueAggregateQuery: true, - hasLogicalPartitionKey: false, - allowDCount: true, - useSystemPrefix: false, - geospatialType: Cosmos.GeospatialType.Geography); - - return tryGetQueryPlan.Result; - } - - private static async Task CreateAndGetOdePipelineAsync(OptimisticDirectExecutionTestInput input, DocumentContainer documentContainer, QueryRequestOptions queryRequestOptions) - { - (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input, queryRequestOptions); - - IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( - documentContainer, - cosmosQueryContextCore, - inputParameters, - NoOpTrace.Singleton); - - Assert.IsNotNull(queryPipelineStage); - - return queryPipelineStage; - } - public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirectExecutionTestInput input) { // gets DocumentContainer IMonadicDocumentContainer monadicDocumentContainer = new InMemoryContainer(input.PartitionKeyDefinition); DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); - QueryRequestOptions queryRequestOptions = new QueryRequestOptions - { - MaxConcurrency = 0, - MaxItemCount = 10, - MaxBufferedItemCount = 7000, - TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: true, new TestInjections.ResponseStats()) - }; + QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input, queryRequestOptions); @@ -541,6 +510,21 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect return Tuple.Create(inputParameters, cosmosQueryContextCore); } + + private static QueryRequestOptions GetQueryRequestOptions(bool enableOptimisticDirectExecution) + { + return new QueryRequestOptions + { + MaxConcurrency = 0, + MaxItemCount = 10, + MaxBufferedItemCount = 7000, + TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: enableOptimisticDirectExecution, new TestInjections.ResponseStats()), + Properties = new Dictionary() + { + { HttpConstants.HttpHeaders.EnumerationDirection, "test" }, + } + }; + } } public sealed class OptimisticDirectExecutionTestOutput : BaselineTestOutput @@ -635,7 +619,7 @@ internal class TestCosmosQueryClient : CosmosQueryClient public override bool ByPassQueryParsing() { - throw new NotImplementedException(); + return false; } public override void ClearSessionTokenCache(string collectionFullName) @@ -660,7 +644,11 @@ public override Task ForceRefreshCollectionCacheAsync(string collectionLink, Can public override Task GetCachedContainerQueryPropertiesAsync(string containerLink, Cosmos.PartitionKey? partitionKey, ITrace trace, CancellationToken cancellationToken) { - return Task.FromResult(new ContainerQueryProperties()); + return Task.FromResult(new ContainerQueryProperties( + "test", + WFConstants.BackendHeaders.EffectivePartitionKeyString, + new PartitionKeyDefinition(), + Cosmos.GeospatialType.Geometry)); } public override Task> GetTargetPartitionKeyRangeByFeedRangeAsync(string resourceLink, string collectionResourceId, PartitionKeyDefinition partitionKeyDefinition, FeedRangeInternal feedRangeInternal, bool forceRefresh, ITrace trace) @@ -680,7 +668,12 @@ public override Task> GetTargetPartitionKeyRangesAsync(s public override Task> GetTargetPartitionKeyRangesByEpkStringAsync(string resourceLink, string collectionResourceId, string effectivePartitionKeyString, bool forceRefresh, ITrace trace) { - throw new NotImplementedException(); + return Task.FromResult(new List{new PartitionKeyRange() + { + MinInclusive = PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, + MaxExclusive = PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey + } + }); } public override Task> TryGetOverlappingRangesAsync(string collectionResourceId, Range range, bool forceRefresh = false) From 60392adc71dc6bfe02166d21c5d14f7a10347051 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Sun, 18 Dec 2022 15:06:11 -0800 Subject: [PATCH 17/37] Added new test to check handling of failing fallback pipeline --- ...OrderByCrossPartitionQueryPipelineStage.cs | 4 +- ...misticDirectExecutionQueryPipelineStage.cs | 37 ++-- .../src/Query/Core/Utils.cs | 20 ++ ...egativeOptimisticDirectExecutionOutput.xml | 13 ++ ...misticDirectExecutionQueryBaselineTests.cs | 180 ++++++++++++++++-- 5 files changed, 208 insertions(+), 46 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/src/Query/Core/Utils.cs diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index bbd3d7be4e..0efa1cd950 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -1153,9 +1153,7 @@ private static bool IsSplitException(Exception exception) exception = exception.InnerException; } - return exception is CosmosException cosmosException - && (cosmosException.StatusCode == HttpStatusCode.Gone) - && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); + return Utils.IsPartitionSplitException(exception); } public void SetCancellationToken(CancellationToken cancellationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 6188e214bb..1fbc8a28bf 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu { using System; using System.Collections.Generic; + using System.Diagnostics; using System.Linq; using System.Threading; using System.Threading.Tasks; @@ -52,34 +53,25 @@ public ValueTask DisposeAsync() public async ValueTask MoveNextAsync(ITrace trace) { - if (this.inner.Failed) - { - this.inner = TryCatch.FromException(this.inner.Exception); - return false; - } - bool success = await this.inner.Result.MoveNextAsync(trace); - bool isGoneException = this.Current.Failed - && this.Current.InnerMostException is CosmosException exception - && exception.StatusCode == System.Net.HttpStatusCode.Gone - && exception.SubStatusCode == (int)SubStatusCodes.PartitionKeyRangeGone; - if (success) { this.SaveContinuation(this.Current.Result.State?.Value); } - else if (isGoneException && this.executionState == ExecutionState.OptimisticDirectExecution) + else if (this.executionState == ExecutionState.OptimisticDirectExecution) { - this.inner = await this.queryPipelineStageFactory(TryUnwrapContinuationToken(this.continuationToken).Result); - this.executionState = ExecutionState.SpecializedDocumentQueryExecution; - - if (this.inner.Failed) + if (Utils.IsPartitionSplitException(this.Current.InnerMostException)) { - this.inner = TryCatch.FromException(this.inner.Exception); - return false; + this.inner = await this.queryPipelineStageFactory(this.TryUnwrapContinuationToken().Result); + if (this.inner.Failed) + { + this.inner = TryCatch.FromException(this.inner.Exception); + return false; + } + + this.executionState = ExecutionState.SpecializedDocumentQueryExecution; + success = await this.inner.Result.MoveNextAsync(trace); } - - success = await this.inner.Result.MoveNextAsync(trace); } return success; @@ -90,9 +82,10 @@ public void SetCancellationToken(CancellationToken cancellationToken) this.inner.Try((pipelineStage) => pipelineStage = this.inner.Result).Result.SetCancellationToken(cancellationToken); } - public static TryCatch TryUnwrapContinuationToken(CosmosElement continuationToken) + public TryCatch TryUnwrapContinuationToken() { - if (!((CosmosObject)continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken)) + Debug.Assert(this.continuationToken is CosmosObject); + if (!((CosmosObject)this.continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken)) { return TryCatch.FromException( new FormatException( diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Utils.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Utils.cs new file mode 100644 index 0000000000..93733f7bf4 --- /dev/null +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Utils.cs @@ -0,0 +1,20 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core +{ + using System; + using System.Collections.Generic; + using System.Text; + + internal static class Utils + { + public static bool IsPartitionSplitException(this Exception ex) + { + return ex is CosmosException cosmosException + && (cosmosException.StatusCode == System.Net.HttpStatusCode.Gone) + && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml index b56c470469..a6905f99b3 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml @@ -38,4 +38,17 @@ false + + + Exception with malformed continuation token + SELECT * FROM c + + /pk + + Hash + + + false + + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index b0b9cb7d23..d421f6f9ce 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -12,9 +12,13 @@ using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query; using Microsoft.Azure.Cosmos.Query.Core; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; @@ -24,7 +28,9 @@ using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Routing; using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; using Newtonsoft.Json; + using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory; [TestClass] public class OptimisticDirectExecutionQueryBaselineTests : BaselineTests @@ -75,8 +81,6 @@ public void PositiveOptimisticDirectExecutionOutput() partitionKeyPath: @"/pk", partitionKeyValue: "a", continuationToken: cosmosElementContinuationToken), - - //TODO: Add non ODE continuation token case }; this.ExecuteTestSuite(testVariations); } @@ -107,10 +111,20 @@ public void NegativeOptimisticDirectExecutionOutput() expectedOptimisticDirectExecution: false, partitionKeyPath: @"/pk", partitionKeyValue: null), + + // Fail to execute due to MalFormedContinuationToken Exception + CreateInput( + description: @"Exception with malformed continuation token", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: "a", + continuationToken: CosmosString.Create("asdf"), + expectException: true), }; this.ExecuteTestSuite(testVariations); } - + // This test confirms that TestInjection.EnableOptimisticDirectExection is set to false from default. // Check test "TestPipelineForDistributedQueryAsync" to understand why this is done [TestMethod] @@ -211,12 +225,104 @@ public async Task TestPipelineForDistributedQueryAsync() Assert.AreEqual(1, result); } + [TestMethod] + public async Task TestHandlingOfFailedFallbackPipeline() + { + List documents = new List(); + int numItems = 100; + int moveNextAsyncCounter = 0; + + List input = new List + { + CreateInput( + description: @"Single Partition Key and Value Field", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: "a", + continuationToken: null) + }; + + CosmosException goneException = await CreateGoneException(); + DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( + numItems, + multiPartition: false, + failureConfigs: new FlakyDocumentContainer.FailureConfigs( + inject429s: false, + injectEmptyPages: false, + shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter == 1 ? goneException : null))); + + QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input[0], queryRequestOptions); + + ContainerQueryProperties containerQueryProperties = new ContainerQueryProperties( + null, + null, + new PartitionKeyDefinition(), + It.IsAny()); + + TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( + documentContainer: inMemoryCollection, + inputParameters: inputParameters, + targetRange: FeedRangeEpk.FullRange, + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), + fallbackQueryPipelineStageFactory: (continuationToken) => + { + InputParameters updatedInputParameters = new InputParameters( + inputParameters.SqlQuerySpec, + CosmosString.Create("asdf"), + inputParameters.InitialFeedRange, + inputParameters.MaxConcurrency, + inputParameters.MaxItemCount, + inputParameters.MaxBufferedItemCount, + inputParameters.PartitionKey, + inputParameters.Properties, + inputParameters.PartitionedQueryExecutionInfo, + inputParameters.ExecutionEnvironment, + inputParameters.ReturnResultsInDeterministicOrder, + inputParameters.ForcePassthrough, + inputParameters.TestInjections); + + Task> tryCreateContext = + CosmosQueryExecutionContextFactory.TryCreateOdeFallbackPipelineAsync( + inMemoryCollection, + cosmosQueryContextCore, + containerQueryProperties, + updatedInputParameters, + NoOpTrace.Singleton, + default); + + return tryCreateContext; + }, + cancellationToken: default); + + while (await monadicQueryPipelineStage.Result.MoveNextAsync(NoOpTrace.Singleton)) + { + moveNextAsyncCounter++; + + TryCatch tryGetPage = monadicQueryPipelineStage.Result.Current; + + if (tryGetPage.Failed) + { + // failure should never come till here. Should be handled before + Assert.Fail(); + } + + documents.AddRange(tryGetPage.Result.Documents); + } + + Assert.IsTrue(monadicQueryPipelineStage.Result.Current.Failed); + Assert.IsTrue(monadicQueryPipelineStage.Result.Current.InnerMostException is MalformedContinuationTokenException); + Assert.AreNotEqual(numItems, documents.Count); + } + // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing // TODO: With the addition of the merge/split support, this queryPipelineStage should be able to return all documents regardless of a gone exception happening private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) { + List documents = new List(); int numItems = 100; - string query = "SELECT * FROM c"; + int moveNextAsyncCounter = 0; OptimisticDirectExecutionTestInput input = CreateInput( description: @"Single Partition Key and Value Field", @@ -225,18 +331,9 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa partitionKeyPath: @"/pk", partitionKeyValue: "a"); - List documents = new List(); - string goneExceptionMessage = $"Epk Range: Partition does not exist at the given range."; - CosmosException goneException = new CosmosException( - message: goneExceptionMessage, - statusCode: System.Net.HttpStatusCode.Gone, - subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, - activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", - requestCharge: default); - + CosmosException goneException = await CreateGoneException(); QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); - int moveNextAsyncCounter = 0; DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( numItems, multiPartition: isMultiPartition, @@ -346,6 +443,17 @@ private static async Task GetOdePipelineAsync(OptimisticDir return queryPipelineStage; } + private static async Task CreateGoneException() + { + string goneExceptionMessage = $"Epk Range: Partition does not exist at the given range."; + return new CosmosException( + message: goneExceptionMessage, + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", + requestCharge: default); + } + private static async Task CreateDocumentContainerAsync( int numItems, bool multiPartition, @@ -416,12 +524,13 @@ private static OptimisticDirectExecutionTestInput CreateInput( bool expectedOptimisticDirectExecution, string partitionKeyPath, string partitionKeyValue, - CosmosElement continuationToken = null) + CosmosElement continuationToken = null, + bool expectException = false) { PartitionKeyBuilder pkBuilder = new PartitionKeyBuilder(); pkBuilder.Add(partitionKeyValue); - return CreateInput(description, query, expectedOptimisticDirectExecution, partitionKeyPath, pkBuilder.Build(), continuationToken); + return CreateInput(description, query, expectedOptimisticDirectExecution, partitionKeyPath, pkBuilder.Build(), continuationToken, expectException); } private static OptimisticDirectExecutionTestInput CreateInput( @@ -430,9 +539,10 @@ private static OptimisticDirectExecutionTestInput CreateInput( bool expectedOptimisticDirectExecution, string partitionKeyPath, Cosmos.PartitionKey partitionKeyValue, - CosmosElement continuationToken = null) + CosmosElement continuationToken = null, + bool expectException = false) { - return new OptimisticDirectExecutionTestInput(description, query, new SqlQuerySpec(query), expectedOptimisticDirectExecution, partitionKeyPath, partitionKeyValue, continuationToken); + return new OptimisticDirectExecutionTestInput(description, query, new SqlQuerySpec(query), expectedOptimisticDirectExecution, partitionKeyPath, partitionKeyValue, continuationToken, expectException); } public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirectExecutionTestInput input) @@ -453,6 +563,11 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect bool result = queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton).Result; + if (input.ExpectException) + { + return new OptimisticDirectExecutionTestOutput(!queryPipelineStage.Current.Failed); + } + if (input.ExpectedOptimisticDirectExecution) { Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); @@ -553,6 +668,7 @@ public sealed class OptimisticDirectExecutionTestInput : BaselineTestInput internal PartitionKeyRangeIdentity PartitionKeyRangeId { get; set; } internal string Query { get; set; } internal CosmosElement ContinuationToken { get; set; } + internal bool ExpectException { get; set; } internal OptimisticDirectExecutionTestInput( string description, @@ -561,7 +677,8 @@ internal OptimisticDirectExecutionTestInput( bool expectedOptimisticDirectExecution, string partitionKeyPath, Cosmos.PartitionKey partitionKeyValue, - CosmosElement continuationToken) + CosmosElement continuationToken, + bool expectException) : base(description) { this.PartitionKeyDefinition = new PartitionKeyDefinition() @@ -578,6 +695,7 @@ internal OptimisticDirectExecutionTestInput( this.Query = query; this.PartitionKeyValue = partitionKeyValue; this.ContinuationToken = continuationToken; + this.ExpectException = expectException; } public override void SerializeAsXml(XmlWriter xmlWriter) @@ -681,9 +799,29 @@ public override Task> TryGetOverlappingRangesAs throw new NotImplementedException(); } - public override Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, CancellationToken cancellationToken) + public override async Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, CancellationToken cancellationToken) { - throw new NotImplementedException(); + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = new PartitionedQueryExecutionInfo() + { + QueryInfo = new QueryInfo() + { + Aggregates = null, + DistinctType = DistinctQueryType.None, + GroupByAliases = null, + GroupByAliasToAggregateType = null, + GroupByExpressions = null, + HasSelectValue = false, + Limit = null, + Offset = null, + OrderBy = null, + OrderByExpressions = null, + RewrittenQuery = null, + Top = null, + }, + QueryRanges = new List>(), + }; + + return TryCatch.FromResult(partitionedQueryExecutionInfo); } } } From f9fbd8abc16c889e440cec96c2f14cd7af1e444a Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Mon, 19 Dec 2022 15:07:14 -0800 Subject: [PATCH 18/37] Code cleanup --- .../Core/Pipeline/CosmosQueryExecutionContextFactory.cs | 2 +- .../OptimisticDirectExecutionQueryPipelineStage.cs | 4 +--- .../Query/OptimisticDirectExecutionQueryBaselineTests.cs | 6 ++---- 3 files changed, 4 insertions(+), 8 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 7028c966bf..267f00d609 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -417,7 +417,7 @@ private static TryCatch TryCreateOptimisticDirectExecutionC cancellationToken: cancellationToken); } - public static async Task> TryCreateOdeFallbackPipelineAsync( + internal static async Task> TryCreateOdeFallbackPipelineAsync( DocumentContainer documentContainer, CosmosQueryContext cosmosQueryContext, ContainerQueryProperties containerQueryProperties, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 1fbc8a28bf..c8793b4a73 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -19,7 +19,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; using Microsoft.Azure.Cosmos.Tracing; - using Microsoft.Azure.Documents; internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage { @@ -119,7 +118,6 @@ public static TryCatch MonadicCreate( } OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, fallbackQueryPipelineStageFactory, inputParameters.InitialUserContinuationToken); - return TryCatch.FromResult(odePipelineStageMonadicCreate); } @@ -132,7 +130,7 @@ private class OptimisticDirectExecutionQueryPipelineImpl : IQueryPipelineStage { private readonly QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator; - internal OptimisticDirectExecutionQueryPipelineImpl( + private OptimisticDirectExecutionQueryPipelineImpl( QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator) { this.queryPartitionRangePageAsyncEnumerator = queryPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(queryPartitionRangePageAsyncEnumerator)); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index d421f6f9ce..dc4fc8619f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -16,7 +16,6 @@ using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; @@ -30,7 +29,6 @@ using Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; using Newtonsoft.Json; - using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory; [TestClass] public class OptimisticDirectExecutionQueryBaselineTests : BaselineTests @@ -250,7 +248,7 @@ public async Task TestHandlingOfFailedFallbackPipeline() failureConfigs: new FlakyDocumentContainer.FailureConfigs( inject429s: false, injectEmptyPages: false, - shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter == 1 ? goneException : null))); + shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter == 1 || moveNextAsyncCounter == 4 ? goneException : null))); QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input[0], queryRequestOptions); @@ -268,7 +266,7 @@ public async Task TestHandlingOfFailedFallbackPipeline() queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), fallbackQueryPipelineStageFactory: (continuationToken) => { - InputParameters updatedInputParameters = new InputParameters( + CosmosQueryExecutionContextFactory.InputParameters updatedInputParameters = new CosmosQueryExecutionContextFactory.InputParameters( inputParameters.SqlQuerySpec, CosmosString.Create("asdf"), inputParameters.InitialFeedRange, From c2308378998dc6f9c11c36ca3f8f83a93c6f64ba Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 21 Dec 2022 13:11:00 -0800 Subject: [PATCH 19/37] Added logic for handling non ODE continuation tokens --- .../CosmosQueryExecutionContextFactory.cs | 40 ++++++++++++++--- ...imisticDirectExecutionContinuationToken.cs | 2 +- ...egativeOptimisticDirectExecutionOutput.xml | 15 ++++++- ...ositiveOptimisticDirectExecutionOutput.xml | 2 +- ...misticDirectExecutionQueryBaselineTests.cs | 44 ++++++++++++++++--- 5 files changed, 87 insertions(+), 16 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 267f00d609..ec3f13a2ca 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -148,13 +148,30 @@ private static async Task> TryCreateCoreContextAsy // Test code added to confirm the correct pipeline is being utilized SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); - return TryCreateOptimisticDirectExecutionContext( + TryCatch pipelineStage = TryCreateOptimisticDirectExecutionContext( documentContainer, cosmosQueryContext, containerQueryProperties, inputParameters, targetRange, cancellationToken); + + // A malformed continuation token exception would happen for 2 reasons here + // 1. the token is actually malformed + // 2. Its a non Ode continuation token + // In both cases, Ode pipeline delegates the work to the Specialized pipeline + // as Ode ppipeline should not take over execution while some other pipeline is already handling it + if (pipelineStage.Failed && pipelineStage.InnerMostException is MalformedContinuationTokenException) + { + return await TryCreateSpecializedPipelineAsync(documentContainer, + cosmosQueryContext, + containerQueryProperties, + inputParameters, + trace, + cancellationToken); + } + + return pipelineStage; } PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; @@ -228,7 +245,7 @@ private static async Task> TryCreateCoreContextAsy } } - partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoFromGatewayOrServiceInteropAsync( + partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoAsync( cosmosQueryContext, inputParameters, containerQueryProperties, @@ -306,6 +323,17 @@ private static async Task> TryCreateFromPartitione targetRange, cancellationToken); + if (tryCreatePipelineStage.Failed && tryCreatePipelineStage.InnerMostException is MalformedContinuationTokenException) + { + tryCreatePipelineStage = CallSpecializedPipeline( + documentContainer, + cosmosQueryContext, + inputParameters, + targetRanges, + partitionedQueryExecutionInfo, + cancellationToken); + } + return tryCreatePipelineStage; } @@ -404,7 +432,7 @@ private static TryCatch TryCreateOptimisticDirectExecutionC // In fallback scenario, the Specialized pipeline is always invoked Task> tryCreateContext = - CosmosQueryExecutionContextFactory.TryCreateOdeFallbackPipelineAsync( + CosmosQueryExecutionContextFactory.TryCreateSpecializedPipelineAsync( documentContainer, cosmosQueryContext, containerQueryProperties, @@ -417,7 +445,7 @@ private static TryCatch TryCreateOptimisticDirectExecutionC cancellationToken: cancellationToken); } - internal static async Task> TryCreateOdeFallbackPipelineAsync( + internal static async Task> TryCreateSpecializedPipelineAsync( DocumentContainer documentContainer, CosmosQueryContext cosmosQueryContext, ContainerQueryProperties containerQueryProperties, @@ -425,7 +453,7 @@ internal static async Task> TryCreateOdeFallbackPi ITrace trace, CancellationToken cancellationToken) { - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoFromGatewayOrServiceInteropAsync( + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoAsync( cosmosQueryContext, inputParameters, containerQueryProperties, @@ -534,7 +562,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx requestCancellationToken: cancellationToken); } - private static async Task GetPartitionedQueryExecutionInfoFromGatewayOrServiceInteropAsync( + private static async Task GetPartitionedQueryExecutionInfoAsync( CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, ContainerQueryProperties containerQueryProperties, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs index c81c378029..e4b4f15bc3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs @@ -43,7 +43,7 @@ public static CosmosElement ToCosmosElement(OptimisticDirectExecutionContinuatio public static TryCatch TryCreateFromCosmosElement(CosmosElement cosmosElement) { CosmosObject cosmosObjectContinuationToken = cosmosElement as CosmosObject; - if (cosmosObjectContinuationToken == null) + if (cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)) { return TryCatch.FromException( new MalformedChangeFeedContinuationTokenException( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml index a6905f99b3..89fe5f4553 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml @@ -40,7 +40,20 @@ - Exception with malformed continuation token + MalformedException with Parallel continuation token + SELECT * FROM c + + /pk + + Hash + + + false + + + + + MalformedException with OrderBy continuation token SELECT * FROM c /pk diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.PositiveOptimisticDirectExecutionOutput.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.PositiveOptimisticDirectExecutionOutput.xml index cedd618124..2482925bd8 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.PositiveOptimisticDirectExecutionOutput.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.PositiveOptimisticDirectExecutionOutput.xml @@ -53,7 +53,7 @@ - Single Partition Key and continuation token + Single Partition Key and Ode continuation token SELECT * FROM c /pk diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index dc4fc8619f..70c9b9f70f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -8,6 +8,7 @@ using System.Threading; using System.Threading.Tasks; using System.Xml; + using System.Xml.Linq; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query; @@ -16,6 +17,8 @@ using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; @@ -37,7 +40,7 @@ public class OptimisticDirectExecutionQueryBaselineTests : BaselineTests("A", "B", true, false)); + + OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( + parallelContinuationToken, + new List() { new OrderByItem(CosmosObject.Create(new Dictionary() { { "item", CosmosArray.Parse("[42, 37]") } })) }, + rid: "rid", + skipCount: 42, + filter: "filter"); + + CosmosElement cosmosElementOrderByContinuationToken = CosmosArray.Create( + new List() + { + OrderByContinuationToken.ToCosmosElement(orderByContinuationToken) + }); + + CosmosElement cosmosElementParallelContinuationToken = ParallelContinuationToken.ToCosmosElement(parallelContinuationToken); + List testVariations = new List { CreateInput( @@ -110,14 +132,22 @@ public void NegativeOptimisticDirectExecutionOutput() partitionKeyPath: @"/pk", partitionKeyValue: null), - // Fail to execute due to MalFormedContinuationToken Exception CreateInput( - description: @"Exception with malformed continuation token", + description: @"MalformedException with Parallel continuation token", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: "a", + continuationToken: cosmosElementParallelContinuationToken, + expectException: true), + + CreateInput( + description: @"MalformedException with OrderBy continuation token", query: "SELECT * FROM c", expectedOptimisticDirectExecution: false, partitionKeyPath: @"/pk", partitionKeyValue: "a", - continuationToken: CosmosString.Create("asdf"), + continuationToken: cosmosElementOrderByContinuationToken, expectException: true), }; this.ExecuteTestSuite(testVariations); @@ -282,7 +312,7 @@ public async Task TestHandlingOfFailedFallbackPipeline() inputParameters.TestInjections); Task> tryCreateContext = - CosmosQueryExecutionContextFactory.TryCreateOdeFallbackPipelineAsync( + CosmosQueryExecutionContextFactory.TryCreateSpecializedPipelineAsync( inMemoryCollection, cosmosQueryContextCore, containerQueryProperties, From c66be0f42048bdba905e218dec16e9e997f7ebdd Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 21 Dec 2022 13:16:17 -0800 Subject: [PATCH 20/37] Moved delegate away from member variables --- .../OptimisticDirectExecutionQueryPipelineStage.cs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index c8793b4a73..30ee446d47 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -29,7 +29,6 @@ private enum ExecutionState } private const string optimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; - public delegate Task> FallbackQueryPipelineStageFactory(CosmosElement continuationToken); private readonly FallbackQueryPipelineStageFactory queryPipelineStageFactory; private TryCatch inner; private CosmosElement continuationToken; @@ -43,6 +42,8 @@ private OptimisticDirectExecutionQueryPipelineStage(TryCatch> FallbackQueryPipelineStageFactory(CosmosElement continuationToken); + public TryCatch Current => this.inner.Try(pipelineStage => pipelineStage.Current); public ValueTask DisposeAsync() From 758c07f907f880a3d8f778eb688670a1caa83c70 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Thu, 22 Dec 2022 15:45:50 -0800 Subject: [PATCH 21/37] Added tests for Merge case --- ...imisticDirectExecutionContinuationToken.cs | 2 +- ...egativeOptimisticDirectExecutionOutput.xml | 6 +-- ...misticDirectExecutionQueryBaselineTests.cs | 54 +++++++++---------- 3 files changed, 29 insertions(+), 33 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs index e4b4f15bc3..4d6164f2ca 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs @@ -46,7 +46,7 @@ public static TryCatch TryCreateFrom if (cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)) { return TryCatch.FromException( - new MalformedChangeFeedContinuationTokenException( + new MalformedContinuationTokenException( message: $"Malformed Continuation Token")); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml index 89fe5f4553..7a8259eefc 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/OptimisticDirectExecutionQueryBaselineTests.NegativeOptimisticDirectExecutionOutput.xml @@ -40,7 +40,7 @@ - MalformedException with Parallel continuation token + Single Partition Key with Parallel continuation token SELECT * FROM c /pk @@ -53,8 +53,8 @@ - MalformedException with OrderBy continuation token - SELECT * FROM c + Single Partition Key with OrderBy continuation token + SELECT * FROM c ORDER BY c._ts /pk diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 70c9b9f70f..a605747d19 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -96,8 +96,8 @@ public void NegativeOptimisticDirectExecutionOutput() OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( parallelContinuationToken, - new List() { new OrderByItem(CosmosObject.Create(new Dictionary() { { "item", CosmosArray.Parse("[42, 37]") } })) }, - rid: "rid", + new List() { new OrderByItem(CosmosObject.Create(new Dictionary() { { "item", CosmosString.Create("asdf") } })) }, + rid: "43223532", skipCount: 42, filter: "filter"); @@ -107,8 +107,6 @@ public void NegativeOptimisticDirectExecutionOutput() OrderByContinuationToken.ToCosmosElement(orderByContinuationToken) }); - CosmosElement cosmosElementParallelContinuationToken = ParallelContinuationToken.ToCosmosElement(parallelContinuationToken); - List testVariations = new List { CreateInput( @@ -133,22 +131,20 @@ public void NegativeOptimisticDirectExecutionOutput() partitionKeyValue: null), CreateInput( - description: @"MalformedException with Parallel continuation token", + description: @"Single Partition Key with Parallel continuation token", query: "SELECT * FROM c", expectedOptimisticDirectExecution: false, partitionKeyPath: @"/pk", partitionKeyValue: "a", - continuationToken: cosmosElementParallelContinuationToken, - expectException: true), - + continuationToken: CosmosArray.Create(new List() { ParallelContinuationToken.ToCosmosElement(parallelContinuationToken) })), + CreateInput( - description: @"MalformedException with OrderBy continuation token", - query: "SELECT * FROM c", + description: @"Single Partition Key with OrderBy continuation token", + query: "SELECT * FROM c ORDER BY c._ts", expectedOptimisticDirectExecution: false, partitionKeyPath: @"/pk", partitionKeyValue: "a", - continuationToken: cosmosElementOrderByContinuationToken, - expectException: true), + continuationToken: cosmosElementOrderByContinuationToken), }; this.ExecuteTestSuite(testVariations); } @@ -552,13 +548,12 @@ private static OptimisticDirectExecutionTestInput CreateInput( bool expectedOptimisticDirectExecution, string partitionKeyPath, string partitionKeyValue, - CosmosElement continuationToken = null, - bool expectException = false) + CosmosElement continuationToken = null) { PartitionKeyBuilder pkBuilder = new PartitionKeyBuilder(); pkBuilder.Add(partitionKeyValue); - return CreateInput(description, query, expectedOptimisticDirectExecution, partitionKeyPath, pkBuilder.Build(), continuationToken, expectException); + return CreateInput(description, query, expectedOptimisticDirectExecution, partitionKeyPath, pkBuilder.Build(), continuationToken); } private static OptimisticDirectExecutionTestInput CreateInput( @@ -567,10 +562,9 @@ private static OptimisticDirectExecutionTestInput CreateInput( bool expectedOptimisticDirectExecution, string partitionKeyPath, Cosmos.PartitionKey partitionKeyValue, - CosmosElement continuationToken = null, - bool expectException = false) + CosmosElement continuationToken = null) { - return new OptimisticDirectExecutionTestInput(description, query, new SqlQuerySpec(query), expectedOptimisticDirectExecution, partitionKeyPath, partitionKeyValue, continuationToken, expectException); + return new OptimisticDirectExecutionTestInput(description, query, new SqlQuerySpec(query), expectedOptimisticDirectExecution, partitionKeyPath, partitionKeyValue, continuationToken); } public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirectExecutionTestInput input) @@ -591,11 +585,6 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect bool result = queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton).Result; - if (input.ExpectException) - { - return new OptimisticDirectExecutionTestOutput(!queryPipelineStage.Current.Failed); - } - if (input.ExpectedOptimisticDirectExecution) { Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); @@ -696,7 +685,6 @@ public sealed class OptimisticDirectExecutionTestInput : BaselineTestInput internal PartitionKeyRangeIdentity PartitionKeyRangeId { get; set; } internal string Query { get; set; } internal CosmosElement ContinuationToken { get; set; } - internal bool ExpectException { get; set; } internal OptimisticDirectExecutionTestInput( string description, @@ -705,8 +693,7 @@ internal OptimisticDirectExecutionTestInput( bool expectedOptimisticDirectExecution, string partitionKeyPath, Cosmos.PartitionKey partitionKeyValue, - CosmosElement continuationToken, - bool expectException) + CosmosElement continuationToken) : base(description) { this.PartitionKeyDefinition = new PartitionKeyDefinition() @@ -723,7 +710,6 @@ internal OptimisticDirectExecutionTestInput( this.Query = query; this.PartitionKeyValue = partitionKeyValue; this.ContinuationToken = continuationToken; - this.ExpectException = expectException; } public override void SerializeAsXml(XmlWriter xmlWriter) @@ -829,6 +815,16 @@ public override Task> TryGetOverlappingRangesAs public override async Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, CancellationToken cancellationToken) { + IReadOnlyList orderByColumn = new List() + { + new SortOrder() + }; + + IReadOnlyList orderByExpression = new List() + { + "ORDER BY" + }; + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = new PartitionedQueryExecutionInfo() { QueryInfo = new QueryInfo() @@ -841,8 +837,8 @@ public override async Task> TryGetPartit HasSelectValue = false, Limit = null, Offset = null, - OrderBy = null, - OrderByExpressions = null, + OrderBy = sqlQuerySpec.QueryText.Contains("ORDER BY") ? orderByColumn : null, + OrderByExpressions = sqlQuerySpec.QueryText.Contains("ORDER BY") ? orderByExpression : null, RewrittenQuery = null, Top = null, }, From d442b9928d6493c90352ca9260d6c3b9f411c397 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 28 Dec 2022 15:21:00 -0800 Subject: [PATCH 22/37] Updated method names --- .../CosmosExceptionExtensions.cs} | 4 +- .../CosmosQueryExecutionContextFactory.cs | 228 ++++++++++-------- ...OrderByCrossPartitionQueryPipelineStage.cs | 2 +- ...misticDirectExecutionQueryPipelineStage.cs | 12 +- ...misticDirectExecutionQueryBaselineTests.cs | 2 +- 5 files changed, 137 insertions(+), 111 deletions(-) rename Microsoft.Azure.Cosmos/src/Query/Core/{Utils.cs => Pipeline/CosmosExceptionExtensions.cs} (86%) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Utils.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs similarity index 86% rename from Microsoft.Azure.Cosmos/src/Query/Core/Utils.cs rename to Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs index 93733f7bf4..93326fe355 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Utils.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs @@ -5,10 +5,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core { using System; - using System.Collections.Generic; - using System.Text; - internal static class Utils + internal static class CosmosExceptionExtensions { public static bool IsPartitionSplitException(this Exception ex) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index ec3f13a2ca..bada456636 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos; + using Microsoft.Azure.Cosmos.Core.Collections; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core; @@ -27,6 +28,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.SqlObjects; using Microsoft.Azure.Cosmos.SqlObjects.Visitors; using Microsoft.Azure.Cosmos.Tracing; + using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory; internal static class CosmosQueryExecutionContextFactory { @@ -136,7 +138,7 @@ private static async Task> TryCreateCoreContextAsy cancellationToken); cosmosQueryContext.ContainerResourceId = containerQueryProperties.ResourceId; - Documents.PartitionKeyRange targetRange = await GetTargetRangeOptimisticDirectExecutionAsync( + Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync( inputParameters, queryPlanFromContinuationToken, cosmosQueryContext, @@ -145,33 +147,14 @@ private static async Task> TryCreateCoreContextAsy if (targetRange != null) { - // Test code added to confirm the correct pipeline is being utilized - SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); - - TryCatch pipelineStage = TryCreateOptimisticDirectExecutionContext( - documentContainer, - cosmosQueryContext, - containerQueryProperties, - inputParameters, - targetRange, - cancellationToken); - - // A malformed continuation token exception would happen for 2 reasons here - // 1. the token is actually malformed - // 2. Its a non Ode continuation token - // In both cases, Ode pipeline delegates the work to the Specialized pipeline - // as Ode ppipeline should not take over execution while some other pipeline is already handling it - if (pipelineStage.Failed && pipelineStage.InnerMostException is MalformedContinuationTokenException) - { - return await TryCreateSpecializedPipelineAsync(documentContainer, - cosmosQueryContext, - containerQueryProperties, - inputParameters, - trace, - cancellationToken); - } - - return pipelineStage; + return await TryCreateExecutionContextAsync( + documentContainer, + cosmosQueryContext, + containerQueryProperties, + inputParameters, + targetRange, + trace, + cancellationToken); } PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; @@ -304,7 +287,7 @@ private static async Task> TryCreateFromPartitione TryCatch tryCreatePipelineStage; - Documents.PartitionKeyRange targetRange = await GetTargetRangeOptimisticDirectExecutionAsync( + Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync( inputParameters, partitionedQueryExecutionInfo, cosmosQueryContext, @@ -313,28 +296,15 @@ private static async Task> TryCreateFromPartitione if (targetRange != null) { - SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); - - tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreateOptimisticDirectExecutionContext( + return await TryCreateExecutionContextAsync( documentContainer, cosmosQueryContext, containerQueryProperties, inputParameters, targetRange, - cancellationToken); - - if (tryCreatePipelineStage.Failed && tryCreatePipelineStage.InnerMostException is MalformedContinuationTokenException) - { - tryCreatePipelineStage = CallSpecializedPipeline( - documentContainer, - cosmosQueryContext, - inputParameters, - targetRanges, - partitionedQueryExecutionInfo, - cancellationToken); - } - - return tryCreatePipelineStage; + trace, + cancellationToken, + partitionedQueryExecutionInfo); } if (createPassthroughQuery) @@ -349,13 +319,74 @@ private static async Task> TryCreateFromPartitione } else { - tryCreatePipelineStage = CallSpecializedPipeline(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken); + tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken); + } + + return tryCreatePipelineStage; + } + + private static async Task> TryCreateExecutionContextAsync( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + ContainerQueryProperties containerQueryProperties, + InputParameters inputParameters, + Documents.PartitionKeyRange targetRange, + ITrace trace, + CancellationToken cancellationToken, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = null) + { + // Test code added to confirm the correct pipeline is being utilized + SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); + + TryCatch tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreateOptimisticDirectExecutionContext( + documentContainer, + cosmosQueryContext, + containerQueryProperties, + inputParameters, + targetRange, + cancellationToken); + + // A malformed continuation token exception would happen for 2 reasons here + // 1. the token is actually malformed + // 2. Its a non Ode continuation token + // In both cases, Ode pipeline delegates the work to the Specialized pipeline + // as Ode ppipeline should not take over execution while some other pipeline is already handling it + if (tryCreatePipelineStage.Failed && tryCreatePipelineStage.InnerMostException is MalformedContinuationTokenException) + { + SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); + + if (partitionedQueryExecutionInfo != null) + { + List targetRanges = new List + { + targetRange + }; + + tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext( + documentContainer, + cosmosQueryContext, + inputParameters, + targetRanges, + partitionedQueryExecutionInfo, + cancellationToken); + } + else + { + tryCreatePipelineStage = await TryCreateSpecializedDocumentQueryExecutionContextAsync( + documentContainer, + cosmosQueryContext, + containerQueryProperties, + inputParameters, + trace, + cancellationToken); + } + } return tryCreatePipelineStage; } - private static TryCatch CallSpecializedPipeline( + private static TryCatch TryCreateSpecializedDocumentQueryExecutionContext( DocumentContainer documentContainer, CosmosQueryContext cosmosQueryContext, InputParameters inputParameters, @@ -399,53 +430,7 @@ private static TryCatch CallSpecializedPipeline( cancellationToken); } - private static TryCatch TryCreateOptimisticDirectExecutionContext( - DocumentContainer documentContainer, - CosmosQueryContext cosmosQueryContext, - ContainerQueryProperties containerQueryProperties, - InputParameters inputParameters, - Documents.PartitionKeyRange targetRange, - CancellationToken cancellationToken) - { - // Return a OptimisticDirectExecution context - return OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( - documentContainer: documentContainer, - inputParameters: inputParameters, - targetRange: new FeedRangeEpk(targetRange.ToRange()), - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), - fallbackQueryPipelineStageFactory: (continuationToken) => - { - InputParameters updatedInputParameters = new InputParameters( - inputParameters.SqlQuerySpec, - continuationToken, - inputParameters.InitialFeedRange, - inputParameters.MaxConcurrency, - inputParameters.MaxItemCount, - inputParameters.MaxBufferedItemCount, - inputParameters.PartitionKey, - inputParameters.Properties, - inputParameters.PartitionedQueryExecutionInfo, - inputParameters.ExecutionEnvironment, - inputParameters.ReturnResultsInDeterministicOrder, - inputParameters.ForcePassthrough, - inputParameters.TestInjections); - - // In fallback scenario, the Specialized pipeline is always invoked - Task> tryCreateContext = - CosmosQueryExecutionContextFactory.TryCreateSpecializedPipelineAsync( - documentContainer, - cosmosQueryContext, - containerQueryProperties, - updatedInputParameters, - NoOpTrace.Singleton, - default); - - return tryCreateContext; - }, - cancellationToken: cancellationToken); - } - - internal static async Task> TryCreateSpecializedPipelineAsync( + internal static async Task> TryCreateSpecializedDocumentQueryExecutionContextAsync( DocumentContainer documentContainer, CosmosQueryContext cosmosQueryContext, ContainerQueryProperties containerQueryProperties, @@ -469,7 +454,7 @@ internal static async Task> TryCreateSpecializedPi inputParameters.InitialFeedRange, trace); - return CallSpecializedPipeline( + return TryCreateSpecializedDocumentQueryExecutionContext( documentContainer, cosmosQueryContext, inputParameters, @@ -478,6 +463,37 @@ internal static async Task> TryCreateSpecializedPi cancellationToken); } + private static TryCatch TryCreateOptimisticDirectExecutionContext( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + ContainerQueryProperties containerQueryProperties, + InputParameters inputParameters, + Documents.PartitionKeyRange targetRange, + CancellationToken cancellationToken) + { + // Return a OptimisticDirectExecution context + return OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + inputParameters: inputParameters, + targetRange: new FeedRangeEpk(targetRange.ToRange()), + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount), + fallbackQueryPipelineStageFactory: (continuationToken) => + { + // In fallback scenario, the Specialized pipeline is always invoked + Task> tryCreateContext = + CosmosQueryExecutionContextFactory.TryCreateSpecializedDocumentQueryExecutionContextAsync( + documentContainer, + cosmosQueryContext, + containerQueryProperties, + inputParameters.WithContinuationToken(continuationToken), + NoOpTrace.Singleton, + default); + + return tryCreateContext; + }, + cancellationToken: cancellationToken); + } + private static TryCatch TryCreatePassthroughQueryExecutionContext( DocumentContainer documentContainer, InputParameters inputParameters, @@ -729,7 +745,7 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP return partitionKeyDefinition; } - private static async Task GetTargetRangeOptimisticDirectExecutionAsync( + private static async Task TryGetTargetRangeOptimisticDirectExecutionAsync( InputParameters inputParameters, PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, CosmosQueryContext cosmosQueryContext, @@ -857,6 +873,24 @@ public InputParameters( public bool ReturnResultsInDeterministicOrder { get; } public TestInjections TestInjections { get; } public bool ForcePassthrough { get; } + + public InputParameters WithContinuationToken(CosmosElement token) + { + return new InputParameters( + this.SqlQuerySpec, + token, + this.InitialFeedRange, + this.MaxConcurrency, + this.MaxItemCount, + this.MaxBufferedItemCount, + this.PartitionKey, + this.Properties, + this.PartitionedQueryExecutionInfo, + this.ExecutionEnvironment, + this.ReturnResultsInDeterministicOrder, + this.ForcePassthrough, + this.TestInjections); + } } internal sealed class AggregateProjectionDetector diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index 0efa1cd950..ef3811b5c7 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -1153,7 +1153,7 @@ private static bool IsSplitException(Exception exception) exception = exception.InnerException; } - return Utils.IsPartitionSplitException(exception); + return CosmosExceptionExtensions.IsPartitionSplitException(exception); } public void SetCancellationToken(CancellationToken cancellationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 30ee446d47..07caf61aef 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -56,20 +56,19 @@ public async ValueTask MoveNextAsync(ITrace trace) bool success = await this.inner.Result.MoveNextAsync(trace); if (success) { - this.SaveContinuation(this.Current.Result.State?.Value); + this.continuationToken = this.Current.Result.State?.Value; } else if (this.executionState == ExecutionState.OptimisticDirectExecution) { - if (Utils.IsPartitionSplitException(this.Current.InnerMostException)) + if (CosmosExceptionExtensions.IsPartitionSplitException(this.Current.InnerMostException)) { this.inner = await this.queryPipelineStageFactory(this.TryUnwrapContinuationToken().Result); + this.executionState = ExecutionState.SpecializedDocumentQueryExecution; if (this.inner.Failed) { - this.inner = TryCatch.FromException(this.inner.Exception); return false; } - this.executionState = ExecutionState.SpecializedDocumentQueryExecution; success = await this.inner.Result.MoveNextAsync(trace); } } @@ -122,11 +121,6 @@ public static TryCatch MonadicCreate( return TryCatch.FromResult(odePipelineStageMonadicCreate); } - private void SaveContinuation(CosmosElement continuationToken) - { - this.continuationToken = continuationToken; - } - private class OptimisticDirectExecutionQueryPipelineImpl : IQueryPipelineStage { private readonly QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index a605747d19..fd7a7953ba 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -308,7 +308,7 @@ public async Task TestHandlingOfFailedFallbackPipeline() inputParameters.TestInjections); Task> tryCreateContext = - CosmosQueryExecutionContextFactory.TryCreateSpecializedPipelineAsync( + CosmosQueryExecutionContextFactory.TryCreateSpecializedDocumentQueryExecutionContextAsync( inMemoryCollection, cosmosQueryContextCore, containerQueryProperties, From d1dcd7e6c6f5617c5c16b9e5f90c746d0e600f15 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Thu, 29 Dec 2022 16:30:03 -0800 Subject: [PATCH 23/37] Added checks for tryCatch --- ...misticDirectExecutionQueryPipelineStage.cs | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 07caf61aef..8d36a3d13a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -48,12 +48,12 @@ private OptimisticDirectExecutionQueryPipelineStage(TryCatch(pipelineStage => pipelineStage.DisposeAsync()).Result; + return this.inner.Failed ? default : this.inner.Result.DisposeAsync(); } public async ValueTask MoveNextAsync(ITrace trace) { - bool success = await this.inner.Result.MoveNextAsync(trace); + bool success = !this.inner.Failed && await this.inner.Result.MoveNextAsync(trace); if (success) { this.continuationToken = this.Current.Result.State?.Value; @@ -78,19 +78,18 @@ public async ValueTask MoveNextAsync(ITrace trace) public void SetCancellationToken(CancellationToken cancellationToken) { - this.inner.Try((pipelineStage) => pipelineStage = this.inner.Result).Result.SetCancellationToken(cancellationToken); + if (this.inner.Failed) + { + throw new ArgumentNullException($"PipelineStage failed with '{this.inner.InnerMostException}'."); + } + + this.inner.Result.SetCancellationToken(cancellationToken); } public TryCatch TryUnwrapContinuationToken() { Debug.Assert(this.continuationToken is CosmosObject); - if (!((CosmosObject)this.continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken)) - { - return TryCatch.FromException( - new FormatException( - $"Unable to convert CosmosObject '{optimisticDirectExecutionToken}' to CosmosElement.")); - } - + ((CosmosObject)this.continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken); CosmosArray cosmosElementFallbackToken = CosmosArray.Create(specializedContinuationToken); return TryCatch.FromResult(cosmosElementFallbackToken); } From 294d17307aed2a5f67b531350241804fc42ad332 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Tue, 3 Jan 2023 12:20:52 -0800 Subject: [PATCH 24/37] Updated SetCancellationToken() to use Try --- .../OptimisticDirectExecutionQueryPipelineStage.cs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 8d36a3d13a..7d23b897e1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -53,7 +53,10 @@ public ValueTask DisposeAsync() public async ValueTask MoveNextAsync(ITrace trace) { - bool success = !this.inner.Failed && await this.inner.Result.MoveNextAsync(trace); + TryCatch result = this.inner.Try((pipelineStage) => pipelineStage = this.inner.Result) + .Try(pipelineStage => pipelineStage.MoveNextAsync(trace)).Try(moveNextValue => moveNextValue.Result); + bool success = result.Succeeded && result.Result; + if (success) { this.continuationToken = this.Current.Result.State?.Value; @@ -78,12 +81,8 @@ public async ValueTask MoveNextAsync(ITrace trace) public void SetCancellationToken(CancellationToken cancellationToken) { - if (this.inner.Failed) - { - throw new ArgumentNullException($"PipelineStage failed with '{this.inner.InnerMostException}'."); - } - - this.inner.Result.SetCancellationToken(cancellationToken); + this.inner.Try((pipelineStage) => pipelineStage = this.inner.Result) + .Try(pipelineStage => pipelineStage.SetCancellationToken(cancellationToken)); } public TryCatch TryUnwrapContinuationToken() From 0621090ab5636830c9ab99925e9bfc40bb7ecd70 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 4 Jan 2023 10:53:24 -0800 Subject: [PATCH 25/37] Updated TryUnwrapContinuationToken() --- .../Pipeline/CosmosExceptionExtensions.cs | 2 +- .../CosmosQueryExecutionContextFactory.cs | 19 +++++++++---------- ...OrderByCrossPartitionQueryPipelineStage.cs | 2 +- ...imisticDirectExecutionContinuationToken.cs | 2 +- ...misticDirectExecutionQueryPipelineStage.cs | 19 +++++++++---------- 5 files changed, 21 insertions(+), 23 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs index 93326fe355..98e9023226 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs @@ -2,7 +2,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // ------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline { using System; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index bada456636..73547b376f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -28,7 +28,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using Microsoft.Azure.Cosmos.SqlObjects; using Microsoft.Azure.Cosmos.SqlObjects.Visitors; using Microsoft.Azure.Cosmos.Tracing; - using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory; internal static class CosmosQueryExecutionContextFactory { @@ -149,6 +148,7 @@ private static async Task> TryCreateCoreContextAsy { return await TryCreateExecutionContextAsync( documentContainer, + null, cosmosQueryContext, containerQueryProperties, inputParameters, @@ -291,23 +291,22 @@ private static async Task> TryCreateFromPartitione inputParameters, partitionedQueryExecutionInfo, cosmosQueryContext, - containerQueryProperties, + containerQueryProperties, trace); if (targetRange != null) { - return await TryCreateExecutionContextAsync( + tryCreatePipelineStage = await TryCreateExecutionContextAsync( documentContainer, + partitionedQueryExecutionInfo, cosmosQueryContext, containerQueryProperties, inputParameters, targetRange, trace, - cancellationToken, - partitionedQueryExecutionInfo); + cancellationToken); } - - if (createPassthroughQuery) + else if (createPassthroughQuery) { SetTestInjectionPipelineType(inputParameters, Passthrough); @@ -327,13 +326,13 @@ private static async Task> TryCreateFromPartitione private static async Task> TryCreateExecutionContextAsync( DocumentContainer documentContainer, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, CosmosQueryContext cosmosQueryContext, ContainerQueryProperties containerQueryProperties, InputParameters inputParameters, Documents.PartitionKeyRange targetRange, ITrace trace, - CancellationToken cancellationToken, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = null) + CancellationToken cancellationToken) { // Test code added to confirm the correct pipeline is being utilized SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); @@ -353,7 +352,7 @@ private static async Task> TryCreateExecutionConte // as Ode ppipeline should not take over execution while some other pipeline is already handling it if (tryCreatePipelineStage.Failed && tryCreatePipelineStage.InnerMostException is MalformedContinuationTokenException) { - SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); + SetTestInjectionPipelineType(inputParameters, Specialized); if (partitionedQueryExecutionInfo != null) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs index ef3811b5c7..10c3b0ee5e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs @@ -1153,7 +1153,7 @@ private static bool IsSplitException(Exception exception) exception = exception.InnerException; } - return CosmosExceptionExtensions.IsPartitionSplitException(exception); + return exception.IsPartitionSplitException(); } public void SetCancellationToken(CancellationToken cancellationToken) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs index 4d6164f2ca..59f87295f3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs @@ -47,7 +47,7 @@ public static TryCatch TryCreateFrom { return TryCatch.FromException( new MalformedContinuationTokenException( - message: $"Malformed Continuation Token")); + message: $"Malformed Continuation Token: Expected OptimisticDirectExecutionToken\r\n")); } TryCatch inner = ParallelContinuationToken.TryCreateFromCosmosElement(cosmosObjectContinuationToken[OptimisticDirectExecutionToken]); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 7d23b897e1..31624fb04a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -59,13 +59,13 @@ public async ValueTask MoveNextAsync(ITrace trace) if (success) { - this.continuationToken = this.Current.Result.State?.Value; + this.continuationToken = this.Current.Succeeded ? this.Current.Result.State?.Value : null; } - else if (this.executionState == ExecutionState.OptimisticDirectExecution) + else if (this.executionState == ExecutionState.OptimisticDirectExecution && result.Succeeded) { - if (CosmosExceptionExtensions.IsPartitionSplitException(this.Current.InnerMostException)) + if (this.Current.InnerMostException.IsPartitionSplitException()) { - this.inner = await this.queryPipelineStageFactory(this.TryUnwrapContinuationToken().Result); + this.inner = await this.queryPipelineStageFactory(this.TryUnwrapContinuationToken()); this.executionState = ExecutionState.SpecializedDocumentQueryExecution; if (this.inner.Failed) { @@ -85,12 +85,11 @@ public void SetCancellationToken(CancellationToken cancellationToken) .Try(pipelineStage => pipelineStage.SetCancellationToken(cancellationToken)); } - public TryCatch TryUnwrapContinuationToken() + private CosmosElement TryUnwrapContinuationToken() { - Debug.Assert(this.continuationToken is CosmosObject); - ((CosmosObject)this.continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken); - CosmosArray cosmosElementFallbackToken = CosmosArray.Create(specializedContinuationToken); - return TryCatch.FromResult(cosmosElementFallbackToken); + CosmosObject cosmosObject = this.continuationToken as CosmosObject; + CosmosElement specializedContinuationToken = cosmosObject[optimisticDirectExecutionToken]; + return CosmosArray.Create(specializedContinuationToken); } public static TryCatch MonadicCreate( @@ -112,7 +111,7 @@ public static TryCatch MonadicCreate( if (pipelineStage.Failed) { - return TryCatch.FromException(pipelineStage.Exception); + return pipelineStage; } OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, fallbackQueryPipelineStageFactory, inputParameters.InitialUserContinuationToken); From 87c9916e093385680e3efa761fca9ec8c4fa5e27 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Fri, 6 Jan 2023 13:59:47 -0800 Subject: [PATCH 26/37] Removed changes in FlakyDocumentContainer.cs --- .../CosmosQueryExecutionContextFactory.cs | 2 +- .../Pagination/FlakyDocumentContainer.cs | 29 +-- ...misticDirectExecutionQueryBaselineTests.cs | 167 +++++++++--------- 3 files changed, 85 insertions(+), 113 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 73547b376f..93e95846c9 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -429,7 +429,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx cancellationToken); } - internal static async Task> TryCreateSpecializedDocumentQueryExecutionContextAsync( + private static async Task> TryCreateSpecializedDocumentQueryExecutionContextAsync( DocumentContainer documentContainer, CosmosQueryContext cosmosQueryContext, ContainerQueryProperties containerQueryProperties, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs index f7f7fbec1c..3ea96d7634 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs @@ -27,9 +27,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination /// internal sealed class FlakyDocumentContainer : IMonadicDocumentContainer { - private FailureConfigs failureConfigs; - private bool isODETest; - private int iterationCount = 0; + private readonly FailureConfigs failureConfigs; private readonly Random random; private static class Throttle @@ -66,12 +64,10 @@ private static class Throttle public FlakyDocumentContainer( IMonadicDocumentContainer documentContainer, - FailureConfigs failureConfigs, - bool isODETest = false) + FailureConfigs failureConfigs) { this.documentContainer = documentContainer ?? throw new ArgumentNullException(nameof(documentContainer)); this.failureConfigs = failureConfigs ?? throw new ArgumentNullException(nameof(failureConfigs)); - this.isODETest = isODETest; this.random = new Random(); } @@ -179,25 +175,10 @@ public async Task> MonadicQueryAsync( state: feedRangeState.State ?? StateForStartedButNoDocumentsReturned)); } - if (this.isODETest) - { - this.iterationCount++; - if (this.iterationCount == 2 || this.iterationCount == 5) - { - Exception failure = await this.ShouldReturnFailure(); - if (failure != null) - { - return TryCatch.FromException(failure); - } - } - } - else + Exception failure = await this.ShouldReturnFailure(); + if (failure != null) { - Exception failure = await this.ShouldReturnFailure(); - if (failure != null) - { - return TryCatch.FromException(failure); - } + return TryCatch.FromException(failure); } return await this.documentContainer.MonadicQueryAsync( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index fd7a7953ba..daefbd9499 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -5,10 +5,10 @@ using System.Collections.ObjectModel; using System.IO; using System.Linq; + using System.Net; using System.Threading; using System.Threading.Tasks; using System.Xml; - using System.Xml.Linq; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query; @@ -20,7 +20,6 @@ using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; @@ -253,91 +252,52 @@ public async Task TestPipelineForDistributedQueryAsync() public async Task TestHandlingOfFailedFallbackPipeline() { List documents = new List(); - int numItems = 100; - int moveNextAsyncCounter = 0; - - List input = new List + MergeTestUtil mergeTest = new MergeTestUtil { - CreateInput( - description: @"Single Partition Key and Value Field", - query: "SELECT * FROM c", - expectedOptimisticDirectExecution: true, - partitionKeyPath: @"/pk", - partitionKeyValue: "a", - continuationToken: null) + MoveNextCounter = 0, + IsFailedFallbackPipelineTest = true }; - CosmosException goneException = await CreateGoneException(); + int numItems = 100; + OptimisticDirectExecutionTestInput input = CreateInput( + description: @"Single Partition Key and Value Field", + query: "SELECT * FROM c", + expectedOptimisticDirectExecution: true, + partitionKeyPath: @"/pk", + partitionKeyValue: "a"); + + QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( numItems, multiPartition: false, failureConfigs: new FlakyDocumentContainer.FailureConfigs( inject429s: false, injectEmptyPages: false, - shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter == 1 || moveNextAsyncCounter == 4 ? goneException : null))); + shouldReturnFailure: mergeTest.ShouldReturnFailure)); - QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); - (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input[0], queryRequestOptions); - - ContainerQueryProperties containerQueryProperties = new ContainerQueryProperties( - null, - null, - new PartitionKeyDefinition(), - It.IsAny()); - - TryCatch monadicQueryPipelineStage = OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( - documentContainer: inMemoryCollection, - inputParameters: inputParameters, - targetRange: FeedRangeEpk.FullRange, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), - fallbackQueryPipelineStageFactory: (continuationToken) => - { - CosmosQueryExecutionContextFactory.InputParameters updatedInputParameters = new CosmosQueryExecutionContextFactory.InputParameters( - inputParameters.SqlQuerySpec, - CosmosString.Create("asdf"), - inputParameters.InitialFeedRange, - inputParameters.MaxConcurrency, - inputParameters.MaxItemCount, - inputParameters.MaxBufferedItemCount, - inputParameters.PartitionKey, - inputParameters.Properties, - inputParameters.PartitionedQueryExecutionInfo, - inputParameters.ExecutionEnvironment, - inputParameters.ReturnResultsInDeterministicOrder, - inputParameters.ForcePassthrough, - inputParameters.TestInjections); - - Task> tryCreateContext = - CosmosQueryExecutionContextFactory.TryCreateSpecializedDocumentQueryExecutionContextAsync( - inMemoryCollection, - cosmosQueryContextCore, - containerQueryProperties, - updatedInputParameters, - NoOpTrace.Singleton, - default); - - return tryCreateContext; - }, - cancellationToken: default); + IQueryPipelineStage queryPipelineStage = await GetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); - while (await monadicQueryPipelineStage.Result.MoveNextAsync(NoOpTrace.Singleton)) + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { - moveNextAsyncCounter++; - - TryCatch tryGetPage = monadicQueryPipelineStage.Result.Current; + mergeTest.MoveNextCounter++; + TryCatch tryGetPage = queryPipelineStage.Current; if (tryGetPage.Failed) { - // failure should never come till here. Should be handled before - Assert.Fail(); + if (mergeTest.MoveNextCounter == 2) + { + Assert.IsTrue(tryGetPage.InnerMostException.Message.Equals("Injected failure")); + Assert.AreNotEqual(numItems, documents.Count); + return; + } + else + { + Assert.Fail(); + } } documents.AddRange(tryGetPage.Result.Documents); } - - Assert.IsTrue(monadicQueryPipelineStage.Result.Current.Failed); - Assert.IsTrue(monadicQueryPipelineStage.Result.Current.InnerMostException is MalformedContinuationTokenException); - Assert.AreNotEqual(numItems, documents.Count); } // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing @@ -345,9 +305,13 @@ public async Task TestHandlingOfFailedFallbackPipeline() private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) { List documents = new List(); - int numItems = 100; - int moveNextAsyncCounter = 0; + MergeTestUtil mergeTest = new MergeTestUtil + { + MoveNextCounter = 0, + IsFailedFallbackPipelineTest = false + }; + int numItems = 100; OptimisticDirectExecutionTestInput input = CreateInput( description: @"Single Partition Key and Value Field", query: "SELECT * FROM c", @@ -355,23 +319,21 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa partitionKeyPath: @"/pk", partitionKeyValue: "a"); - CosmosException goneException = await CreateGoneException(); QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( numItems, multiPartition: isMultiPartition, failureConfigs: new FlakyDocumentContainer.FailureConfigs( inject429s: false, injectEmptyPages: false, - shouldReturnFailure: () => Task.FromResult(moveNextAsyncCounter == 1 || moveNextAsyncCounter == 4 ? goneException : null))); + shouldReturnFailure: mergeTest.ShouldReturnFailure)); IQueryPipelineStage queryPipelineStage = await GetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { - moveNextAsyncCounter++; - if (moveNextAsyncCounter == 1) + mergeTest.MoveNextCounter++; + if (mergeTest.MoveNextCounter == 1) { Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); } @@ -467,17 +429,6 @@ private static async Task GetOdePipelineAsync(OptimisticDir return queryPipelineStage; } - private static async Task CreateGoneException() - { - string goneExceptionMessage = $"Epk Range: Partition does not exist at the given range."; - return new CosmosException( - message: goneExceptionMessage, - statusCode: System.Net.HttpStatusCode.Gone, - subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, - activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", - requestCharge: default); - } - private static async Task CreateDocumentContainerAsync( int numItems, bool multiPartition, @@ -496,7 +447,7 @@ private static async Task CreateDocumentContainerAsync( IMonadicDocumentContainer monadicDocumentContainer = new InMemoryContainer(partitionKeyDefinition); if (failureConfigs != null) { - monadicDocumentContainer = new FlakyDocumentContainer(monadicDocumentContainer, failureConfigs, isODETest: true); + monadicDocumentContainer = new FlakyDocumentContainer(monadicDocumentContainer, failureConfigs); } DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); @@ -657,6 +608,46 @@ private static QueryRequestOptions GetQueryRequestOptions(bool enableOptimisticD } }; } + + private class MergeTestUtil + { + public int MoveNextCounter { get; set; } + + public bool GoneExceptionCreated { get; set; } + + public bool InjectionFailureCreated { get; set; } + + public bool IsFailedFallbackPipelineTest { get; set; } + + public async Task ShouldReturnFailure() + { + + if (this.MoveNextCounter == 1 && !this.GoneExceptionCreated) + { + this.GoneExceptionCreated = true; + return new CosmosException( + message: $"Epk Range: Partition does not exist at the given range.", + statusCode: System.Net.HttpStatusCode.Gone, + subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone, + activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", + requestCharge: default); + } + else if (this.IsFailedFallbackPipelineTest && this.GoneExceptionCreated && !this.InjectionFailureCreated) + { + this.InjectionFailureCreated = true; + return new CosmosException( + message: "Injected failure", + statusCode: HttpStatusCode.TooManyRequests, + subStatusCode: 3200, + activityId: "111fad5b-d9cb-469f-a165-70867728950e", + requestCharge: 0); + } + else + { + return null; + } + } + } } public sealed class OptimisticDirectExecutionTestOutput : BaselineTestOutput @@ -744,7 +735,7 @@ public override void SerializeAsXml(XmlWriter xmlWriter) } } } - + internal class TestCosmosQueryClient : CosmosQueryClient { public override Action OnExecuteScalarQueryCallback => throw new NotImplementedException(); From c4b0546f8ef8c702b4a1528ea79696a76e8dbbc7 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Fri, 6 Jan 2023 14:35:11 -0800 Subject: [PATCH 27/37] Removed unused imports --- .../Query/OptimisticDirectExecutionQueryBaselineTests.cs | 2 -- 1 file changed, 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index daefbd9499..431c441206 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -13,7 +13,6 @@ using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query; using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline; @@ -29,7 +28,6 @@ using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Routing; using Microsoft.VisualStudio.TestTools.UnitTesting; - using Moq; using Newtonsoft.Json; [TestClass] From 3de111c75cdbd54ec26800a5e2f62fbc8f8615ac Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Fri, 6 Jan 2023 14:36:19 -0800 Subject: [PATCH 28/37] Updated comments --- .../Query/OptimisticDirectExecutionQueryBaselineTests.cs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 431c441206..8b131fe559 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -298,8 +298,7 @@ public async Task TestHandlingOfFailedFallbackPipeline() } } - // it creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing - // TODO: With the addition of the merge/split support, this queryPipelineStage should be able to return all documents regardless of a gone exception happening + // Creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) { List documents = new List(); From 7543ec72182b107cadc9f97527c7d831fb2c9644 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Sun, 8 Jan 2023 18:08:41 -0800 Subject: [PATCH 29/37] Fixed comments and cleaned up test code --- .../Query/Core/Monads/TryCatch{TResult}.cs | 16 ++ ...misticDirectExecutionQueryPipelineStage.cs | 16 +- ...misticDirectExecutionQueryBaselineTests.cs | 159 +++++++----------- 3 files changed, 88 insertions(+), 103 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs index 0fb56995e2..38d0f4d1fc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs @@ -126,6 +126,22 @@ public async Task> TryAsync( return matchResult; } + public async ValueTask> TryAsync( + Func> onSuccess) + { + TryCatch matchResult; + if (this.Succeeded) + { + matchResult = TryCatch.FromResult(await onSuccess(this.either.FromRight(default))); + } + else + { + matchResult = TryCatch.FromException(this.either.FromLeft(default)); + } + + return matchResult; + } + public TryCatch Try( Func> onSuccess) { diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 31624fb04a..f417ece786 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -53,16 +53,17 @@ public ValueTask DisposeAsync() public async ValueTask MoveNextAsync(ITrace trace) { - TryCatch result = this.inner.Try((pipelineStage) => pipelineStage = this.inner.Result) - .Try(pipelineStage => pipelineStage.MoveNextAsync(trace)).Try(moveNextValue => moveNextValue.Result); - bool success = result.Succeeded && result.Result; + TryCatch hasNext = await this.inner.TryAsync(pipelineStage => pipelineStage.MoveNextAsync(trace)); + bool success = hasNext.Succeeded && hasNext.Result; if (success) { this.continuationToken = this.Current.Succeeded ? this.Current.Result.State?.Value : null; } - else if (this.executionState == ExecutionState.OptimisticDirectExecution && result.Succeeded) + else if (this.executionState == ExecutionState.OptimisticDirectExecution && hasNext.Succeeded) { + // only the case where hasNext.Succeeded = true and hasNext.Result = false should enter this code block + if (this.Current.InnerMostException.IsPartitionSplitException()) { this.inner = await this.queryPipelineStageFactory(this.TryUnwrapContinuationToken()); @@ -81,15 +82,14 @@ public async ValueTask MoveNextAsync(ITrace trace) public void SetCancellationToken(CancellationToken cancellationToken) { - this.inner.Try((pipelineStage) => pipelineStage = this.inner.Result) - .Try(pipelineStage => pipelineStage.SetCancellationToken(cancellationToken)); + this.inner.Try(pipelineStage => pipelineStage.SetCancellationToken(cancellationToken)); } private CosmosElement TryUnwrapContinuationToken() { CosmosObject cosmosObject = this.continuationToken as CosmosObject; - CosmosElement specializedContinuationToken = cosmosObject[optimisticDirectExecutionToken]; - return CosmosArray.Create(specializedContinuationToken); + CosmosElement backendContinuationToken = cosmosObject[optimisticDirectExecutionToken]; + return CosmosArray.Create(backendContinuationToken); } public static TryCatch MonadicCreate( diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 8b131fe559..9b3a577ec2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -222,6 +222,15 @@ public async Task TestPipelineForGoneExceptionOnSingleAndMultiplePartitionAsync( Assert.IsTrue(await ExecuteGoneExceptionOnODEPipeline(isMultiPartition: true)); } + // test to check if failing fallback pipeline is handled properly + [TestMethod] + public async Task TestHandlingOfFailedFallbackPipelineOnSingleAndMultiplePartitionAsync() + { + Assert.IsTrue(await TestHandlingOfFailedFallbackPipeline(isMultiPartition: false)); + + Assert.IsTrue(await TestHandlingOfFailedFallbackPipeline(isMultiPartition: true)); + } + // The reason we have the below test is to show the missing capabilities of the OptimisticDirectExecution pipeline. // Currently this pipeline cannot handle distributed queries as it does not have the logic to sum up the values it gets from the backend in partial results. // This functionality is available for other pipelines such as the ParallelCrossPartitionQueryPipelineStage. @@ -246,69 +255,79 @@ public async Task TestPipelineForDistributedQueryAsync() Assert.AreEqual(1, result); } - [TestMethod] - public async Task TestHandlingOfFailedFallbackPipeline() + // Creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing + private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) { + int numItems = 100; List documents = new List(); - MergeTestUtil mergeTest = new MergeTestUtil + QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); + (MergeTestUtil mergeTest, IQueryPipelineStage queryPipelineStage) = await CreateFallbackPipelineTestInfra(numItems, isFailedFallbackPipelineTest: false, isMultiPartition, queryRequestOptions); + + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { - MoveNextCounter = 0, - IsFailedFallbackPipelineTest = true - }; + if (mergeTest.MoveNextCounter == 1) + { + Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); + } + else + { + Assert.AreNotEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); + } - int numItems = 100; - OptimisticDirectExecutionTestInput input = CreateInput( - description: @"Single Partition Key and Value Field", - query: "SELECT * FROM c", - expectedOptimisticDirectExecution: true, - partitionKeyPath: @"/pk", - partitionKeyValue: "a"); + TryCatch tryGetPage = queryPipelineStage.Current; - QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); - DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( - numItems, - multiPartition: false, - failureConfigs: new FlakyDocumentContainer.FailureConfigs( - inject429s: false, - injectEmptyPages: false, - shouldReturnFailure: mergeTest.ShouldReturnFailure)); + if (tryGetPage.Failed) + { + // failure should never come till here. Should be handled before + Assert.Fail(); + } - IQueryPipelineStage queryPipelineStage = await GetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); + documents.AddRange(tryGetPage.Result.Documents); + } + + Assert.AreEqual(numItems, documents.Count); + return true; + } + + private static async Task TestHandlingOfFailedFallbackPipeline(bool isMultiPartition) + { + int numItems = 100; + List documents = new List(); + QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); + (MergeTestUtil mergeTest, IQueryPipelineStage queryPipelineStage) = await CreateFallbackPipelineTestInfra(numItems, isFailedFallbackPipelineTest: true, isMultiPartition, queryRequestOptions); while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { - mergeTest.MoveNextCounter++; TryCatch tryGetPage = queryPipelineStage.Current; - if (tryGetPage.Failed) { - if (mergeTest.MoveNextCounter == 2) + if (mergeTest.MoveNextCounter == 3) { Assert.IsTrue(tryGetPage.InnerMostException.Message.Equals("Injected failure")); Assert.AreNotEqual(numItems, documents.Count); - return; + return true; } else { Assert.Fail(); + return false; } } documents.AddRange(tryGetPage.Result.Documents); } + + return false; } - // Creates a gone exception after the first MoveNexyAsync() call. This allows for the pipeline to return some documents before failing - private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPartition) + private static async Task<(MergeTestUtil, IQueryPipelineStage)> CreateFallbackPipelineTestInfra(int numItems, bool isFailedFallbackPipelineTest, bool isMultiPartition, QueryRequestOptions queryRequestOptions) { List documents = new List(); MergeTestUtil mergeTest = new MergeTestUtil { - MoveNextCounter = 0, - IsFailedFallbackPipelineTest = false + IsFailedFallbackPipelineTest = isFailedFallbackPipelineTest }; - int numItems = 100; OptimisticDirectExecutionTestInput input = CreateInput( description: @"Single Partition Key and Value Field", query: "SELECT * FROM c", @@ -316,7 +335,6 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa partitionKeyPath: @"/pk", partitionKeyValue: "a"); - QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); DocumentContainer inMemoryCollection = await CreateDocumentContainerAsync( numItems, multiPartition: isMultiPartition, @@ -327,31 +345,7 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa IQueryPipelineStage queryPipelineStage = await GetOdePipelineAsync(input, inMemoryCollection, queryRequestOptions); - while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) - { - mergeTest.MoveNextCounter++; - if (mergeTest.MoveNextCounter == 1) - { - Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); - } - else - { - Assert.AreNotEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); - } - - TryCatch tryGetPage = queryPipelineStage.Current; - - if (tryGetPage.Failed) - { - // failure should never come till here. Should be handled before - Assert.Fail(); - } - - documents.AddRange(tryGetPage.Result.Documents); - } - - Assert.AreEqual(100, documents.Count); - return true; + return (mergeTest, queryPipelineStage); } private async Task GetPipelineAndDrainAsync(OptimisticDirectExecutionTestInput input, int numItems, bool isMultiPartition, int expectedContinuationTokenCount) @@ -396,7 +390,7 @@ private async Task GetPipelineAndDrainAsync(OptimisticDirectExecutionTestIn return documents.Count; } - private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(string querySpecJsonString, PartitionKeyDefinition pkDefinition) + internal static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(string querySpecJsonString, PartitionKeyDefinition pkDefinition) { TryCatch tryGetQueryPlan = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfo( querySpecJsonString: querySpecJsonString, @@ -608,18 +602,18 @@ private static QueryRequestOptions GetQueryRequestOptions(bool enableOptimisticD private class MergeTestUtil { - public int MoveNextCounter { get; set; } + public int MoveNextCounter { get; private set; } - public bool GoneExceptionCreated { get; set; } + public bool GoneExceptionCreated { get; private set; } - public bool InjectionFailureCreated { get; set; } + public bool TooManyRequestsFailureCreated { get; private set; } public bool IsFailedFallbackPipelineTest { get; set; } public async Task ShouldReturnFailure() { - - if (this.MoveNextCounter == 1 && !this.GoneExceptionCreated) + this.MoveNextCounter++; + if (this.MoveNextCounter == 2 && !this.GoneExceptionCreated) { this.GoneExceptionCreated = true; return new CosmosException( @@ -629,9 +623,9 @@ public async Task ShouldReturnFailure() activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", requestCharge: default); } - else if (this.IsFailedFallbackPipelineTest && this.GoneExceptionCreated && !this.InjectionFailureCreated) + else if (this.IsFailedFallbackPipelineTest && this.GoneExceptionCreated && !this.TooManyRequestsFailureCreated) { - this.InjectionFailureCreated = true; + this.TooManyRequestsFailureCreated = true; return new CosmosException( message: "Injected failure", statusCode: HttpStatusCode.TooManyRequests, @@ -802,37 +796,12 @@ public override Task> TryGetOverlappingRangesAs } public override async Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, CancellationToken cancellationToken) - { - IReadOnlyList orderByColumn = new List() - { - new SortOrder() - }; - - IReadOnlyList orderByExpression = new List() - { - "ORDER BY" - }; - - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = new PartitionedQueryExecutionInfo() - { - QueryInfo = new QueryInfo() - { - Aggregates = null, - DistinctType = DistinctQueryType.None, - GroupByAliases = null, - GroupByAliasToAggregateType = null, - GroupByExpressions = null, - HasSelectValue = false, - Limit = null, - Offset = null, - OrderBy = sqlQuerySpec.QueryText.Contains("ORDER BY") ? orderByColumn : null, - OrderByExpressions = sqlQuerySpec.QueryText.Contains("ORDER BY") ? orderByExpression : null, - RewrittenQuery = null, - Top = null, - }, - QueryRanges = new List>(), - }; + { + CosmosSerializerCore serializerCore = new(); + using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, Documents.ResourceType.Document)); + string sqlQuerySpecJsonString = streamReader.ReadToEnd(); + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = OptimisticDirectExecutionQueryBaselineTests.GetPartitionedQueryExecutionInfo(sqlQuerySpecJsonString, partitionKeyDefinition); return TryCatch.FromResult(partitionedQueryExecutionInfo); } } From 134dcc82de3ca4272378760742bfcac977e9ea4a Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Sun, 8 Jan 2023 23:34:43 -0800 Subject: [PATCH 30/37] Added CosmosElement null check in TryUnwrapContinuationToken() --- .../OptimisticDirectExecutionQueryPipelineStage.cs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index f417ece786..e30d936e30 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -88,8 +88,13 @@ public void SetCancellationToken(CancellationToken cancellationToken) private CosmosElement TryUnwrapContinuationToken() { CosmosObject cosmosObject = this.continuationToken as CosmosObject; - CosmosElement backendContinuationToken = cosmosObject[optimisticDirectExecutionToken]; - return CosmosArray.Create(backendContinuationToken); + if (cosmosObject != null) + { + CosmosElement backendContinuationToken = cosmosObject[optimisticDirectExecutionToken]; + return CosmosArray.Create(backendContinuationToken); + } + + return null; } public static TryCatch MonadicCreate( From 9d37c12ba0d1135c4d2461e367093725f0f83ae0 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Mon, 9 Jan 2023 10:46:15 -0800 Subject: [PATCH 31/37] Removed FlakyDocumentContainer.cs from pull request --- .../Pagination/FlakyDocumentContainer.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs index 3ea96d7634..cc038ff4ed 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/FlakyDocumentContainer.cs @@ -190,8 +190,8 @@ public async Task> MonadicQueryAsync( } public async Task> MonadicChangeFeedAsync( - FeedRangeState feedRangeState, - ChangeFeedPaginationOptions changeFeedPaginationOptions, + FeedRangeState feedRangeState, + ChangeFeedPaginationOptions changeFeedPaginationOptions, ITrace trace, CancellationToken cancellationToken) { @@ -289,7 +289,7 @@ public sealed class FailureConfigs public delegate Task ShouldReturnFailure(); public FailureConfigs( - bool inject429s, + bool inject429s, bool injectEmptyPages, Exception throwException = null, Exception returnFailure = null) From a009e539556a0a21509e38c34aeccf0519fe01c0 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Mon, 9 Jan 2023 11:37:02 -0800 Subject: [PATCH 32/37] Removed unused imports --- .../Core/Pipeline/CosmosQueryExecutionContextFactory.cs | 5 ++--- .../OptimisticDirectExecutionQueryPipelineStage.cs | 2 +- .../Query/OptimisticDirectExecutionQueryBaselineTests.cs | 7 +++---- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 93e95846c9..ec1db1f7f0 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos; - using Microsoft.Azure.Cosmos.Core.Collections; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core; @@ -148,7 +147,7 @@ private static async Task> TryCreateCoreContextAsy { return await TryCreateExecutionContextAsync( documentContainer, - null, + partitionedQueryExecutionInfo: null, cosmosQueryContext, containerQueryProperties, inputParameters, @@ -349,7 +348,7 @@ private static async Task> TryCreateExecutionConte // 1. the token is actually malformed // 2. Its a non Ode continuation token // In both cases, Ode pipeline delegates the work to the Specialized pipeline - // as Ode ppipeline should not take over execution while some other pipeline is already handling it + // as Ode pipeline should not take over execution while some other pipeline is already handling it if (tryCreatePipelineStage.Failed && tryCreatePipelineStage.InnerMostException is MalformedContinuationTokenException) { SetTestInjectionPipelineType(inputParameters, Specialized); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index e30d936e30..20756fe061 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -62,7 +62,7 @@ public async ValueTask MoveNextAsync(ITrace trace) } else if (this.executionState == ExecutionState.OptimisticDirectExecution && hasNext.Succeeded) { - // only the case where hasNext.Succeeded = true and hasNext.Result = false should enter this code block + // only if hasNext.Succeeded = true and hasNext.Result = false should this code block be hit if (this.Current.InnerMostException.IsPartitionSplitException()) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 9b3a577ec2..fde851d82a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -18,7 +18,6 @@ using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; @@ -261,7 +260,7 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa int numItems = 100; List documents = new List(); QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); - (MergeTestUtil mergeTest, IQueryPipelineStage queryPipelineStage) = await CreateFallbackPipelineTestInfra(numItems, isFailedFallbackPipelineTest: false, isMultiPartition, queryRequestOptions); + (MergeTestUtil mergeTest, IQueryPipelineStage queryPipelineStage) = await CreateFallbackPipelineTestInfrastructure(numItems, isFailedFallbackPipelineTest: false, isMultiPartition, queryRequestOptions); while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { @@ -294,7 +293,7 @@ private static async Task TestHandlingOfFailedFallbackPipeline(bool isMult int numItems = 100; List documents = new List(); QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: true); - (MergeTestUtil mergeTest, IQueryPipelineStage queryPipelineStage) = await CreateFallbackPipelineTestInfra(numItems, isFailedFallbackPipelineTest: true, isMultiPartition, queryRequestOptions); + (MergeTestUtil mergeTest, IQueryPipelineStage queryPipelineStage) = await CreateFallbackPipelineTestInfrastructure(numItems, isFailedFallbackPipelineTest: true, isMultiPartition, queryRequestOptions); while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) { @@ -320,7 +319,7 @@ private static async Task TestHandlingOfFailedFallbackPipeline(bool isMult return false; } - private static async Task<(MergeTestUtil, IQueryPipelineStage)> CreateFallbackPipelineTestInfra(int numItems, bool isFailedFallbackPipelineTest, bool isMultiPartition, QueryRequestOptions queryRequestOptions) + private static async Task<(MergeTestUtil, IQueryPipelineStage)> CreateFallbackPipelineTestInfrastructure(int numItems, bool isFailedFallbackPipelineTest, bool isMultiPartition, QueryRequestOptions queryRequestOptions) { List documents = new List(); MergeTestUtil mergeTest = new MergeTestUtil From 400d4f58894526aa67c8ce0981b83d1c70873522 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Mon, 9 Jan 2023 15:27:25 -0800 Subject: [PATCH 33/37] Updated TryUnwrapContinuationToken() --- .../Core/Pipeline/CosmosQueryExecutionContextFactory.cs | 1 - .../OptimisticDirectExecutionQueryPipelineStage.cs | 8 ++++---- .../Query/OptimisticDirectExecutionQueryBaselineTests.cs | 4 ++-- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index ec1db1f7f0..82cde7df15 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -378,7 +378,6 @@ private static async Task> TryCreateExecutionConte trace, cancellationToken); } - } return tryCreatePipelineStage; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 20756fe061..a10f35fd29 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -62,8 +62,7 @@ public async ValueTask MoveNextAsync(ITrace trace) } else if (this.executionState == ExecutionState.OptimisticDirectExecution && hasNext.Succeeded) { - // only if hasNext.Succeeded = true and hasNext.Result = false should this code block be hit - + Debug.Assert(hasNext.Result == false); if (this.Current.InnerMostException.IsPartitionSplitException()) { this.inner = await this.queryPipelineStageFactory(this.TryUnwrapContinuationToken()); @@ -87,10 +86,11 @@ public void SetCancellationToken(CancellationToken cancellationToken) private CosmosElement TryUnwrapContinuationToken() { - CosmosObject cosmosObject = this.continuationToken as CosmosObject; - if (cosmosObject != null) + if (this.continuationToken != null) { + CosmosObject cosmosObject = this.continuationToken as CosmosObject; CosmosElement backendContinuationToken = cosmosObject[optimisticDirectExecutionToken]; + Debug.Assert(backendContinuationToken != null); return CosmosArray.Create(backendContinuationToken); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index fde851d82a..cc342fa6ce 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -278,7 +278,7 @@ private static async Task ExecuteGoneExceptionOnODEPipeline(bool isMultiPa if (tryGetPage.Failed) { // failure should never come till here. Should be handled before - Assert.Fail(); + Assert.Fail("Unexpected error. Gone Exception should not reach till here"); } documents.AddRange(tryGetPage.Result.Documents); @@ -308,7 +308,7 @@ private static async Task TestHandlingOfFailedFallbackPipeline(bool isMult } else { - Assert.Fail(); + Assert.Fail("Fallback pipeline failure not handled correctly"); return false; } } From 617911202be70bd7eb81ea5ea801314f4cb3a148 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Tue, 10 Jan 2023 10:50:27 -0800 Subject: [PATCH 34/37] Update MoveNextAsync() call in OptimisticDirectExecutionQueryBaselineTests.cs --- .../Query/OptimisticDirectExecutionQueryBaselineTests.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index cc342fa6ce..d1a9c3191c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -524,8 +524,8 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect inputParameters, NoOpTrace.Singleton); - bool result = queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton).Result; - + bool result = queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton).AsTask().GetAwaiter().GetResult(); + if (input.ExpectedOptimisticDirectExecution) { Assert.AreEqual(TestInjections.PipelineType.OptimisticDirectExecution, queryRequestOptions.TestSettings.Stats.PipelineType.Value); From c94db624a488f28caec111570db87fdfc6f9205e Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 11 Jan 2023 13:11:59 -0800 Subject: [PATCH 35/37] Made MergeTestUtil.IsFailedFallbackPipelineTest a readonly property --- ...timisticDirectExecutionQueryBaselineTests.cs | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index d1a9c3191c..c0644045e5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -322,10 +322,7 @@ private static async Task TestHandlingOfFailedFallbackPipeline(bool isMult private static async Task<(MergeTestUtil, IQueryPipelineStage)> CreateFallbackPipelineTestInfrastructure(int numItems, bool isFailedFallbackPipelineTest, bool isMultiPartition, QueryRequestOptions queryRequestOptions) { List documents = new List(); - MergeTestUtil mergeTest = new MergeTestUtil - { - IsFailedFallbackPipelineTest = isFailedFallbackPipelineTest - }; + MergeTestUtil mergeTest = new MergeTestUtil(isFailedFallbackPipelineTest); OptimisticDirectExecutionTestInput input = CreateInput( description: @"Single Partition Key and Value Field", @@ -565,7 +562,7 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect partitionedQueryExecutionInfo: partitionedQueryExecutionInfo, executionEnvironment: null, returnResultsInDeterministicOrder: null, - forcePassthrough: true, + forcePassthrough: false, testInjections: queryRequestOptions.TestSettings); string databaseId = "db1234"; @@ -590,11 +587,10 @@ private static QueryRequestOptions GetQueryRequestOptions(bool enableOptimisticD { MaxConcurrency = 0, MaxItemCount = 10, - MaxBufferedItemCount = 7000, TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: enableOptimisticDirectExecution, new TestInjections.ResponseStats()), Properties = new Dictionary() { - { HttpConstants.HttpHeaders.EnumerationDirection, "test" }, + { HttpConstants.HttpHeaders.EnumerationDirection, ""}, } }; } @@ -607,7 +603,12 @@ private class MergeTestUtil public bool TooManyRequestsFailureCreated { get; private set; } - public bool IsFailedFallbackPipelineTest { get; set; } + public bool IsFailedFallbackPipelineTest { get; } + + public MergeTestUtil(bool isFailedFallbackPipelineTest) + { + this.IsFailedFallbackPipelineTest = isFailedFallbackPipelineTest; + } public async Task ShouldReturnFailure() { From 2416a89cea8840244331622b07a59ad3df28e260 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Thu, 12 Jan 2023 16:07:21 -0800 Subject: [PATCH 36/37] Added IsPartitionSplitException() overload to take CosmosElement --- .../Pipeline/CosmosExceptionExtensions.cs | 14 +++++++++-- .../CosmosQueryExecutionContextFactory.cs | 25 +++++++++++-------- ...misticDirectExecutionQueryBaselineTests.cs | 9 +++---- 3 files changed, 30 insertions(+), 18 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs index 98e9023226..dca1ef488d 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs @@ -9,10 +9,20 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline internal static class CosmosExceptionExtensions { public static bool IsPartitionSplitException(this Exception ex) + { + if (ex != null) + { + return IsPartitionSplitException(ex as CosmosException); + } + + return false; + } + + public static bool IsPartitionSplitException(this CosmosException ex) { return ex is CosmosException cosmosException - && (cosmosException.StatusCode == System.Net.HttpStatusCode.Gone) - && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); + && (cosmosException.StatusCode == System.Net.HttpStatusCode.Gone) + && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone); } } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 82cde7df15..2476d9d701 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -305,19 +305,22 @@ private static async Task> TryCreateFromPartitione trace, cancellationToken); } - else if (createPassthroughQuery) - { - SetTestInjectionPipelineType(inputParameters, Passthrough); - - tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext( - documentContainer, - inputParameters, - targetRanges, - cancellationToken); - } else { - tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken); + if (createPassthroughQuery) + { + SetTestInjectionPipelineType(inputParameters, Passthrough); + + tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext( + documentContainer, + inputParameters, + targetRanges, + cancellationToken); + } + else + { + tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken); + } } return tryCreatePipelineStage; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index c0644045e5..6737a3d338 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -623,7 +623,8 @@ public async Task ShouldReturnFailure() activityId: "0f8fad5b-d9cb-469f-a165-70867728950e", requestCharge: default); } - else if (this.IsFailedFallbackPipelineTest && this.GoneExceptionCreated && !this.TooManyRequestsFailureCreated) + + if (this.IsFailedFallbackPipelineTest && this.GoneExceptionCreated && !this.TooManyRequestsFailureCreated) { this.TooManyRequestsFailureCreated = true; return new CosmosException( @@ -633,10 +634,8 @@ public async Task ShouldReturnFailure() activityId: "111fad5b-d9cb-469f-a165-70867728950e", requestCharge: 0); } - else - { - return null; - } + + return null; } } } From c129afe518ebaebc1a9da724b25b366bc4a6fee2 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 18 Jan 2023 13:04:42 -0800 Subject: [PATCH 37/37] Fixed bug regarding syntax error queries --- ...misticDirectExecutionQueryPipelineStage.cs | 25 ++++++++----------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index a10f35fd29..0664603024 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -55,25 +55,22 @@ public async ValueTask MoveNextAsync(ITrace trace) { TryCatch hasNext = await this.inner.TryAsync(pipelineStage => pipelineStage.MoveNextAsync(trace)); bool success = hasNext.Succeeded && hasNext.Result; + bool isPartitionSplitException = hasNext.Succeeded && this.Current.Failed && this.Current.InnerMostException.IsPartitionSplitException(); - if (success) + if (success && !isPartitionSplitException) { this.continuationToken = this.Current.Succeeded ? this.Current.Result.State?.Value : null; } - else if (this.executionState == ExecutionState.OptimisticDirectExecution && hasNext.Succeeded) + else if (isPartitionSplitException && this.executionState == ExecutionState.OptimisticDirectExecution) { - Debug.Assert(hasNext.Result == false); - if (this.Current.InnerMostException.IsPartitionSplitException()) + this.inner = await this.queryPipelineStageFactory(this.TryUnwrapContinuationToken()); + this.executionState = ExecutionState.SpecializedDocumentQueryExecution; + if (this.inner.Failed) { - this.inner = await this.queryPipelineStageFactory(this.TryUnwrapContinuationToken()); - this.executionState = ExecutionState.SpecializedDocumentQueryExecution; - if (this.inner.Failed) - { - return false; - } - - success = await this.inner.Result.MoveNextAsync(trace); + return false; } + + success = await this.inner.Result.MoveNextAsync(trace); } return success; @@ -161,8 +158,8 @@ public async ValueTask MoveNextAsync(ITrace trace) TryCatch partitionPage = this.queryPartitionRangePageAsyncEnumerator.Current; if (partitionPage.Failed) { - this.Current = partitionPage; - return false; + this.Current = TryCatch.FromException(partitionPage.Exception); + return true; } QueryPage backendQueryPage = partitionPage.Result;