Skip to content

Commit

Permalink
UserAgent: Fixes race condition in ClientId and max value for UserAge…
Browse files Browse the repository at this point in the history 
…ntString (#2574)

Issue introduced in #2552

The cap on the ClientId was removed in the above PR, as a result there was a increase in size of the Kusto tables.
The ClientId is also being made thread safe by introducing a clientId ref for every CosmosClient instance
  • Loading branch information
@asketagarwal
asketagarwal authored Jun 25, 2021
1 parent ae7df79 commit fd1b2b6
Show file tree
Hide file tree
Showing 11 changed files with 82 additions and 44 deletions.
2 changes: 1 addition & 1 deletion Microsoft.Azure.Cosmos/src/Batch/BatchAsyncContainerExecutor.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion Microsoft.Azure.Cosmos/src/ConnectionPolicy.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
9 changes: 5 additions & 4 deletions Microsoft.Azure.Cosmos/src/CosmosClient.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
7 changes: 4 additions & 3 deletions Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -646,7 +646,7 @@ internal CosmosClientOptions Clone()
return cloneConfiguration;
}

internal virtual ConnectionPolicy GetConnectionPolicy()
internal virtual ConnectionPolicy GetConnectionPolicy(int clientId)
{
this.ValidateDirectTCPSettings();
this.ValidateLimitToEndpointSettings();
Expand All @@ -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,
Expand Down Expand Up @@ -808,7 +808,7 @@ private void ValidateDirectTCPSettings()
}
}

internal UserAgentContainer CreateUserAgentContainerWithFeatures()
internal UserAgentContainer CreateUserAgentContainerWithFeatures(int clientId)
{
CosmosClientOptionsFeatures features = CosmosClientOptionsFeatures.NoFeatures;
if (this.AllowBulkExecution)
Expand All @@ -830,6 +830,7 @@ internal UserAgentContainer CreateUserAgentContainerWithFeatures()
string regionConfiguration = this.GetRegionConfiguration();

return new UserAgentContainer(
clientId: clientId,
features: featureString,
regionConfiguration: regionConfiguration,
suffix: this.ApplicationName);
Expand Down
2 changes: 1 addition & 1 deletion Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
9 changes: 5 additions & 4 deletions Microsoft.Azure.Cosmos/src/UserAgentContainer.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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;
Expand All @@ -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);
Expand All @@ -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))
{
Expand Down
2 changes: 1 addition & 1 deletion Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/HeadersValidationTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
67 changes: 50 additions & 17 deletions Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/UserAgentTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down Expand Up @@ -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]);
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
{
Expand All @@ -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)
Expand All @@ -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);
Expand All @@ -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;
Expand Down
2 changes: 2 additions & 0 deletions ...Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
Loading

0 comments on commit fd1b2b6

Please sign in to comment.