[Internal] Query: Adds OptimisticDirectExecute and RequiresDistribution headers

Microsoft.Azure.Cosmos/src/Headers/CosmosMessageHeadersInternal.cs

@@ -177,6 +177,18 @@ public virtual string EndEpk
             set => this.SetProperty(HttpConstants.HttpHeaders.EndEpk, value);
         }
 
+        public virtual string OptimisticDirectExecute
+        {
+            get => this.GetValueOrDefault(HttpConstants.HttpHeaders.OptimisticDirectExecute);
+            set => this.SetProperty(HttpConstants.HttpHeaders.OptimisticDirectExecute, value);
+        }
+
+        public virtual string RequiresDistribution
+        {
+            get => this.GetValueOrDefault(HttpConstants.HttpHeaders.RequiresDistribution);
+            set => this.SetProperty(HttpConstants.HttpHeaders.RequiresDistribution, value);
+        }
+
         public abstract INameValueCollection INameValueCollection { get; }
 
         public virtual string this[string headerName] 

Microsoft.Azure.Cosmos/src/Headers/Headers.cs

@@ -171,6 +171,18 @@ internal virtual string PartitionKey
             set => this.CosmosMessageHeaders.PartitionKey = value;
         }
 
+        internal virtual string OptimisticDirectExecute
+        {
+            get => this.CosmosMessageHeaders.OptimisticDirectExecute;
+            set => this.CosmosMessageHeaders.OptimisticDirectExecute = value;
+        }
+
+        internal virtual string RequiresDistribution
+        {
+            get => this.CosmosMessageHeaders.RequiresDistribution;
+            set => this.CosmosMessageHeaders.RequiresDistribution = value;
+        }
+
         internal virtual string PartitionKeyRangeId
         {
             get => this.CosmosMessageHeaders.PartitionKeyRangeId;

Microsoft.Azure.Cosmos/src/Headers/StoreRequestHeaders.cs

@@ -46,6 +46,12 @@ public override string XDate
             set => this.requestNameValueCollection.XDate = value;
         }
 
