[Internal] Thin Client Integration: Adds Thin Client Enablement Through Account Config and Reduce Request Headers for Proxy

Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs

@@ -28,7 +28,8 @@ internal sealed class ClientRetryPolicy : IDocumentClientRetryPolicy
         private readonly GlobalEndpointManager globalEndpointManager;
         private readonly GlobalPartitionEndpointManager partitionKeyRangeLocationCache;
         private readonly bool enableEndpointDiscovery;
-        private readonly bool isPartitionLevelFailoverEnabled;
+        private readonly bool isPartitionLevelFailoverEnabled;
+        private readonly bool isThinClientEnabled;
         private int failoverRetryCount;
 
         private int sessionTokenRetryCount;
@@ -45,7 +46,8 @@ public ClientRetryPolicy(
             GlobalPartitionEndpointManager partitionKeyRangeLocationCache,
             RetryOptions retryOptions,
             bool enableEndpointDiscovery,
-            bool isPartitionLevelFailoverEnabled)
+            bool isPartitionLevelFailoverEnabled,
+            bool isThinClientEnabled)
         {
             this.throttlingRetry = new ResourceThrottleRetryPolicy(
                 retryOptions.MaxRetryAttemptsOnThrottledRequests,
@@ -59,7 +61,8 @@ public ClientRetryPolicy(
             this.serviceUnavailableRetryCount = 0;
             this.canUseMultipleWriteLocations = false;
             this.isMultiMasterWriteRequest = false;
-            this.isPartitionLevelFailoverEnabled = isPartitionLevelFailoverEnabled;
+            this.isPartitionLevelFailoverEnabled = isPartitionLevelFailoverEnabled;
+            this.isThinClientEnabled = isThinClientEnabled;
         }
 
         /// <summary> 
@@ -226,7 +229,7 @@ public void OnBeforeSendRequest(DocumentServiceRequest request)
 
             // Resolve the endpoint for the request and pin the resolution to the resolved endpoint
             // This enables marking the endpoint unavailability on endpoint failover/unreachability
-            this.locationEndpoint = ConfigurationManager.IsThinClientEnabled(defaultValue: false)
+            this.locationEndpoint = this.isThinClientEnabled
                 && ThinClientStoreModel.IsOperationSupportedByThinClient(request)
                 ? this.globalEndpointManager.ResolveThinClientEndpoint(request)
                 : this.globalEndpointManager.ResolveServiceEndpoint(request);

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -126,16 +126,16 @@ internal partial class DocumentClient : IDisposable, IAuthorizationTokenProvider
         private readonly bool isReplicaAddressValidationEnabled;
         private readonly bool enableAsyncCacheExceptionNoSharing;
 
-        private readonly bool isThinClientEnabled;
-
         //Fault Injection
         private readonly IChaosInterceptorFactory chaosInterceptorFactory;
         private readonly IChaosInterceptor chaosInterceptor;
 
         private bool isChaosInterceptorInititalized = false;
 
         //Auth
-        internal readonly AuthorizationTokenProvider cosmosAuthorization;
+        internal readonly AuthorizationTokenProvider cosmosAuthorization;
+
+        private bool isThinClientEnabled = ConfigurationManager.IsThinClientEnabled(defaultValue: false);
 
         // Gateway has backoff/retry logic to hide transient errors.
         private RetryPolicy retryPolicy;
@@ -256,7 +256,6 @@ public DocumentClient(Uri serviceEndpoint,
                 cancellationToken: this.cancellationTokenSource.Token,
                 enableAsyncCacheExceptionNoSharing: this.enableAsyncCacheExceptionNoSharing);
             this.isReplicaAddressValidationEnabled = ConfigurationManager.IsReplicaAddressValidationEnabled(connectionPolicy);
-            this.isThinClientEnabled = ConfigurationManager.IsThinClientEnabled(defaultValue: false);
         }
 
         /// <summary>
@@ -514,7 +513,6 @@ internal DocumentClient(Uri serviceEndpoint,
                 enableAsyncCacheExceptionNoSharing: this.enableAsyncCacheExceptionNoSharing);
             this.chaosInterceptorFactory = chaosInterceptorFactory;
             this.chaosInterceptor = chaosInterceptorFactory?.CreateInterceptor(this);
-            this.isThinClientEnabled = ConfigurationManager.IsThinClientEnabled(defaultValue: false);
 
             this.Initialize(
                 serviceEndpoint: serviceEndpoint,
@@ -526,8 +524,7 @@ internal DocumentClient(Uri serviceEndpoint,
                 storeClientFactory: storeClientFactory,
                 cosmosClientId: cosmosClientId,
                 remoteCertificateValidationCallback: remoteCertificateValidationCallback,
-                cosmosClientTelemetryOptions: cosmosClientTelemetryOptions,
-                enableThinClientMode: this.isThinClientEnabled);
+                cosmosClientTelemetryOptions: cosmosClientTelemetryOptions);
         }
 
         /// <summary>
