[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]);