PriorityBasedExecution: Adds PriorityLevel in CosmosClientOptions (#4262

) * Added PriorityLevel in CosmosClientoptions * Addressed comments on PR * Changed remark to remarks * Corrected null check * Updated contracts * reverted changes to PartitionKeyHashBaselineTest.Lists.xml * Update Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs Co-authored-by: Matias Quaranta <[email protected]> * Update Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs Co-authored-by: Matias Quaranta <[email protected]> * Making priority level as GA contract * fixing the preview contract API * resetting few misc files * Builder and contract changes: --------- Co-authored-by: Achint Agrawal <[email protected]> Co-authored-by: Debdatta Kunda <[email protected]> Co-authored-by: Kiran Kumar Kolli <[email protected]> Co-authored-by: Matias Quaranta <[email protected]>
Azure · Jan 31, 2024 · d7a954f · d7a954f
1 parent 7b18607
commit d7a954f
Show file tree

Hide file tree

Showing 12 changed files with 271 additions and 105 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs b/Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs
@@ -280,6 +280,16 @@ 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"/>
+        public 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 +741,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 +862,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 +886,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 +915,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 +945,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)

diff --git a/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs b/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs
@@ -442,6 +442,22 @@ 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"/>
+        public 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>

diff --git a/Microsoft.Azure.Cosmos/src/Handler/ClientPipelineBuilder.cs b/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())

diff --git a/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs b/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.Value;
+            }
+
+            if (priorityLevel.HasValue)
+            {
+                requestMessage.Headers.Set(HttpConstants.HttpHeaders.PriorityLevel, priorityLevel.ToString());
+            }
+        }
+
         internal static bool ShouldSetNoContentResponseHeaders(RequestOptions requestOptions,
             CosmosClientOptions clientOptions,
             OperationType operationType,

diff --git a/Microsoft.Azure.Cosmos/src/RequestOptions/RequestOptions.cs b/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.
@@ -55,12 +56,7 @@ public class RequestOptions
         /// configured RU/s, low priority requests start getting throttled first to allow execution of mission critical workloads.
         /// </remarks>
         /// <seealso href="https://aka.ms/CosmosDB/PriorityBasedExecution"/>
-#if PREVIEW
-        public
-#else
-        internal
-#endif
-        PriorityLevel? PriorityLevel { get; set; }
+        public PriorityLevel? PriorityLevel { get; set; }
 
         /// <summary>
         /// Threshold values for Distributed Tracing. 

diff --git a/Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs b/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);
 

diff --git a/Microsoft.Azure.Cosmos/src/Resource/Settings/PriorityLevel.cs b/Microsoft.Azure.Cosmos/src/Resource/Settings/PriorityLevel.cs
@@ -18,13 +18,7 @@ namespace Microsoft.Azure.Cosmos
     /// configured RU/s, low priority requests start getting throttled first to allow execution of mission critical workloads.
     /// </remarks>
     /// <seealso href="https://aka.ms/CosmosDB/PriorityBasedExecution"/>
-
-#if PREVIEW
-    public
-#else
-    internal
-#endif
-    enum PriorityLevel
+    public enum PriorityLevel
     {
         /// <summary> 
         /// High Priority

diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetPreviewSDKAPI.json b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetPreviewSDKAPI.json
@@ -466,52 +466,6 @@
         }
       },
       "NestedTypes": {}
-    },
-    "Microsoft.Azure.Cosmos.PriorityLevel;System.Enum;IsAbstract:False;IsSealed:True;IsInterface:False;IsEnum:True;IsClass:False;IsValueType:True;IsNested:False;IsGenericType:False;IsSerializable:True": {
-      "Subclasses": {},
-      "Members": {
-        "Int32 value__": {
-          "Type": "Field",
-          "Attributes": [],
-          "MethodInfo": "Int32 value__;IsInitOnly:False;IsStatic:False;"
-        },
-        "Microsoft.Azure.Cosmos.PriorityLevel High": {
-          "Type": "Field",
-          "Attributes": [],
-          "MethodInfo": "Microsoft.Azure.Cosmos.PriorityLevel High;IsInitOnly:False;IsStatic:True;"
-        },
-        "Microsoft.Azure.Cosmos.PriorityLevel Low": {
-          "Type": "Field",
-          "Attributes": [],
-          "MethodInfo": "Microsoft.Azure.Cosmos.PriorityLevel Low;IsInitOnly:False;IsStatic:True;"
-        }
-      },
-      "NestedTypes": {}
-    },
-    "Microsoft.Azure.Cosmos.RequestOptions;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": {}
     }
   },
   "Members": {},