@@ -712,8 +709,7 @@ internal virtual void Initialize(Uri serviceEndpoint,
             TokenCredential tokenCredential = null,
             string cosmosClientId = null,
             RemoteCertificateValidationCallback remoteCertificateValidationCallback = null,
-            CosmosClientTelemetryOptions cosmosClientTelemetryOptions = null,
-            bool enableThinClientMode = false)
+            CosmosClientTelemetryOptions cosmosClientTelemetryOptions = null)
         {
             if (serviceEndpoint == null)
             {
@@ -1062,6 +1058,8 @@ private async Task<bool> GetInitializationTaskAsync(IStoreClientFactory storeCli
                 this.ConnectionPolicy.EnablePartitionLevelFailover = this.accountServiceConfiguration.AccountProperties.EnablePartitionLevelFailover.Value;
             }
 
+            this.isThinClientEnabled = (this.accountServiceConfiguration.AccountProperties?.ThinClientWritableLocationsInternal?.Count ?? 0) > 0;
+
             this.ConnectionPolicy.EnablePartitionLevelCircuitBreaker |= this.ConnectionPolicy.EnablePartitionLevelFailover;
             this.ConnectionPolicy.UserAgentContainer.AppendFeatures(this.GetUserAgentFeatures());
             this.InitializePartitionLevelFailoverWithDefaultHedging();
@@ -1072,13 +1070,15 @@ private async Task<bool> GetInitializationTaskAsync(IStoreClientFactory storeCli
                     ? new GlobalPartitionEndpointManagerCore(
                         this.GlobalEndpointManager,
                         this.ConnectionPolicy.EnablePartitionLevelFailover,
-                        this.ConnectionPolicy.EnablePartitionLevelCircuitBreaker)
+                        this.ConnectionPolicy.EnablePartitionLevelCircuitBreaker,
+                        this.isThinClientEnabled)
                     : GlobalPartitionEndpointManagerNoOp.Instance;
 
             this.retryPolicy = new RetryPolicy(
                 globalEndpointManager: this.GlobalEndpointManager,
                 connectionPolicy: this.ConnectionPolicy,
-                partitionKeyRangeLocationCache: this.PartitionKeyRangeLocation);
+                partitionKeyRangeLocationCache: this.PartitionKeyRangeLocation,
+                isThinClientEnabled: this.isThinClientEnabled);
 
             this.ResetSessionTokenRetryPolicy = this.retryPolicy;
 
@@ -1106,25 +1106,29 @@ private async Task<bool> GetInitializationTaskAsync(IStoreClientFactory storeCli
 
             gatewayStoreModel.SetCaches(this.partitionKeyRangeCache, this.collectionCache);
 
-            if (this.ConnectionPolicy.ConnectionMode == ConnectionMode.Gateway && this.isThinClientEnabled)
-            {
-                ThinClientStoreModel thinClientStoreModel = new (
-                    endpointManager: this.GlobalEndpointManager,
-                    this.PartitionKeyRangeLocation,
-                    this.sessionContainer,
-                    (Cosmos.ConsistencyLevel)this.accountServiceConfiguration.DefaultConsistencyLevel,
-                    this.eventSource,
-                    this.serializerSettings,
-                    this.httpClient,
-                    isPartitionLevelFailoverEnabled: this.ConnectionPolicy.EnablePartitionLevelFailover || this.ConnectionPolicy.EnablePartitionLevelCircuitBreaker);
-
-                thinClientStoreModel.SetCaches(this.partitionKeyRangeCache, this.collectionCache);
-
-                this.StoreModel = thinClientStoreModel;
-            }
-            else if (this.ConnectionPolicy.ConnectionMode == ConnectionMode.Gateway)
-            {
-                this.StoreModel = this.GatewayStoreModel;
+            if (this.ConnectionPolicy.ConnectionMode == ConnectionMode.Gateway)
+            {
+                if (this.isThinClientEnabled)
+                {
+                    ThinClientStoreModel thinClientStoreModel = new (
+                        endpointManager: this.GlobalEndpointManager,
+                        this.PartitionKeyRangeLocation,
+                        this.sessionContainer,
+                        (Cosmos.ConsistencyLevel)this.accountServiceConfiguration.DefaultConsistencyLevel,
+                        this.eventSource,
+                        this.serializerSettings,
+                        this.httpClient,
+                        this.ConnectionPolicy.UserAgentContainer,
+                        isPartitionLevelFailoverEnabled: this.ConnectionPolicy.EnablePartitionLevelFailover || this.ConnectionPolicy.EnablePartitionLevelCircuitBreaker);
+
+                    thinClientStoreModel.SetCaches(this.partitionKeyRangeCache, this.collectionCache);
+
+                    this.StoreModel = thinClientStoreModel;
+                }
+                else
+                {
+                    this.StoreModel = this.GatewayStoreModel;
+                }
             }
             else
             {

Microsoft.Azure.Cosmos/src/RetryPolicy.cs

@@ -13,7 +13,8 @@ internal sealed class RetryPolicy : IRetryPolicyFactory
         private readonly GlobalPartitionEndpointManager partitionKeyRangeLocationCache;
         private readonly GlobalEndpointManager globalEndpointManager;
         private readonly bool enableEndpointDiscovery;
-        private readonly bool isPartitionLevelFailoverEnabled;
+        private readonly bool isPartitionLevelFailoverEnabled;
+        private readonly bool isThinClientEnabled;
         private readonly RetryOptions retryOptions;
 
         /// <summary>
@@ -22,13 +23,15 @@ internal sealed class RetryPolicy : IRetryPolicyFactory
         public RetryPolicy(
             GlobalEndpointManager globalEndpointManager,
             ConnectionPolicy connectionPolicy,
-            GlobalPartitionEndpointManager partitionKeyRangeLocationCache)
+            GlobalPartitionEndpointManager partitionKeyRangeLocationCache,
+            bool isThinClientEnabled)
         {
             this.enableEndpointDiscovery = connectionPolicy.EnableEndpointDiscovery;
             this.isPartitionLevelFailoverEnabled = connectionPolicy.EnablePartitionLevelFailover;
             this.globalEndpointManager = globalEndpointManager;
             this.retryOptions = connectionPolicy.RetryOptions;
-            this.partitionKeyRangeLocationCache = partitionKeyRangeLocationCache;
+            this.partitionKeyRangeLocationCache = partitionKeyRangeLocationCache;
+            this.isThinClientEnabled = isThinClientEnabled;
         }
 
         /// <summary>
@@ -41,7 +44,8 @@ public IDocumentClientRetryPolicy GetRequestPolicy()
                 this.partitionKeyRangeLocationCache,
                 this.retryOptions,
                 this.enableEndpointDiscovery,
-                this.isPartitionLevelFailoverEnabled);
+                this.isPartitionLevelFailoverEnabled,
+                this.isThinClientEnabled);
 
             return clientRetryPolicy;
         }

