PriorityBasedExecution: Adds PriorityLevel in CosmosClientOptions

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -280,6 +280,21 @@ public ConnectionMode ConnectionMode
         /// </summary>
         public ConsistencyLevel? ConsistencyLevel { get; set; }
 
+        /// <summary>
+        /// Sets the priority level for requests created using cosmos client.
+        /// </summary>
+        /// <remarks>
+        /// If priority level is also set at request level in <see cref="RequestOptions.PriorityLevel"/>, that priority is used.
+        /// If <see cref="AllowBulkExecution"/> is set to true in CosmosClientOptions, priority level set on the CosmosClient is used.
+        /// </remarks>
+        /// <seealso href="https://aka.ms/CosmosDB/PriorityBasedExecution"/>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        PriorityLevel? PriorityLevel { get; set; }
+
         /// <summary>
         /// Gets or sets the maximum number of retries in the case where the request fails
         /// because the Azure Cosmos DB service has applied rate limiting on the client.
@@ -731,8 +746,8 @@ internal Protocol ConnectionProtocol
         /// Flag that controls whether CPU monitoring thread is created to enrich timeout exceptions with additional diagnostic. Default value is true.
         /// </summary>
         internal bool? EnableCpuMonitor { get; set; }
-
-        /// <summary>
+
+        /// <summary>
         /// Flag indicates the value of DisableServerCertificateValidation flag set at connection string level.Default it is false.
         /// </summary>
         internal bool DisableServerCertificateValidation { get; set; }
@@ -852,13 +867,13 @@ internal virtual ConnectionPolicy GetConnectionPolicy(int clientId)
             return (Documents.ConsistencyLevel)this.ConsistencyLevel.Value;
         }
 
-        internal static string GetAccountEndpoint(string connectionString)
-        {
+        internal static string GetAccountEndpoint(string connectionString)
+        {
             return CosmosClientOptions.GetValueFromConnectionString<string>(connectionString, CosmosClientOptions.ConnectionStringAccountEndpoint, null);
-        }
-
-        internal static string GetAccountKey(string connectionString)
-        {
+        }
+
+        internal static string GetAccountKey(string connectionString)
+        {
             return CosmosClientOptions.GetValueFromConnectionString<string>(connectionString, CosmosClientOptions.ConnectionStringAccountKey, null);
         }
 
@@ -876,21 +891,21 @@ internal static CosmosClientOptions GetCosmosClientOptionsWithCertificateFlag(st
             }
 
             return clientOptions;
-        }
-
+        }
+
         private static T GetValueFromConnectionString<T>(string connectionString, string keyName, T defaultValue)
