PriorityBasedExecution: Adds PriorityLevel in CosmosClientOptions

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -278,6 +278,17 @@ public ConnectionMode ConnectionMode
         /// If this is not set the database account consistency level will be used for all requests.
         /// </summary>
         public ConsistencyLevel? ConsistencyLevel { get; set; }
+
+        /// <summary>
+        /// Sets the PriorityLevel for requests created using cosmos client.
+        /// If PriorityLevel is also specified at request level in RequestOptions, that PriorityLevel takes precedence.
+        /// </summary>
+#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
@@ -842,6 +853,16 @@ internal virtual ConnectionPolicy GetConnectionPolicy(int clientId)
 
             return (Documents.ConsistencyLevel)this.ConsistencyLevel.Value;
         }
+
+        internal Documents.PriorityLevel? GetDocumentsPriorityLevel()
+        {
+            if (!this.PriorityLevel.HasValue)
+            {
+                return null;
+            }
+
+            return (Documents.PriorityLevel)this.PriorityLevel.Value;
+        }
 
         internal static string GetAccountEndpoint(string connectionString)
         {

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

@@ -435,6 +435,24 @@ public CosmosClientBuilder WithConsistencyLevel(Cosmos.ConsistencyLevel consiste
 
         }
 
+        /// <summary>
+        /// Sets the PriorityLevel for requests created using cosmos client.
+        /// If PriorityLevel is also specified at request level in RequestOptions, that PriorityLevel takes precedence.
+        /// </summary>
+        /// <param name="priorityLevel">The desired consistency 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 @@ internal class RequestInvokerHandler : RequestHandler
             }
 
             await this.ValidateAndSetConsistencyLevelAsync(request);
