UserAgent: Fixes race condition in ClientId and max value for UserAgentString

Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs

@@ -67,7 +67,7 @@ public BatchAsyncContainerExecutor(
             this.maxServerRequestBodyLength = maxServerRequestBodyLength;
             this.maxServerRequestOperationCount = maxServerRequestOperationCount;
             this.timerWheel = TimerWheel.CreateTimerWheel(BatchAsyncContainerExecutor.TimerWheelResolution, BatchAsyncContainerExecutor.TimerWheelBucketCount);
-            this.retryOptions = cosmosClientContext.ClientOptions.GetConnectionPolicy().RetryOptions;
+            this.retryOptions = cosmosClientContext.ClientOptions.GetConnectionPolicy(cosmosClientContext.Client.ClientId).RetryOptions;
         }
 
         public virtual async Task<TransactionalBatchOperationResult> AddAsync(

Microsoft.Azure.Cosmos/src/ConnectionPolicy.cs

@@ -39,7 +39,7 @@ public ConnectionPolicy()
             this.ConnectionMode = ConnectionMode.Gateway;
             this.MaxConcurrentFanoutRequests = defaultMaxConcurrentFanoutRequests;
             this.MediaReadMode = MediaReadMode.Buffered;
-            this.UserAgentContainer = new UserAgentContainer();
+            this.UserAgentContainer = new UserAgentContainer(clientId: 0);
             this.preferredLocations = new ObservableCollection<string>();
             this.EnableEndpointDiscovery = true;
             this.MaxConnectionLimit = defaultMaxConcurrentConnectionLimit;

Microsoft.Azure.Cosmos/src/CosmosClient.cs

@@ -250,11 +250,11 @@ internal CosmosClient(
 
             clientOptions ??= new CosmosClientOptions();
 
+            this.ClientId = this.IncrementNumberOfClientsCreated();
             this.ClientContext = ClientContextCore.Create(
                 this,
                 clientOptions);
-
-            this.IncrementNumberOfClientsCreated();
+
             this.ClientConfigurationTraceDatum = new ClientConfigurationTraceDatum(this.ClientContext, DateTime.UtcNow);
         }
 
@@ -479,6 +479,7 @@ internal CosmosClient(
         internal RequestInvokerHandler RequestHandler => this.ClientContext.RequestHandler;
         internal CosmosClientContext ClientContext { get; }
         internal ClientConfigurationTraceDatum ClientConfigurationTraceDatum { get; }
+        internal int ClientId { get; }
 
         /// <summary>
         /// Reads the <see cref="Microsoft.Azure.Cosmos.AccountProperties"/> for the Azure Cosmos DB account.
@@ -1195,9 +1196,9 @@ private Task InitializeContainersAsync(IReadOnlyList<(string databaseId, string
             }
         }
 
-        private void IncrementNumberOfClientsCreated()
+        private int IncrementNumberOfClientsCreated()
         {
-            Interlocked.Increment(ref numberOfClientsCreated);
+            return Interlocked.Increment(ref numberOfClientsCreated);
         }
 
         private async Task InitializeContainerAsync(string databaseId, string containerId, CancellationToken cancellationToken = default)

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -646,7 +646,7 @@ internal CosmosClientOptions Clone()
             return cloneConfiguration;
         }
 
-        internal virtual ConnectionPolicy GetConnectionPolicy()
+        internal virtual ConnectionPolicy GetConnectionPolicy(int clientId)
         {
             this.ValidateDirectTCPSettings();
             this.ValidateLimitToEndpointSettings();
@@ -657,7 +657,7 @@ internal virtual ConnectionPolicy GetConnectionPolicy()
                 RequestTimeout = this.RequestTimeout,
                 ConnectionMode = this.ConnectionMode,
                 ConnectionProtocol = this.ConnectionProtocol,
-                UserAgentContainer = this.CreateUserAgentContainerWithFeatures(),
+                UserAgentContainer = this.CreateUserAgentContainerWithFeatures(clientId),
                 UseMultipleWriteLocations = true,
                 IdleTcpConnectionTimeout = this.IdleTcpConnectionTimeout,
                 OpenTcpConnectionTimeout = this.OpenTcpConnectionTimeout,
@@ -808,7 +808,7 @@ private void ValidateDirectTCPSettings()
             }
         }
 
