ThroughputBucketing: Adds Throughput Bucket to RequestOptions and ClientOptions

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -940,6 +940,21 @@ public IFaultInjector FaultInjector
             }
         }
 
+        /// <summary>
+        /// Sets the throughput bucket for requests created using cosmos client.
+        /// </summary>
+        /// <remarks>
+        /// If throughput bucket is also set at request level in <see cref="RequestOptions.ThroughputBucket"/>, that throughput bucket is used.
+        /// If <see cref="AllowBulkExecution"/> is set to true in CosmosClientOptions, throughput bucket can only be set at client level.
+        /// </remarks>
+        /// <seealso href="https://aka.ms/cosmsodb-bucketing"/>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        int? ThroughputBucket { get; set; }
+
         internal IChaosInterceptorFactory ChaosInterceptorFactory { get; set; }
 
         internal void SetSerializerIfNotConfigured(CosmosSerializer serializer)

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -7096,6 +7096,12 @@ private INameValueCollection GetRequestHeaders(
             {
                 headers.Set(HttpConstants.HttpHeaders.PreserveFullContent, bool.TrueString);
             }
+
+            if (options.ThroughputBucket.HasValue)
+            {
+                headers.Set(HttpConstants.HttpHeaders.ThroughputBucket, options.ThroughputBucket?.ToString(CultureInfo.InvariantCulture));
+            }
+
             return headers;
         }
 

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

@@ -802,5 +802,26 @@ public CosmosClientBuilder WithClientTelemetryOptions(CosmosClientTelemetryOptio
             this.clientOptions.CosmosClientTelemetryOptions = options;
             return this;
         }
+
+        /// <summary>
+        /// Sets the throughput bucket for requests created using cosmos client.
+        /// </summary>
+        /// <remarks>
+        /// If throughput bucket is also set at request level in <see cref="RequestOptions.ThroughputBucket"/>, that throughput bucket is used.
+        /// If <see cref="WithBulkExecution(bool)"/> is set to true, throughput bucket can only be set at client level.
+        /// </remarks>
+        /// <param name="throughputBucket">The desired throughput bucket for the client.</param>
+        /// <returns>The current <see cref="CosmosClientBuilder"/>.</returns>
+        /// <seealso href="https://aka.ms/cosmsodb-bucketing"/>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        CosmosClientBuilder WithThroughputBucket(int throughputBucket)
+        {
+            this.clientOptions.ThroughputBucket = throughputBucket;
+            return this;
+        }
     }
 }

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

@@ -20,6 +20,7 @@ internal class ClientPipelineBuilder
         private readonly RequestHandler invalidPartitionExceptionRetryHandler;
         private readonly RequestHandler transportHandler;
         private readonly TelemetryHandler telemetryHandler;
+        private readonly int? requestedClientThroughputBucket;
 
         private IReadOnlyCollection<RequestHandler> customHandlers;
         private RequestHandler retryHandler;
@@ -29,11 +30,13 @@ public ClientPipelineBuilder(
             ConsistencyLevel? requestedClientConsistencyLevel,
             PriorityLevel? requestedClientPriorityLevel,
             IReadOnlyCollection<RequestHandler> customHandlers,
-            TelemetryToServiceHelper telemetryToServiceHelper)
+            TelemetryToServiceHelper telemetryToServiceHelper,
+            int? requestedClientThroughputBucket)
         {
             this.client = client ?? throw new ArgumentNullException(nameof(client));
             this.requestedClientConsistencyLevel = requestedClientConsistencyLevel;
             this.requestedPriorityLevel = requestedClientPriorityLevel;
+            this.requestedClientThroughputBucket = requestedClientThroughputBucket;
             this.transportHandler = new TransportHandler(client);
             Debug.Assert(this.transportHandler.InnerHandler == null, nameof(this.transportHandler));
 
@@ -153,7 +156,8 @@ public RequestInvokerHandler Build()
             RequestInvokerHandler root = new RequestInvokerHandler(
                 this.client,
                 this.requestedClientConsistencyLevel,
-                this.requestedPriorityLevel);
+                this.requestedPriorityLevel,
+                this.requestedClientThroughputBucket);
 
             RequestHandler current = root;
             if (this.CustomHandlers != null && this.CustomHandlers.Any())

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