+        public override string OptimisticDirectExecute
+        {
+            get => this.requestNameValueCollection.OptimisticDirectExecute;
+            set => this.requestNameValueCollection.OptimisticDirectExecute = value;
+        }
+
         public override string ConsistencyLevel
         {
             get => this.requestNameValueCollection.ConsistencyLevel;

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

@@ -258,16 +258,17 @@ public async Task<TryCatch<QueryPage>> MonadicQueryAsync(
             }
 
             QueryRequestOptions queryRequestOptions = this.queryRequestOptions == null ? new QueryRequestOptions() : this.queryRequestOptions;
+            AdditionalRequestHeaders additionalRequestHeaders = new AdditionalRequestHeaders(this.correlatedActivityId, isContinuationExpected: false, optimisticDirectExecute: queryPaginationOptions.OptimisticDirectExecute);
+
             TryCatch<QueryPage> monadicQueryPage = await this.cosmosQueryClient.ExecuteItemQueryAsync(
                 this.resourceLink,
                 this.resourceType,
                 Documents.OperationType.Query,
-                this.correlatedActivityId,
                 feedRangeState.FeedRange,
                 queryRequestOptions,
+                additionalRequestHeaders,
                 sqlQuerySpec,
                 feedRangeState.State == null ? null : ((CosmosString)feedRangeState.State.Value).Value,
-                isContinuationExpected: false,
                 queryPaginationOptions.PageSizeLimit ?? int.MaxValue,
                 trace,
                 cancellationToken);

Microsoft.Azure.Cosmos/src/Query/Core/AdditionalRequestHeaders.cs

@@ -0,0 +1,36 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Query.Core
+{
+    using System;
+
+    /// <summary>
+    /// Captures additional headers to be sent in the web request.
+    /// </summary>
+    internal sealed class AdditionalRequestHeaders
+    {
+        /// <summary>
+        /// Enable backend to link multiple activityIds that belong to the same operation.
+        /// </summary>
+        public Guid CorrelatedActivityId { get; }
+
+        /// <summary>
+        /// Let the backend know if continuations are expected for a certain query.
+        /// </summary>
+        public bool IsContinuationExpected { get; }
+
+        /// <summary>
+        /// Let the backend know if a query is utilizing the OptimisticDirectExecute pipeline.
+        /// </summary>
+        public bool OptimisticDirectExecute { get; }
+
+        public AdditionalRequestHeaders(Guid correlatedActivityId, bool isContinuationExpected, bool optimisticDirectExecute)
+        {
+            this.CorrelatedActivityId = correlatedActivityId;
+            this.IsContinuationExpected = isContinuationExpected;
+            this.OptimisticDirectExecute = optimisticDirectExecute;
+        }
+    }
+}

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

@@ -477,7 +477,6 @@ private static TryCatch<IQueryPipelineStage> TryCreateOptimisticDirectExecutionC
                 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

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

@@ -19,6 +19,7 @@ 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
     {
@@ -28,18 +29,25 @@ private enum ExecutionState
             SpecializedDocumentQueryExecution,
         }
 
-        private const string optimisticDirectExecutionToken = "OptimisticDirectExecutionToken";
+        private const string OptimisticDirectExecutionToken = "OptimisticDirectExecutionToken";
         private readonly FallbackQueryPipelineStageFactory queryPipelineStageFactory;
         private TryCatch<IQueryPipelineStage> inner;
         private CosmosElement continuationToken;
         private ExecutionState executionState;
+        private bool? previousRequiresDistribution;
 
         private OptimisticDirectExecutionQueryPipelineStage(TryCatch<IQueryPipelineStage> inner, FallbackQueryPipelineStageFactory queryPipelineStageFactory, CosmosElement continuationToken)
         {
             this.inner = inner;
             this.queryPipelineStageFactory = queryPipelineStageFactory;
             this.continuationToken = continuationToken;
             this.executionState = ExecutionState.OptimisticDirectExecution;
+
+            if (this.continuationToken != null)
+            {
+                this.previousRequiresDistribution = false;
+            }
+
         }
 
         public delegate Task<TryCatch<IQueryPipelineStage>> FallbackQueryPipelineStageFactory(CosmosElement continuationToken);
@@ -55,22 +63,43 @@ public async ValueTask<bool> MoveNextAsync(ITrace trace)
         {
             TryCatch<bool> 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 && !isPartitionSplitException)
-            {
-                this.continuationToken = this.Current.Succeeded ? this.Current.Result.State?.Value : null;
-            }
-            else if (isPartitionSplitException && this.executionState == ExecutionState.OptimisticDirectExecution)
+            if (this.executionState == ExecutionState.OptimisticDirectExecution)
             {
-                this.inner = await this.queryPipelineStageFactory(this.TryUnwrapContinuationToken());
-                this.executionState = ExecutionState.SpecializedDocumentQueryExecution;
-                if (this.inner.Failed)
+                bool isPartitionSplitException = hasNext.Succeeded && this.Current.Failed && this.Current.InnerMostException.IsPartitionSplitException();
+                if (success && !isPartitionSplitException)
                 {
-                    return false;
+                    this.continuationToken = this.Current.Succeeded ? this.Current.Result.State?.Value : null;
+                    if (this.continuationToken != null)
+                    {
+                        bool requiresDistribution;
+                        if (this.Current.Result.AdditionalHeaders.TryGetValue(HttpConstants.HttpHeaders.RequiresDistribution, out string requiresDistributionHeaderValue))
+                        {
+                            requiresDistribution = bool.Parse(requiresDistributionHeaderValue);
+                        }
+                        else
+                        {
+                            requiresDistribution = true;
+                        }
+
+                        if (this.previousRequiresDistribution.HasValue && this.previousRequiresDistribution != requiresDistribution)
+                        {
+                            // We should never enter this if statement as requiresDistribution flag can never switch mid execution.
+                            // Hence, this exception should never be thrown.
+                            throw new InvalidOperationException($"Unexpected switch in {HttpConstants.HttpHeaders.RequiresDistribution} value. Previous value : {this.previousRequiresDistribution} Current value : {requiresDistribution}.");
+                        }
+
+                        if (requiresDistribution)
+                        {
+                            success = await this.SwitchToFallbackPipelineAsync(continuationToken: null, trace);
+                        }
+
+                        this.previousRequiresDistribution = requiresDistribution;
+                    }
+                }
+                else if (isPartitionSplitException)
+                {
+                    success = await this.SwitchToFallbackPipelineAsync(continuationToken: this.continuationToken, trace);
                 }
-
-                success = await this.inner.Result.MoveNextAsync(trace);
             }
 
             return success;
@@ -81,32 +110,46 @@ public void SetCancellationToken(CancellationToken cancellationToken)
             this.inner.Try(pipelineStage => pipelineStage.SetCancellationToken(cancellationToken));
         }
 