Microsoft.Azure.Cosmos/src/Routing/GlobalPartitionEndpointManagerCore.cs

@@ -53,6 +53,11 @@ internal sealed class GlobalPartitionEndpointManagerCore : GlobalPartitionEndpoi
         /// </summary>
         private readonly bool isPartitionLevelFailoverEnabled;
 
+        /// <summary>
+        /// A readonly boolean flag used to determine if thinclient is enabled.
+        /// </summary>
+        private readonly bool isThinClientEnabled;
+
         /// <summary>
         /// A readonly boolean flag used to determine if partition level circuit breaker is enabled.
         /// </summary>
@@ -93,13 +98,16 @@ internal sealed class GlobalPartitionEndpointManagerCore : GlobalPartitionEndpoi
         /// <param name="globalEndpointManager">An instance of <see cref="GlobalEndpointManager"/>.</param>
         /// <param name="isPartitionLevelFailoverEnabled">A boolean flag indicating if partition level failover is enabled.</param>
         /// <param name="isPartitionLevelCircuitBreakerEnabled">A boolean flag indicating if partition level circuit breaker is enabled.</param>
+        /// <param name="isThinClientEnabled">A boolean flag indicating if thinclient is enabled.</param>
         public GlobalPartitionEndpointManagerCore(
             IGlobalEndpointManager globalEndpointManager,
             bool isPartitionLevelFailoverEnabled = false,
-            bool isPartitionLevelCircuitBreakerEnabled = false)
+            bool isPartitionLevelCircuitBreakerEnabled = false,
+            bool isThinClientEnabled = false)
         {
             this.isPartitionLevelFailoverEnabled = isPartitionLevelFailoverEnabled;
             this.isPartitionLevelCircuitBreakerEnabled = isPartitionLevelCircuitBreakerEnabled;
+            this.isThinClientEnabled = isThinClientEnabled;
             this.globalEndpointManager = globalEndpointManager ?? throw new ArgumentNullException(nameof(globalEndpointManager));
             this.InitializeAndStartCircuitBreakerFailbackBackgroundRefresh();
         }