@@ -30,19 +30,22 @@ internal class RequestInvokerHandler : RequestHandler
         private readonly CosmosClient client;
         private readonly Cosmos.ConsistencyLevel? RequestedClientConsistencyLevel;
         private readonly Cosmos.PriorityLevel? RequestedClientPriorityLevel;
+        private readonly int? RequestedClientThroughputBucket;
 
         private bool? IsLocalQuorumConsistency;
         private Cosmos.ConsistencyLevel? AccountConsistencyLevel = null;
 
         public RequestInvokerHandler(
             CosmosClient client,
             Cosmos.ConsistencyLevel? requestedClientConsistencyLevel,
-            Cosmos.PriorityLevel? requestedClientPriorityLevel)
+            Cosmos.PriorityLevel? requestedClientPriorityLevel,
+            int? requestedClientThroughputBucket)
         {
             this.client = client;
 
             this.RequestedClientConsistencyLevel = requestedClientConsistencyLevel;       
             this.RequestedClientPriorityLevel = requestedClientPriorityLevel;
+            this.RequestedClientThroughputBucket = requestedClientThroughputBucket;
         }
 
         public override async Task<ResponseMessage> SendAsync(
@@ -76,6 +79,7 @@ public override async Task<ResponseMessage> SendAsync(
 
             await this.ValidateAndSetConsistencyLevelAsync(request);
             this.SetPriorityLevel(request);
+            this.ValidateAndSetThroughputBucket(request);
 
             (bool isError, ResponseMessage errorResponse) = await this.EnsureValidClientAsync(request, request.Trace);
             if (isError)
@@ -513,6 +517,32 @@ private void SetPriorityLevel(RequestMessage requestMessage)
             }
         }
 
+        /// <summary>
+        /// Set the ThroughputBucket in the request headers
+        /// </summary>
+        /// <param name="requestMessage"></param>
+        private void ValidateAndSetThroughputBucket(RequestMessage requestMessage)
+        {
+            int? throughputBucket = this.RequestedClientThroughputBucket;
+            RequestOptions promotedRequestOptions = requestMessage.RequestOptions;
+
+            if (promotedRequestOptions?.ThroughputBucket.HasValue == true)
+            {
+                if (this.client.ClientOptions.AllowBulkExecution)
+                {
+                    throw new ArgumentException($"{nameof(requestMessage.RequestOptions.ThroughputBucket)} cannot be set in " +
+                        $"{nameof(requestMessage.RequestOptions)} when {nameof(this.client.ClientOptions.AllowBulkExecution)} is set to true. " +
+                        $"Instead, set {nameof(this.client.ClientOptions.ThroughputBucket)} only in {nameof(this.client.ClientOptions)}.");
+                }
+                throughputBucket = promotedRequestOptions.ThroughputBucket.Value;
+            }
+
+            if (throughputBucket.HasValue)
+            {
+                requestMessage.Headers.Set(HttpConstants.HttpHeaders.ThroughputBucket, throughputBucket.ToString());
+            }
+        }
+
         internal static bool ShouldSetNoContentResponseHeaders(RequestOptions requestOptions,
             CosmosClientOptions clientOptions,
             OperationType operationType,

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

@@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos
 {
     using System;
     using System.Collections.Generic;
+    using System.Globalization;
     using Microsoft.Azure.Documents;
 
     /// <summary>
@@ -114,6 +115,21 @@ public class RequestOptions
         internal virtual ConsistencyLevel? BaseConsistencyLevel { get; set; }
 
         internal bool DisablePointOperationDiagnostics { get; set; }
+
+        /// <summary>
+        /// Gets or sets the throughput bucket for a request.
+        /// </summary>
+        /// <remarks>
+        /// If <see cref="CosmosClientOptions.AllowBulkExecution"/> is set to true on CosmosClient,
+        /// <see cref="RequestOptions.ThroughputBucket"/> cannot be set in RequestOptions.
+        /// </remarks>
+        /// <seealso href="https://aka.ms/cosmsodb-bucketing"/>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        int? ThroughputBucket { get; set; }
 
         /// <summary>
         /// Fill the CosmosRequestMessage headers with the set properties
@@ -143,6 +159,11 @@ internal virtual void PopulateRequestOptions(RequestMessage request)
             {
                 request.Headers.Add(HttpConstants.HttpHeaders.PriorityLevel, this.PriorityLevel.ToString());
             }
+
+            if (this.ThroughputBucket.HasValue)
+            {
+                request.Headers.Add(HttpConstants.HttpHeaders.ThroughputBucket, this.ThroughputBucket?.ToString(CultureInfo.InvariantCulture));
+            }
 
             this.AddRequestHeaders?.Invoke(request.Headers);
         }

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

@@ -123,7 +123,8 @@ internal static CosmosClientContext Create(
                     clientOptions.ConsistencyLevel,
                     clientOptions.PriorityLevel,
                     clientOptions.CustomHandlers,
-                    telemetryToServiceHelper: documentClient.telemetryToServiceHelper);
+                    telemetryToServiceHelper: documentClient.telemetryToServiceHelper,
+                    clientOptions.ThroughputBucket);
 
                 requestInvokerHandler = clientPipelineBuilder.Build();
             }

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