-        private CosmosElement TryUnwrapContinuationToken()
+        private static CosmosElement TryUnwrapContinuationToken(CosmosElement continuationToken)
         {
-            if (this.continuationToken != null)
+            if (continuationToken == null) return null;
+
+            CosmosObject cosmosObject = continuationToken as CosmosObject;
+            CosmosElement backendContinuationToken = cosmosObject[OptimisticDirectExecutionToken];
+            Debug.Assert(backendContinuationToken != null);
+
+            return CosmosArray.Create(backendContinuationToken);
+        }
+
+        private async Task<bool> SwitchToFallbackPipelineAsync(CosmosElement continuationToken, ITrace trace)
+        {
+            Debug.Assert(this.executionState == ExecutionState.OptimisticDirectExecution, "OptimisticDirectExecuteQueryPipelineStage Assert!", "Only OptimisticDirectExecute pipeline can create this fallback pipeline");
+            this.executionState = ExecutionState.SpecializedDocumentQueryExecution;
+            this.inner = continuationToken != null
+                ? await this.queryPipelineStageFactory(TryUnwrapContinuationToken(continuationToken))
+                : await this.queryPipelineStageFactory(null);
+
+            if (this.inner.Failed)
             {
-                CosmosObject cosmosObject = this.continuationToken as CosmosObject;
-                CosmosElement backendContinuationToken = cosmosObject[optimisticDirectExecutionToken];
-                Debug.Assert(backendContinuationToken != null);
-                return CosmosArray.Create(backendContinuationToken);
+                return false;
             }
 
-            return null;
+            return await this.inner.Result.MoveNextAsync(trace);
         }
 
         public static TryCatch<IQueryPipelineStage> MonadicCreate(
             DocumentContainer documentContainer,
             CosmosQueryExecutionContextFactory.InputParameters inputParameters,
             FeedRangeEpk targetRange,
-            QueryPaginationOptions queryPaginationOptions,
             FallbackQueryPipelineStageFactory fallbackQueryPipelineStageFactory,
             CancellationToken cancellationToken)
         {
+            QueryPaginationOptions paginationOptions = new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount, optimisticDirectExecute: true);
             TryCatch<IQueryPipelineStage> pipelineStage = OptimisticDirectExecutionQueryPipelineImpl.MonadicCreate(
                 documentContainer: documentContainer,
                 sqlQuerySpec: inputParameters.SqlQuerySpec,
                 targetRange: targetRange,
-                queryPaginationOptions: queryPaginationOptions,
+                queryPaginationOptions: paginationOptions,
                 partitionKey: inputParameters.PartitionKey,
                 continuationToken: inputParameters.InitialUserContinuationToken,
                 cancellationToken: cancellationToken);
@@ -120,7 +163,7 @@ public static TryCatch<IQueryPipelineStage> MonadicCreate(
             return TryCatch<IQueryPipelineStage>.FromResult(odePipelineStageMonadicCreate);
         }
 