-        internal UserAgentContainer CreateUserAgentContainerWithFeatures()
+        internal UserAgentContainer CreateUserAgentContainerWithFeatures(int clientId)
         {
             CosmosClientOptionsFeatures features = CosmosClientOptionsFeatures.NoFeatures;
             if (this.AllowBulkExecution)
@@ -830,6 +830,7 @@ internal UserAgentContainer CreateUserAgentContainerWithFeatures()
             string regionConfiguration = this.GetRegionConfiguration();
 
             return new UserAgentContainer(
+                        clientId: clientId,
                         features: featureString,
                         regionConfiguration: regionConfiguration,
                         suffix: this.ApplicationName);

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

@@ -71,7 +71,7 @@ internal static CosmosClientContext Create(
                apitype: clientOptions.ApiType,
                sendingRequestEventArgs: clientOptions.SendingRequestEventArgs,
                transportClientHandlerFactory: clientOptions.TransportClientHandlerFactory,
-               connectionPolicy: clientOptions.GetConnectionPolicy(),
+               connectionPolicy: clientOptions.GetConnectionPolicy(cosmosClient.ClientId),
                enableCpuMonitor: clientOptions.EnableCpuMonitor,
                storeClientFactory: clientOptions.StoreClientFactory,
                desiredConsistencyLevel: clientOptions.GetDocumentsConsistencyLevel(),

Microsoft.Azure.Cosmos/src/UserAgentContainer.cs

@@ -11,14 +11,18 @@ namespace Microsoft.Azure.Cosmos
     internal class UserAgentContainer : Documents.UserAgentContainer
     {
         private const int MaxOperatingSystemString = 30;
+        private const int MaxClientId = 10;
         private readonly string cosmosBaseUserAgent;
+        private readonly string clientId;
 
         public UserAgentContainer(
+            int clientId,
             string features = null,
             string regionConfiguration = "NS",
             string suffix = null) 
                : base()
         {
+            this.clientId = System.Math.Min(clientId, UserAgentContainer.MaxClientId).ToString();
             this.cosmosBaseUserAgent = this.CreateBaseUserAgentString(
                 features: features,
                 regionConfiguration: regionConfiguration);
@@ -30,15 +34,13 @@ public UserAgentContainer(
         protected virtual void GetEnvironmentInformation(
             out string clientVersion,
             out string directVersion,
-            out string clientId,
             out string processArchitecture,
             out string operatingSystem,
             out string runtimeFramework)
         {
             EnvironmentInformation environmentInformation = new EnvironmentInformation();
             clientVersion = environmentInformation.ClientVersion;
             directVersion = environmentInformation.DirectVersion;
-            clientId = CosmosClient.numberOfClientsCreated.ToString();
             processArchitecture = environmentInformation.ProcessArchitecture;
             operatingSystem = environmentInformation.OperatingSystem;
             runtimeFramework = environmentInformation.RuntimeFramework;
@@ -51,7 +53,6 @@ private string CreateBaseUserAgentString(
             this.GetEnvironmentInformation(
                 out string clientVersion,
                 out string directVersion,
-                out string clientId,
                 out string processArchitecture,
                 out string operatingSystem,
                 out string runtimeFramework);
@@ -63,7 +64,7 @@ private string CreateBaseUserAgentString(
 
             // Regex replaces all special characters with empty space except . - | since they do not cause format exception for the user agent string.
             // Do not change the cosmos-netstandard-sdk as it is required for reporting
-            string baseUserAgent = $"cosmos-netstandard-sdk/{clientVersion}" + Regex.Replace($"|{directVersion}|{clientId}|{processArchitecture}|{operatingSystem}|{runtimeFramework}|", @"[^0-9a-zA-Z\.\|\-]+", " ");
+            string baseUserAgent = $"cosmos-netstandard-sdk/{clientVersion}" + Regex.Replace($"|{directVersion}|{this.clientId}|{processArchitecture}|{operatingSystem}|{runtimeFramework}|", @"[^0-9a-zA-Z\.\|\-]+", " ");
 
             if (!string.IsNullOrEmpty(regionConfiguration))
             {

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/HeadersValidationTests.cs

@@ -584,7 +584,7 @@ public void ValidateCustomUserAgentHeader()
             ConnectionPolicy policy = new ConnectionPolicy();
             policy.UserAgentSuffix = suffix;
 
-            string expectedUserAgent = new Cosmos.UserAgentContainer().BaseUserAgent + suffix;
+            string expectedUserAgent = new Cosmos.UserAgentContainer(clientId: 0).BaseUserAgent + suffix;
             string actualUserAgent = policy.UserAgentContainer.UserAgent;
             int startIndexOfClientCounter = expectedUserAgent.IndexOfNth('|', 2) + 1;
 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/UserAgentTests.cs

@@ -36,7 +36,7 @@ public async Task ValidateUserAgentHeaderWithMacOs(bool useMacOs)
 
             using (CosmosClient client = TestCommon.CreateCosmosClient(clientOptions))
             {
-                Cosmos.UserAgentContainer userAgentContainer = client.ClientOptions.GetConnectionPolicy().UserAgentContainer;
+                Cosmos.UserAgentContainer userAgentContainer = client.ClientOptions.GetConnectionPolicy(client.ClientId).UserAgentContainer;
 
                 string userAgentString = userAgentContainer.UserAgent;
                 Assert.IsTrue(userAgentString.Contains(suffix));
@@ -83,14 +83,14 @@ public void ValidateUniqueClientIdHeader()
         public void VerifyUserAgentContent()
         {
             EnvironmentInformation envInfo = new EnvironmentInformation();
-            Cosmos.UserAgentContainer userAgentContainer = new Cosmos.UserAgentContainer();
+            Cosmos.UserAgentContainer userAgentContainer = new Cosmos.UserAgentContainer(clientId: 0);
             string serialization = userAgentContainer.UserAgent;
 
             Assert.IsTrue(serialization.Contains(envInfo.ProcessArchitecture));
             string[] values = serialization.Split('|');
             Assert.AreEqual($"cosmos-netstandard-sdk/{envInfo.ClientVersion}", values[0]);
             Assert.AreEqual(envInfo.DirectVersion, values[1]);
-            Assert.AreEqual(CosmosClient.numberOfClientsCreated.ToString(), values[2]);
+            Assert.AreEqual("0", values[2]);
             Assert.AreEqual(envInfo.ProcessArchitecture, values[3]);
             Assert.AreEqual(envInfo.OperatingSystem, values[4]);
             Assert.AreEqual(envInfo.RuntimeFramework, values[5]);
@@ -173,24 +173,56 @@ await this.ValidateUserAgentStringAsync(
         [TestMethod]
         public void VerifyDefaultUserAgentContainsRegionConfig()
         {
-            UserAgentContainer userAgentContainer = new Cosmos.UserAgentContainer();
+            UserAgentContainer userAgentContainer = new Cosmos.UserAgentContainer(clientId: 0);
             Assert.IsTrue(userAgentContainer.UserAgent.Contains("|NS|"));
         }
 
         [TestMethod]
         public void VerifyClientIDForUserAgentString()
         {
             CosmosClient.numberOfClientsCreated = 0; // reset
-            for (int i = 0; i < 10; i++)
+            const int max = 10;
+            for (int i = 1; i < max + 5; i++)
             {
                 using (CosmosClient client = TestCommon.CreateCosmosClient())
                 {
                     string userAgentString = client.DocumentClient.ConnectionPolicy.UserAgentContainer.UserAgent;
-                    Assert.AreEqual(userAgentString.Split('|')[2], i.ToString());
+                    if (i <= max)
+                    {
+                        Assert.AreEqual(userAgentString.Split('|')[2], i.ToString());
+                    }
+                    else
+                    {
+                        Assert.AreEqual(userAgentString.Split('|')[2], max.ToString());
+                    }
                 }
             }
         }
 
+        [TestMethod]
+        public async Task VerifyClientIDIncrements_Concurrent()
+        {
+            CosmosClient.numberOfClientsCreated = 0; // reset
+            const int max = 10;
+            List<int> expected = new List<int>();
+            for (int i = 1; i < max + 5; i++)
+            {
+                expected.Add(i > max ? max : i);
+            }
+
+            List<Task<CosmosClient>> tasks = new List<Task<CosmosClient>>();
+            for (int i = 1; i < max + 5; i++)
+            {
+                tasks.Add(Task.Factory.StartNew(() => TestCommon.CreateCosmosClient()));
+            }
+
+            await Task.WhenAll(tasks);
+            List<int> actual = tasks.Select(r => 
+                        int.Parse(r.Result.DocumentClient.ConnectionPolicy.UserAgentContainer.UserAgent.Split('|')[2])).ToList();
+            actual.Sort();
+            CollectionAssert.AreEqual(expected, actual);
+        }
+
         private async Task ValidateUserAgentStringAsync(
             CosmosClientOptions cosmosClientOptions,
             string userAgentContentToValidate,
@@ -241,7 +273,7 @@ public async Task VerifyUserAgentWithFeatures(bool setApplicationName, bool useM
 
             using (CosmosClient client = TestCommon.CreateCosmosClient(cosmosClientOptions))
             {
-                Cosmos.UserAgentContainer userAgentContainer = client.ClientOptions.GetConnectionPolicy().UserAgentContainer;
+                Cosmos.UserAgentContainer userAgentContainer = client.ClientOptions.GetConnectionPolicy(client.ClientId).UserAgentContainer;
 
                 string userAgentString = userAgentContainer.UserAgent;
                 if (setApplicationName)
@@ -272,7 +304,7 @@ public async Task VerifyUserAgentWithFeatures(bool setApplicationName, bool useM
 
             using (CosmosClient client = TestCommon.CreateCosmosClient(cosmosClientOptions))
             {
-                Cosmos.UserAgentContainer userAgentContainer = client.ClientOptions.GetConnectionPolicy().UserAgentContainer;
+                Cosmos.UserAgentContainer userAgentContainer = client.ClientOptions.GetConnectionPolicy(client.ClientId).UserAgentContainer;
 
                 string userAgentString = userAgentContainer.UserAgent;
                 if (setApplicationName)
@@ -290,17 +322,17 @@ public async Task VerifyUserAgentWithFeatures(bool setApplicationName, bool useM
 
         private sealed class MacUserAgentStringClientOptions : CosmosClientOptions
         {
-            internal override ConnectionPolicy GetConnectionPolicy()
+            internal override ConnectionPolicy GetConnectionPolicy(int clientId)
             {
-                ConnectionPolicy connectionPolicy = base.GetConnectionPolicy();
-                MacOsUserAgentContainer userAgent = this.CreateUserAgentContainer();
+                ConnectionPolicy connectionPolicy = base.GetConnectionPolicy(clientId);
+                MacOsUserAgentContainer userAgent = this.CreateUserAgentContainer(clientId);
 
                 connectionPolicy.UserAgentContainer = userAgent;
 
                 return connectionPolicy;
             }
 
-            internal MacOsUserAgentContainer CreateUserAgentContainer()
+            internal MacOsUserAgentContainer CreateUserAgentContainer(int clientId)
             {
                 CosmosClientOptionsFeatures features = CosmosClientOptionsFeatures.NoFeatures;
                 if (this.AllowBulkExecution)
@@ -320,17 +352,20 @@ internal MacOsUserAgentContainer CreateUserAgentContainer()
                 }
 
                 return new MacOsUserAgentContainer(
+                            clientId: clientId,
                             features: featureString,
                             suffix: this.ApplicationName);
             }
         }
 
         private sealed class MacOsUserAgentContainer : Cosmos.UserAgentContainer
         {
-            public MacOsUserAgentContainer(string features = null,
+            public MacOsUserAgentContainer(int clientId,
+                                            string features = null,
                                             string regionConfiguration = "N",
                                             string suffix = null)
-                                               : base(features,
+                                               : base(clientId,
+                                                     features,
                                                      regionConfiguration,
                                                      suffix)
             {
@@ -339,7 +374,6 @@ public MacOsUserAgentContainer(string features = null,
             protected override void GetEnvironmentInformation(
                 out string clientVersion,
                 out string directVersion,
-                out string clientId,
                 out string processArchitecture,
                 out string operatingSystem,
                 out string runtimeFramework)
@@ -350,7 +384,6 @@ protected override void GetEnvironmentInformation(
                 base.GetEnvironmentInformation(
                     clientVersion: out clientVersion,
                     directVersion: out directVersion,
-                    clientId: out clientId,
                     processArchitecture: out processArchitecture,
                     operatingSystem: out _,
                     runtimeFramework: out runtimeFramework);
@@ -369,7 +402,7 @@ private CosmosClientOptions SetEnvironmentInformation(bool useMacOs)
 
         private string GetClientIdFromCosmosClient(CosmosClient client)
         {
-            Cosmos.UserAgentContainer userAgentContainer = client.ClientOptions.GetConnectionPolicy().UserAgentContainer;
+            Cosmos.UserAgentContainer userAgentContainer = client.ClientOptions.GetConnectionPolicy(client.ClientId).UserAgentContainer;
             string userAgentString = userAgentContainer.UserAgent;
             string clientId = userAgentString.Split('|')[2];
             return clientId;

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs

@@ -351,6 +351,8 @@ private Mock<CosmosClientContext> MockClientContext()
                .Returns<string, RequestOptions, Func<ITrace, Task<object>>, TraceComponent, TraceLevel>(
                 (operationName, requestOptions, func, comp, level) => func(NoOpTrace.Singleton));
 
+            mockContext.Setup(x => x.Client).Returns(MockCosmosUtil.CreateMockCosmosClient());
+
             return mockContext;
         }