@@ -403,12 +403,31 @@
           ],
           "MethodInfo": "Microsoft.Azure.Cosmos.AvailabilityStrategy get_AvailabilityStrategy();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         },
+        "System.Nullable`1[System.Int32] get_ThroughputBucket()[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
+          "Type": "Method",
+          "Attributes": [
+            "CompilerGeneratedAttribute"
+          ],
+          "MethodInfo": "System.Nullable`1[System.Int32] get_ThroughputBucket();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "System.Nullable`1[System.Int32] ThroughputBucket": {
+          "Type": "Property",
+          "Attributes": [],
+          "MethodInfo": "System.Nullable`1[System.Int32] ThroughputBucket;CanRead:True;CanWrite:True;System.Nullable`1[System.Int32] get_ThroughputBucket();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;Void set_ThroughputBucket(System.Nullable`1[System.Int32]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
         "Void set_AvailabilityStrategy(Microsoft.Azure.Cosmos.AvailabilityStrategy)[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
           "Type": "Method",
           "Attributes": [
             "CompilerGeneratedAttribute"
           ],
           "MethodInfo": "Void set_AvailabilityStrategy(Microsoft.Azure.Cosmos.AvailabilityStrategy);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "Void set_ThroughputBucket(System.Nullable`1[System.Int32])[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
+          "Type": "Method",
+          "Attributes": [
+            "CompilerGeneratedAttribute"
+          ],
+          "MethodInfo": "Void set_ThroughputBucket(System.Nullable`1[System.Int32]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         }
       },
       "NestedTypes": {}
@@ -474,6 +493,11 @@
           "Type": "Method",
           "Attributes": [],
           "MethodInfo": "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithAvailabilityStrategy(Microsoft.Azure.Cosmos.AvailabilityStrategy);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithThroughputBucket(Int32)": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithThroughputBucket(Int32);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         }
       },
       "NestedTypes": {}
@@ -741,12 +765,31 @@
           ],
           "MethodInfo": "Microsoft.Azure.Cosmos.AvailabilityStrategy get_AvailabilityStrategy();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         },
+        "System.Nullable`1[System.Int32] get_ThroughputBucket()[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
+          "Type": "Method",
+          "Attributes": [
+            "CompilerGeneratedAttribute"
+          ],
+          "MethodInfo": "System.Nullable`1[System.Int32] get_ThroughputBucket();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "System.Nullable`1[System.Int32] ThroughputBucket": {
+          "Type": "Property",
+          "Attributes": [],
+          "MethodInfo": "System.Nullable`1[System.Int32] ThroughputBucket;CanRead:True;CanWrite:True;System.Nullable`1[System.Int32] get_ThroughputBucket();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;Void set_ThroughputBucket(System.Nullable`1[System.Int32]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
         "Void set_AvailabilityStrategy(Microsoft.Azure.Cosmos.AvailabilityStrategy)[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
           "Type": "Method",
           "Attributes": [
             "CompilerGeneratedAttribute"
           ],
           "MethodInfo": "Void set_AvailabilityStrategy(Microsoft.Azure.Cosmos.AvailabilityStrategy);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "Void set_ThroughputBucket(System.Nullable`1[System.Int32])[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
+          "Type": "Method",
+          "Attributes": [
+            "CompilerGeneratedAttribute"
+          ],
+          "MethodInfo": "Void set_ThroughputBucket(System.Nullable`1[System.Int32]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         }
       },
       "NestedTypes": {}