+            this.SetPriorityLevel(request);
+
             (bool isError, ResponseMessage errorResponse) = await this.EnsureValidClientAsync(request, request.Trace);
             if (isError)
             {
@@ -431,6 +436,29 @@ 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 = null;
+            RequestOptions promotedRequestOptions = requestMessage.RequestOptions;
+            if (promotedRequestOptions != null && promotedRequestOptions.PriorityLevel.HasValue)
+            {
+                priorityLevel = promotedRequestOptions.PriorityLevel;
+            }
+            else if (this.RequestedClientPriorityLevel.HasValue)
+            {
+                priorityLevel = this.RequestedClientPriorityLevel;
+            }
+
+            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/Resource/ClientContextCore.cs

@@ -120,6 +120,7 @@ private static RemoteCertificateValidationCallback SslCustomValidationCallBack(F
                 ClientPipelineBuilder clientPipelineBuilder = new ClientPipelineBuilder(
                     cosmosClient,
                     clientOptions.ConsistencyLevel,
+                    clientOptions.PriorityLevel,
                     clientOptions.CustomHandlers,
                     telemetryToServiceHelper: documentClient.telemetryToServiceHelper);
 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosClientOptionsUnitTests.cs

@@ -52,6 +52,7 @@ public void VerifyCosmosConfigurationPropertiesGetUpdated()
             Cosmos.PortReuseMode portReuseMode = Cosmos.PortReuseMode.PrivatePortPool;
             IWebProxy webProxy = new TestWebProxy();
             Cosmos.ConsistencyLevel consistencyLevel = Cosmos.ConsistencyLevel.ConsistentPrefix;
+            Cosmos.PriorityLevel priorityLevel = Cosmos.PriorityLevel.Low;
 
             CosmosClientBuilder cosmosClientBuilder = new CosmosClientBuilder(
                 accountEndpoint: endpoint,
@@ -80,6 +81,7 @@ public void VerifyCosmosConfigurationPropertiesGetUpdated()
             Assert.IsTrue(clientOptions.EnableTcpConnectionEndpointRediscovery);
             Assert.IsNull(clientOptions.HttpClientFactory);
             Assert.AreNotEqual(consistencyLevel, clientOptions.ConsistencyLevel);
+            Assert.AreNotEqual(priorityLevel, clientOptions.PriorityLevel);
             Assert.IsFalse(clientOptions.EnablePartitionLevelFailover);
             Assert.IsFalse(clientOptions.EnableAdvancedReplicaSelectionForTcp.HasValue);
 
@@ -115,7 +117,8 @@ public void VerifyCosmosConfigurationPropertiesGetUpdated()
                 .WithThrottlingRetryOptions(maxRetryWaitTime, maxRetryAttemptsOnThrottledRequests)
                 .WithBulkExecution(true)
                 .WithSerializerOptions(cosmosSerializerOptions)
-                .WithConsistencyLevel(consistencyLevel);
+                .WithConsistencyLevel(consistencyLevel)
+                .WithPriorityLevel(priorityLevel);
 
             cosmosClient = cosmosClientBuilder.Build(new MockDocumentClient());
             clientOptions = cosmosClient.ClientOptions;
@@ -138,6 +141,7 @@ public void VerifyCosmosConfigurationPropertiesGetUpdated()
             Assert.IsTrue(object.ReferenceEquals(webProxy, clientOptions.WebProxy));
             Assert.IsTrue(clientOptions.AllowBulkExecution);
             Assert.AreEqual(consistencyLevel, clientOptions.ConsistencyLevel);
+            Assert.AreEqual(priorityLevel, clientOptions.PriorityLevel);
             Assert.IsFalse(clientOptions.EnablePartitionLevelFailover);
             Assert.IsTrue(clientOptions.EnableAdvancedReplicaSelectionForTcp.HasValue && clientOptions.EnableAdvancedReplicaSelectionForTcp.Value);
 
@@ -236,6 +240,8 @@ public void CosmosClientOptions_WhenPartitionLevelFailoverEnabledAndPreferredReg
                 };
 
                 Cosmos.ConsistencyLevel consistencyLevel = Cosmos.ConsistencyLevel.ConsistentPrefix;
+                Cosmos.PriorityLevel priorityLevel = Cosmos.PriorityLevel.Low;
+
                 CosmosClientBuilder cosmosClientBuilder = new(
                     accountEndpoint: endpoint,
                     authKeyOrResourceToken: key);
@@ -248,7 +254,8 @@ public void CosmosClientOptions_WhenPartitionLevelFailoverEnabledAndPreferredReg
                     .WithApiType(apiType)
                     .WithThrottlingRetryOptions(maxRetryWaitTime, maxRetryAttemptsOnThrottledRequests)
                     .WithSerializerOptions(cosmosSerializerOptions)
-                    .WithConsistencyLevel(consistencyLevel);
+                    .WithConsistencyLevel(consistencyLevel)
+                    .WithPriorityLevel(priorityLevel);
 
                 if (!useEnvironmentVariable)
                 {
@@ -300,6 +307,7 @@ public void CosmosClientOptions_WhenPartitionLevelFailoverEnabledAndPreferredReg
                 };
 
                 Cosmos.ConsistencyLevel consistencyLevel = Cosmos.ConsistencyLevel.ConsistentPrefix;
+                Cosmos.PriorityLevel priorityLevel = Cosmos.PriorityLevel.Low;
                 CosmosClientBuilder cosmosClientBuilder = new(
                     accountEndpoint: endpoint,
                     authKeyOrResourceToken: key);
@@ -313,6 +321,7 @@ public void CosmosClientOptions_WhenPartitionLevelFailoverEnabledAndPreferredReg
                     .WithThrottlingRetryOptions(maxRetryWaitTime, maxRetryAttemptsOnThrottledRequests)
                     .WithSerializerOptions(cosmosSerializerOptions)
                     .WithConsistencyLevel(consistencyLevel)
+                    .WithPriorityLevel(priorityLevel)
                     .WithPartitionLevelFailoverEnabled()
                     .WithApplicationPreferredRegions(
                         new List<string>()
@@ -370,6 +379,32 @@ public void VerifyConsisentencyLevels()
 
             Assert.IsNull(cosmosClientOptionsNull.GetDocumentsConsistencyLevel());
         }
+
+        [TestMethod]
+        public void VerifyPriorityLevels()
+        {
+            List<Cosmos.PriorityLevel> cosmosLevels = Enum.GetValues(typeof(Cosmos.PriorityLevel)).Cast<Cosmos.PriorityLevel>().ToList();
+            List<Documents.PriorityLevel> documentLevels = Enum.GetValues(typeof(Documents.PriorityLevel)).Cast<Documents.PriorityLevel>().ToList();
+            CollectionAssert.AreEqual(cosmosLevels, documentLevels, new EnumComparer(), "Document priority level is different from cosmos priority level");
+
+            foreach (Cosmos.PriorityLevel priorityLevel in cosmosLevels)
+            {
+                CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
+                {
+                    PriorityLevel = priorityLevel
+                };
+
+                Assert.AreEqual((int)priorityLevel, (int)cosmosClientOptions.GetDocumentsPriorityLevel());
+                Assert.AreEqual(priorityLevel.ToString(), cosmosClientOptions.GetDocumentsPriorityLevel().ToString());
+            }
+
+            CosmosClientOptions cosmosClientOptionsNull = new CosmosClientOptions()
+            {
+                PriorityLevel = null
+            };
+
+            Assert.IsNull(cosmosClientOptionsNull.GetDocumentsPriorityLevel());
+        }
 
         [TestMethod]
         public void VerifyPortReuseModeIsSyncedWithDirect()