[Internal] Query: Adds Split Support for Ode (#3572)

* Added tests to test different aspects of merge/split support with OptimisticDirectExecution pipeline. Tests check for gone exception handling, pipeline switching etc.

* Added gone exception simulation tests.

* Added new tests and improved test infra

* Removed ParalleContEvocation test. Fixed comments

* Removed CreateParallelCrossPartitionPipelineStateAsync() as it is not being used anymore

* Removed while loop in CreateDocumentContainerAsync()

* Fixed comments.

* Updated ExecuteGoneExceptionOnODEPipeline()

* Added type Assert for ExecuteGoneExceptionOnODEPipeline()

* Replaced try-catch with if statement in MoveNextAsync()

* Added delegate to access TryCreateCoreContextAsync()

* Added check to confirm Ode pipeline is not called in fallback plan

* Updated method name from OptimisticDirectExecutionContext() to TryCreateOptimisticDirectExecutionContext()

* Using delegate instead of Func<>.

* Ode fallback plan always calls Specialized pipeline

* Using ServiceInterop/Gateway to get QueryPlan for Specialized Pipeline

* Added new test to check handling of failing fallback pipeline

* Code cleanup

* Added logic for handling non ODE continuation tokens

* Moved delegate away from member variables

* Added tests for Merge case

* Updated method names

* Added checks for tryCatch

* Updated SetCancellationToken() to use Try

* Updated TryUnwrapContinuationToken()

* Removed changes in FlakyDocumentContainer.cs

* Removed unused imports

* Updated comments

* Fixed comments and cleaned up test code

* Added CosmosElement null check in TryUnwrapContinuationToken()

* Removed FlakyDocumentContainer.cs from pull request

* Removed unused imports

* Updated TryUnwrapContinuationToken()

* Update MoveNextAsync() call in OptimisticDirectExecutionQueryBaselineTests.cs

* Made MergeTestUtil.IsFailedFallbackPipelineTest a readonly property

* Added IsPartitionSplitException() overload to take CosmosElement

* Fixed bug regarding syntax error queries

Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs

@@ -126,6 +126,38 @@ public async Task<TryCatch<T>> TryAsync<T>(
  return matchResult;
  }
 
+ public async ValueTask<TryCatch<T>> TryAsync<T>(
+ Func<TResult, ValueTask<T>> onSuccess)
+ {
+ TryCatch<T> matchResult;
+ if (this.Succeeded)
+ {
+ matchResult = TryCatch<T>.FromResult(await onSuccess(this.either.FromRight(default)));
+ }
+ else
+ {
+ matchResult = TryCatch<T>.FromException(this.either.FromLeft(default));
+ }
+
+ return matchResult;
+ }
+
+ public TryCatch<T> Try<T>(
+ Func<TResult, TryCatch<T>> onSuccess)
+ {
+ TryCatch<T> matchResult;
+ if (this.Succeeded)
+ {
+ matchResult = onSuccess(this.either.FromRight(default));
+ }
+ else
+ {
+ matchResult = TryCatch<T>.FromException(this.either.FromLeft(default));
+ }
+
+ return matchResult;
+ }
+
  public TryCatch<TResult> Catch(
  Action<Exception> onError)
  {

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs

@@ -0,0 +1,28 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline
+{
+ using System;
+
+ 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);
+ }
+ }
+}

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 exception.IsPartitionSplitException();
  }
 
  public void SetCancellationToken(CancellationToken cancellationToken)

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs

@@ -43,11 +43,11 @@ public static CosmosElement ToCosmosElement(OptimisticDirectExecutionContinuatio
  public static TryCatch<OptimisticDirectExecutionContinuationToken> TryCreateFromCosmosElement(CosmosElement cosmosElement)
  {
  CosmosObject cosmosObjectContinuationToken = cosmosElement as CosmosObject;
- if (cosmosObjectContinuationToken == null)
+ if (cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken))
  {
  return TryCatch<OptimisticDirectExecutionContinuationToken>.FromException(
- new MalformedChangeFeedContinuationTokenException(
- message: $"Malformed Continuation Token"));
+ new MalformedContinuationTokenException(
+ message: $"Malformed Continuation Token: Expected OptimisticDirectExecutionToken\r\n"));
  }
 
  TryCatch<ParallelContinuationToken> inner = ParallelContinuationToken.TryCreateFromCosmosElement(cosmosObjectContinuationToken[OptimisticDirectExecutionToken]);