-        private class OptimisticDirectExecutionQueryPipelineImpl : IQueryPipelineStage
+        private sealed class OptimisticDirectExecutionQueryPipelineImpl : IQueryPipelineStage
         {
             private readonly QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator;
 
@@ -226,7 +269,6 @@ public static TryCatch<IQueryPipelineStage> MonadicCreate(
                 }
 
                 FeedRangeState<QueryState> feedRangeState = monadicExtractState.Result;
-
                 QueryPartitionRangePageAsyncEnumerator partitionPageEnumerator = new QueryPartitionRangePageAsyncEnumerator(
                     documentContainer,
                     sqlQuerySpec,

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Pagination/QueryPaginationOptions.cs

@@ -15,6 +15,8 @@ internal sealed class QueryPaginationOptions : PaginationOptions
     {
         public static readonly QueryPaginationOptions Default = new QueryPaginationOptions();
 
+        public bool OptimisticDirectExecute { get; }
+
         public static readonly ImmutableHashSet<string> BannedHeaders = new HashSet<string>()
         {
             HttpConstants.HttpHeaders.Continuation,
@@ -29,10 +31,12 @@ internal sealed class QueryPaginationOptions : PaginationOptions
 
         public QueryPaginationOptions(
             int? pageSizeHint = null,
+            bool optimisticDirectExecute = false,
             JsonSerializationFormat? jsonSerializationFormat = null,
             Dictionary<string, string> additionalHeaders = null)
             : base(pageSizeHint, jsonSerializationFormat, additionalHeaders)
         {
+            this.OptimisticDirectExecute = optimisticDirectExecute;
         }
 
         protected override ImmutableHashSet<string> BannedAdditionalHeaders => BannedHeaders;

Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs

@@ -54,12 +54,11 @@ public abstract Task<TryCatch<QueryPage>> ExecuteItemQueryAsync(
             string resourceUri,
             Documents.ResourceType resourceType,
             Documents.OperationType operationType,
-            Guid clientQueryCorrelationId,
             FeedRange feedRange,
             QueryRequestOptions requestOptions,
+            AdditionalRequestHeaders additionalRequestHeaders,
             SqlQuerySpec sqlQuerySpec,
             string continuationToken,
-            bool isContinuationExpected,
             int pageSize,
             ITrace trace,
             CancellationToken cancellationToken);

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs

@@ -119,12 +119,11 @@ public override async Task<TryCatch<QueryPage>> ExecuteItemQueryAsync(
             string resourceUri,
             ResourceType resourceType,
             OperationType operationType,
-            Guid clientQueryCorrelationId,
             FeedRange feedRange,
             QueryRequestOptions requestOptions,
+            AdditionalRequestHeaders additionalRequestHeaders,
             SqlQuerySpec sqlQuerySpec,
             string continuationToken,
-            bool isContinuationExpected,
             int pageSize,
             ITrace trace,
             CancellationToken cancellationToken)
@@ -143,13 +142,14 @@ public override async Task<TryCatch<QueryPage>> ExecuteItemQueryAsync(
                 {
                     cosmosRequestMessage.Headers.Add(
                         HttpConstants.HttpHeaders.IsContinuationExpected,
-                        isContinuationExpected.ToString());
+                        additionalRequestHeaders.IsContinuationExpected.ToString());
                     QueryRequestOptions.FillContinuationToken(
                         cosmosRequestMessage,
                         continuationToken);
                     cosmosRequestMessage.Headers.Add(HttpConstants.HttpHeaders.ContentType, MediaTypes.QueryJson);
                     cosmosRequestMessage.Headers.Add(HttpConstants.HttpHeaders.IsQuery, bool.TrueString);
-                    cosmosRequestMessage.Headers.Add(WFConstants.BackendHeaders.CorrelatedActivityId, clientQueryCorrelationId.ToString());
+                    cosmosRequestMessage.Headers.Add(WFConstants.BackendHeaders.CorrelatedActivityId, additionalRequestHeaders.CorrelatedActivityId.ToString());
+                    cosmosRequestMessage.Headers.Add(HttpConstants.HttpHeaders.OptimisticDirectExecute, additionalRequestHeaders.OptimisticDirectExecute.ToString());
                 },
                 trace: trace,
                 cancellationToken: cancellationToken);

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryContextCore.cs

@@ -51,16 +51,17 @@ internal override Task<TryCatch<QueryPage>> ExecuteQueryAsync(
             ITrace trace,
             CancellationToken cancellationToken)
         {
+            AdditionalRequestHeaders additionalRequestHeaders = new AdditionalRequestHeaders(this.CorrelatedActivityId, isContinuationExpected, optimisticDirectExecute: false);
+
             return this.QueryClient.ExecuteItemQueryAsync(
                 resourceUri: this.ResourceLink,
                 resourceType: this.ResourceTypeEnum,
                 operationType: this.OperationTypeEnum,
-                clientQueryCorrelationId: this.CorrelatedActivityId,
                 requestOptions: queryRequestOptions,
+                additionalRequestHeaders: additionalRequestHeaders,
                 sqlQuerySpec: querySpecForInit,
                 continuationToken: continuationToken,
                 feedRange: feedRange,
-                isContinuationExpected: isContinuationExpected,
                 pageSize: pageSize,
                 trace: trace,
                 cancellationToken: cancellationToken);

Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs

@@ -44,7 +44,7 @@ public class QueryRequestOptions : RequestOptions
         public bool? EnableLowPrecisionOrderBy { get; set; }
 
         /// <summary>
-        /// Gets or sets the option for customers to opt in for direct (optimistic) execution of the query..
+        /// Gets or sets the option for customers to opt in for direct (optimistic) execution of the query.
         /// </summary>
         /// <value>
         /// Direct (optimistic) execution offers improved performance for several kinds of queries such as a single partition streaming query.