-        {
-            if (connectionString == null)
-            {
-                throw new ArgumentNullException(nameof(connectionString));
-            }
-
-            DbConnectionStringBuilder builder = new DbConnectionStringBuilder { ConnectionString = connectionString };
-            if (builder.TryGetValue(keyName, out object value))
-            {
-                string keyNameValue = value as string;
-                if (!string.IsNullOrEmpty(keyNameValue))
-                {
+        {
+            if (connectionString == null)
+            {
+                throw new ArgumentNullException(nameof(connectionString));
+            }
+
+            DbConnectionStringBuilder builder = new DbConnectionStringBuilder { ConnectionString = connectionString };
+            if (builder.TryGetValue(keyName, out object value))
+            {
+                string keyNameValue = value as string;
+                if (!string.IsNullOrEmpty(keyNameValue))
+                {
                     try
                     {
                         return (T)Convert.ChangeType(value, typeof(T));
@@ -905,9 +920,9 @@ private static T GetValueFromConnectionString<T>(string connectionString, string
             if (defaultValue != null)
             {
                 return defaultValue;
-            }
-
-            throw new ArgumentException("The connection string is missing a required property: " + keyName);
+            }
+
+            throw new ArgumentException("The connection string is missing a required property: " + keyName);
         }
 
         private void ValidateLimitToEndpointSettings()
@@ -935,8 +950,8 @@ private void ValidatePartitionLevelFailoverSettings()
             {
                 throw new ArgumentException($"{nameof(this.ApplicationPreferredRegions)} is required when {nameof(this.EnablePartitionLevelFailover)} is enabled.");
             }
-        }
-
+        }
+
         private void ValidateAndSetServerCallbackSettings()
         {
             if (this.DisableServerCertificateValidation && this.ServerCertificateCustomValidationCallback != null)

Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs

@@ -442,6 +442,27 @@ public CosmosClientBuilder WithConsistencyLevel(Cosmos.ConsistencyLevel consiste
 
         }
 
+        /// <summary>
+        /// Sets the priority level for requests created using cosmos client.
+        /// </summary>
+        /// <remarks>
+        /// If priority level is also set at request level in <see cref="RequestOptions.PriorityLevel"/>, that priority is used.
+        /// If <see cref="WithBulkExecution(bool)"/> is set to true, priority level set on the CosmosClient is used.
+        /// </remarks>
+        /// <param name="priorityLevel">The desired priority level for the client.</param>
+        /// <returns>The current <see cref="CosmosClientBuilder"/>.</returns>
+        /// <seealso href="https://aka.ms/CosmosDB/PriorityBasedExecution"/>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        CosmosClientBuilder WithPriorityLevel(Cosmos.PriorityLevel priorityLevel)
+        {
+            this.clientOptions.PriorityLevel = priorityLevel;
+            return this;
+        }
+
         /// <summary>
         /// Sets the connection mode to Gateway. This is used by the client when connecting to the Azure Cosmos DB service.
         /// </summary>

Microsoft.Azure.Cosmos/src/Handler/ClientPipelineBuilder.cs

@@ -15,6 +15,7 @@ internal class ClientPipelineBuilder
     {
         private readonly CosmosClient client;
         private readonly ConsistencyLevel? requestedClientConsistencyLevel;
+        private readonly PriorityLevel? requestedPriorityLevel;
         private readonly DiagnosticsHandler diagnosticsHandler;
         private readonly RequestHandler invalidPartitionExceptionRetryHandler;
         private readonly RequestHandler transportHandler;
@@ -26,11 +27,13 @@ internal class ClientPipelineBuilder
         public ClientPipelineBuilder(
             CosmosClient client,
             ConsistencyLevel? requestedClientConsistencyLevel,
+            PriorityLevel? requestedClientPriorityLevel,
             IReadOnlyCollection<RequestHandler> customHandlers,
             TelemetryToServiceHelper telemetryToServiceHelper)
         {
             this.client = client ?? throw new ArgumentNullException(nameof(client));
             this.requestedClientConsistencyLevel = requestedClientConsistencyLevel;
+            this.requestedPriorityLevel = requestedClientPriorityLevel;
             this.transportHandler = new TransportHandler(client);
             Debug.Assert(this.transportHandler.InnerHandler == null, nameof(this.transportHandler));
 
@@ -149,7 +152,8 @@ public RequestInvokerHandler Build()
         {
             RequestInvokerHandler root = new RequestInvokerHandler(
                 this.client,
-                this.requestedClientConsistencyLevel);
+                this.requestedClientConsistencyLevel,
+                this.requestedPriorityLevel);
 
             RequestHandler current = root;
             if (this.CustomHandlers != null && this.CustomHandlers.Any())

Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs

@@ -27,16 +27,19 @@ internal class RequestInvokerHandler : RequestHandler
 
         private readonly CosmosClient client;
         private readonly Cosmos.ConsistencyLevel? RequestedClientConsistencyLevel;
+        private readonly Cosmos.PriorityLevel? RequestedClientPriorityLevel;
 
         private bool? IsLocalQuorumConsistency;
         private Cosmos.ConsistencyLevel? AccountConsistencyLevel = null;
 
         public RequestInvokerHandler(
             CosmosClient client,
-            Cosmos.ConsistencyLevel? requestedClientConsistencyLevel)
+            Cosmos.ConsistencyLevel? requestedClientConsistencyLevel,
+            Cosmos.PriorityLevel? requestedClientPriorityLevel)
         {
             this.client = client;
             this.RequestedClientConsistencyLevel = requestedClientConsistencyLevel;       
+            this.RequestedClientPriorityLevel = requestedClientPriorityLevel;
         }
 
         public override async Task<ResponseMessage> SendAsync(
@@ -66,6 +69,8 @@ public override async Task<ResponseMessage> SendAsync(
             }
 
             await this.ValidateAndSetConsistencyLevelAsync(request);
+            this.SetPriorityLevel(request);
+
             (bool isError, ResponseMessage errorResponse) = await this.EnsureValidClientAsync(request, request.Trace);
             if (isError)
             {
@@ -431,6 +436,25 @@ private async Task ValidateAndSetConsistencyLevelAsync(RequestMessage requestMes
             }
         }
 
+        /// <summary>
+        /// Set the PriorityLevel in the request headers
+        /// </summary>
+        /// <param name="requestMessage"></param>
+        private void SetPriorityLevel(RequestMessage requestMessage)
+        {
+            Cosmos.PriorityLevel? priorityLevel = this.RequestedClientPriorityLevel;
+            RequestOptions promotedRequestOptions = requestMessage.RequestOptions;
+            if (promotedRequestOptions?.PriorityLevel.HasValue == true)
+            {
+                priorityLevel = promotedRequestOptions.PriorityLevel;
+            }
+
+            if (priorityLevel.HasValue)
+            {
+                requestMessage.Headers.Set(HttpConstants.HttpHeaders.PriorityLevel, priorityLevel.ToString());
+            }
+        }
+
         internal static bool ShouldSetNoContentResponseHeaders(RequestOptions requestOptions,
             CosmosClientOptions clientOptions,
             OperationType operationType,

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

@@ -47,6 +47,7 @@ public class RequestOptions
         /// <remarks>
         /// Setting priority level only has an effect if Priority Based Execution is enabled.
         /// If it is not enabled, the priority level is ignored by the backend.
+        /// If <see cref="CosmosClientOptions.AllowBulkExecution"/> is set to true on CosmosClient, priority level set in RequestOptions is ignored.
         /// Default PriorityLevel for each request is treated as High. It can be explicitly set to Low for some requests.
         /// When Priority based execution is enabled, if there are more requests than the configured RU/S in a second, 
         /// then Cosmos DB will throttle low priority requests to allow high priority requests to execute.

Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs

@@ -120,6 +120,7 @@ internal static CosmosClientContext Create(
                 ClientPipelineBuilder clientPipelineBuilder = new ClientPipelineBuilder(
                     cosmosClient,
                     clientOptions.ConsistencyLevel,
+                    clientOptions.PriorityLevel,
                     clientOptions.CustomHandlers,
                     telemetryToServiceHelper: documentClient.telemetryToServiceHelper);
 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetPreviewSDKAPI.json

@@ -352,6 +352,31 @@
       },
       "NestedTypes": {}
     },
+    "Microsoft.Azure.Cosmos.CosmosClientOptions;System.Object;IsAbstract:False;IsSealed:False;IsInterface:False;IsEnum:False;IsClass:True;IsValueType:False;IsNested:False;IsGenericType:False;IsSerializable:False": {
+      "Subclasses": {},
+      "Members": {
+        "System.Nullable`1[Microsoft.Azure.Cosmos.PriorityLevel] get_PriorityLevel()[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
+          "Type": "Method",
+          "Attributes": [
+            "CompilerGeneratedAttribute"
+          ],
+          "MethodInfo": "System.Nullable`1[Microsoft.Azure.Cosmos.PriorityLevel] get_PriorityLevel();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "System.Nullable`1[Microsoft.Azure.Cosmos.PriorityLevel] PriorityLevel": {
+          "Type": "Property",
+          "Attributes": [],
+          "MethodInfo": "System.Nullable`1[Microsoft.Azure.Cosmos.PriorityLevel] PriorityLevel;CanRead:True;CanWrite:True;System.Nullable`1[Microsoft.Azure.Cosmos.PriorityLevel] get_PriorityLevel();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;Void set_PriorityLevel(System.Nullable`1[Microsoft.Azure.Cosmos.PriorityLevel]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "Void set_PriorityLevel(System.Nullable`1[Microsoft.Azure.Cosmos.PriorityLevel])[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
+          "Type": "Method",
+          "Attributes": [
+            "CompilerGeneratedAttribute"
+          ],
+          "MethodInfo": "Void set_PriorityLevel(System.Nullable`1[Microsoft.Azure.Cosmos.PriorityLevel]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        }
+      },
+      "NestedTypes": {}
+    },
     "Microsoft.Azure.Cosmos.CosmosLinqSerializer;Microsoft.Azure.Cosmos.CosmosSerializer;IsAbstract:True;IsSealed:False;IsInterface:False;IsEnum:False;IsClass:True;IsValueType:False;IsNested:False;IsGenericType:False;IsSerializable:False": {
       "Subclasses": {},
       "Members": {
@@ -454,6 +479,17 @@
       },
       "NestedTypes": {}
     },
+    "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder;System.Object;IsAbstract:False;IsSealed:False;IsInterface:False;IsEnum:False;IsClass:True;IsValueType:False;IsNested:False;IsGenericType:False;IsSerializable:False": {
+      "Subclasses": {},
+      "Members": {
+        "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithPriorityLevel(Microsoft.Azure.Cosmos.PriorityLevel)": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithPriorityLevel(Microsoft.Azure.Cosmos.PriorityLevel);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        }
+      },
+      "NestedTypes": {}
+    },
     "Microsoft.Azure.Cosmos.Linq.CosmosLinqExtensions;System.Object;IsAbstract:True;IsSealed:True;IsInterface:False;IsEnum:False;IsClass:True;IsValueType:False;IsNested:False;IsGenericType:False;IsSerializable:False": {
       "Subclasses": {},
       "Members": {