@@ -169,7 +177,7 @@ public override bool TryMarkEndpointUnavailableForPartitionKeyRange(
             {
                 // For multi master write accounts, since all the regions are treated as write regions, the next locations to fail over
                 // will be the preferred read regions that are configured in the application preferred regions in the CosmosClientOptions.
-                ReadOnlyCollection<Uri> nextLocations = ConfigurationManager.IsThinClientEnabled(defaultValue: false) && ThinClientStoreModel.IsOperationSupportedByThinClient(request)
+                ReadOnlyCollection<Uri> nextLocations = this.isThinClientEnabled && ThinClientStoreModel.IsOperationSupportedByThinClient(request)
                     ? this.globalEndpointManager.ThinClientReadEndpoints
                     : this.globalEndpointManager.ReadEndpoints;
 
@@ -183,7 +191,7 @@ public override bool TryMarkEndpointUnavailableForPartitionKeyRange(
             else if (this.IsRequestEligibleForPerPartitionAutomaticFailover(request))
             {
                 // For any single master write accounts, the next locations to fail over will be the read regions configured at the account level.
-                ReadOnlyCollection<Uri> nextLocations = ConfigurationManager.IsThinClientEnabled(defaultValue: false) && ThinClientStoreModel.IsOperationSupportedByThinClient(request)
+                ReadOnlyCollection<Uri> nextLocations = this.isThinClientEnabled && ThinClientStoreModel.IsOperationSupportedByThinClient(request)
                     ? this.globalEndpointManager.ThinClientReadEndpoints
                     : this.globalEndpointManager.AccountReadEndpoints;
 

Microsoft.Azure.Cosmos/src/ThinClientConstants.cs

@@ -12,7 +12,7 @@ internal static class ThinClientConstants
         public const string RoutedViaProxy = "x-ms-thinclient-route-via-proxy";
         public const string ProxyStartEpk = "x-ms-thinclient-range-min";
         public const string ProxyEndEpk = "x-ms-thinclient-range-max";
-
+        public const string UserAgent = "x-ms-user-agent";
         public const string ProxyOperationType = "x-ms-thinclient-proxy-operation-type";
         public const string ProxyResourceType = "x-ms-thinclient-proxy-resource-type";
         public const string EffectivePartitionKey = "x-ms-effective-partition-key";

Microsoft.Azure.Cosmos/src/ThinClientStoreClient.cs

@@ -5,7 +5,8 @@
 namespace Microsoft.Azure.Cosmos
 {
     using System;
-    using System.Collections.Concurrent;
+    using System.Collections.Concurrent;
+    using System.Diagnostics;
     using System.IO;
     using System.Net.Http;
     using System.Threading;
@@ -24,19 +25,22 @@ internal class ThinClientStoreClient : GatewayStoreClient
     {
         private readonly bool isPartitionLevelFailoverEnabled;
         private readonly ObjectPool<BufferProviderWrapper> bufferProviderWrapperPool;
+        private readonly UserAgentContainer userAgentContainer;
 
         public ThinClientStoreClient(
-            CosmosHttpClient httpClient,
-            ICommunicationEventSource eventSource,
-            JsonSerializerSettings serializerSettings = null,
-            bool isPartitionLevelFailoverEnabled = false)
+            CosmosHttpClient httpClient,
+            UserAgentContainer userAgentContainer,
+            ICommunicationEventSource eventSource,
+            bool isPartitionLevelFailoverEnabled = false,
+            JsonSerializerSettings serializerSettings = null)
             : base(httpClient,
                   eventSource,
                   serializerSettings,
                   isPartitionLevelFailoverEnabled)
         {
-            this.bufferProviderWrapperPool = new ObjectPool<BufferProviderWrapper>(() => new BufferProviderWrapper());
+            this.bufferProviderWrapperPool = new ObjectPool<BufferProviderWrapper>(() => new BufferProviderWrapper());
             this.isPartitionLevelFailoverEnabled = isPartitionLevelFailoverEnabled;
+            this.userAgentContainer = userAgentContainer;
         }
 
         public override async Task<DocumentServiceResponse> InvokeAsync(
@@ -129,7 +133,17 @@ private async ValueTask<HttpRequestMessage> PrepareRequestForProxyAsync(
 
                 requestMessage.Content = new StreamContent(contentStream);
                 requestMessage.Content.Headers.ContentLength = contentStream.Length;
-
+
+                requestMessage.Headers.Clear();
+                requestMessage.Headers.TryAddWithoutValidation(
+                    ThinClientConstants.UserAgent,
+                    this.userAgentContainer.UserAgent);
+
+                Guid activityId = Trace.CorrelationManager.ActivityId;
+                Debug.Assert(activityId != Guid.Empty);
+                requestMessage.Headers.TryAddWithoutValidation(
+                    HttpConstants.HttpHeaders.ActivityId, activityId.ToString());
+
                 requestMessage.RequestUri = thinClientEndpoint;
                 requestMessage.Method = HttpMethod.Post;
 

Microsoft.Azure.Cosmos/src/ThinClientStoreModel.cs

@@ -30,7 +30,8 @@ public ThinClientStoreModel(
             ConsistencyLevel defaultConsistencyLevel,
             DocumentClientEventSource eventSource,
             JsonSerializerSettings serializerSettings,
-            CosmosHttpClient httpClient,
+            CosmosHttpClient httpClient,
+            UserAgentContainer userAgentContainer,
             bool isPartitionLevelFailoverEnabled = false)
             : base(endpointManager,
                   sessionContainer,
@@ -42,10 +43,11 @@ public ThinClientStoreModel(
                   isPartitionLevelFailoverEnabled)
         {
             this.thinClientStoreClient = new ThinClientStoreClient(
-                httpClient,
-                eventSource,
-                serializerSettings,
-                isPartitionLevelFailoverEnabled);
+                httpClient,
+                userAgentContainer,
+                eventSource,
+                isPartitionLevelFailoverEnabled,
+                serializerSettings);
 
             this.isPartitionLevelFailoverEnabled = isPartitionLevelFailoverEnabled;
             this.globalPartitionEndpointManager = globalPartitionEndpointManager;