User agent suffixing for PPAF and PPCB enablement. (Azure#45197)

* Added user agent suffixing if PPAF / PPCB are opted into at client-level.

* Added user agent suffixing if PPAF / PPCB are opted into at client-level.

* Added user agent suffixing if PPAF / PPCB are opted into at client-level.

* Added user agent suffixing if PPAF / PPCB are opted into at client-level.

* Added user agent suffixing if PPAF / PPCB are opted into at client-level.

* Addressing review comments.

* Addressing review comments.

Author: jeet1995

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/UserAgentContainerTest.java

@@ -6,14 +6,33 @@
 import com.azure.cosmos.CosmosAsyncClient;
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class UserAgentContainerTest {
     private final static String SPACE = " ";
     private final static int TIMEOUT = 40000;
 
+    @DataProvider(name = "userAgentSuffixFeatureEnablementFlags")
+    public Object[][] userAgentSuffixFeatureEnablementFlags() {
+        return new Object[][]{
+            {Arrays.asList(UserAgentFeatureFlags.PerPartitionAutomaticFailover, UserAgentFeatureFlags.PerPartitionCircuitBreaker), "|F3"},
+            {Arrays.asList(UserAgentFeatureFlags.PerPartitionCircuitBreaker, UserAgentFeatureFlags.PerPartitionAutomaticFailover), "|F3"},
+            {Arrays.asList(UserAgentFeatureFlags.PerPartitionAutomaticFailover), "|F1"},
+            {Arrays.asList(UserAgentFeatureFlags.PerPartitionCircuitBreaker), "|F2"},
+            {new ArrayList<>(), ""},
+            {null, ""}
+        };
+    }
+
     @Test(groups = {"unit"})
     public void userAgentContainerSetSuffix() {
         String expectedStringFixedPart = getUserAgentFixedPart();
@@ -53,6 +72,25 @@ public void userAgentContainerSetSuffix() {
         }
     }
 
+    @Test(groups = {"unit"}, dataProvider = "userAgentSuffixFeatureEnablementFlags")
+    public void userAgentContainerSetSuffixWithFeatureEnablementFlags(
+        List<UserAgentFeatureFlags> userAgentFeatureFlags,
+        String expectedUserAgentFeatureFlagSuffix) {
+        String expectedStringFixedPart = getUserAgentFixedPart();
+        Set<UserAgentFeatureFlags> uniqueUserAgentFeatureFlags = userAgentFeatureFlags != null ? new HashSet<>(userAgentFeatureFlags) : null;
+
+        //With suffix less than 64 character
+        String userProvidedSuffix = "test-application-id";
+        UserAgentContainer userAgentContainer = new UserAgentContainer();
+        userAgentContainer.setSuffix(userProvidedSuffix);
+        userAgentContainer.setFeatureEnabledFlagsAsSuffix(uniqueUserAgentFeatureFlags);
+
+        String expectedString = expectedStringFixedPart + SPACE + userProvidedSuffix +
+            ((!Strings.isNullOrEmpty(expectedUserAgentFeatureFlagSuffix)) ?
+             expectedUserAgentFeatureFlagSuffix : Strings.Emtpy);
+        assertThat(userAgentContainer.getUserAgent()).isEqualTo(expectedString);
+    }
+
     @Test(groups = {"emulator"}, timeOut = TIMEOUT)
     public void UserAgentIntegration() {
         String userProvidedSuffix = "test-application-id";

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -15,6 +15,7 @@
 
 #### Other Changes
 * Added the `vectorIndexShardKeys` to the vectorIndexSpec for QuantizedFlat and DiskANN vector search. - [PR 44007](https://github.com/Azure/azure-sdk-for-java/pull/44007)
+* Added user agent suffixing if Per-Partition Automatic Failover or Per-Partition Circuit Breaker are enabled at client scope. - [PR 45197](https://github.com/Azure/azure-sdk-for-java/pull/45197)
 
 ### 4.68.0 (2025-03-20)
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java

@@ -111,13 +111,15 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
@@ -564,7 +566,8 @@ private RxDocumentClientImpl(URI serviceEndpoint,
             this.userAgentContainer = new UserAgentContainer();
 
             String userAgentSuffix = this.connectionPolicy.getUserAgentSuffix();
-            if (userAgentSuffix != null && userAgentSuffix.length() > 0) {
+
+            if (userAgentSuffix != null && !userAgentSuffix.isEmpty()) {
                 userAgentContainer.setSuffix(userAgentSuffix);
             }
 
@@ -786,6 +789,8 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func
                     && readConsistencyStrategy != ReadConsistencyStrategy.SESSION
                     && !sessionCapturingOverrideEnabled);
             this.sessionContainer.setDisableSessionCapturing(updatedDisableSessionCapturing);
+
+            this.addUserAgentSuffix(this.userAgentContainer, EnumSet.allOf(UserAgentFeatureFlags.class));
         } catch (Exception e) {
             logger.error("unexpected failure in initializing client.", e);
             close();
@@ -1360,6 +1365,19 @@ private static <T> Flux<FeedResponse<T>> getFeedResponseFluxWithTimeout(
             });
     }
 
+    private void addUserAgentSuffix(UserAgentContainer userAgentContainer, Set<UserAgentFeatureFlags> userAgentFeatureFlags) {
+
+        if (!this.globalPartitionEndpointManagerForPerPartitionAutomaticFailover.isPerPartitionAutomaticFailoverEnabled()) {
+            userAgentFeatureFlags.remove(UserAgentFeatureFlags.PerPartitionAutomaticFailover);
+        }
+
+        if (!this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getCircuitBreakerConfig().isPartitionLevelCircuitBreakerEnabled()) {
+            userAgentFeatureFlags.remove(UserAgentFeatureFlags.PerPartitionCircuitBreaker);
+        }
+
+        userAgentContainer.setFeatureEnabledFlagsAsSuffix(userAgentFeatureFlags);
+    }
+
     @Override
     public Flux<FeedResponse<Database>> queryDatabases(String query, QueryFeedOperationState state) {
         return queryDatabases(new SqlQuerySpec(query), state);

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/UserAgentContainer.java

@@ -3,10 +3,8 @@
 
 package com.azure.cosmos.implementation;
 
-import java.net.IDN;
-import java.net.URLEncoder;
-import java.nio.charset.StandardCharsets;
 import java.text.Normalizer;
+import java.util.Set;
 import java.util.regex.Pattern;
 
 /**
@@ -37,6 +35,20 @@ public String getSuffix() {
         return this.suffix;
     }
 
+    public void setFeatureEnabledFlagsAsSuffix(Set<UserAgentFeatureFlags> userAgentFeatureFlags) {
+        if (userAgentFeatureFlags == null || userAgentFeatureFlags.isEmpty()) {
+            return;
+        }
+
+        int value = 0;
+
+        for (UserAgentFeatureFlags userAgentFeatureFlag : userAgentFeatureFlags) {
+            value += userAgentFeatureFlag.getValue();
+        }
+
+        this.userAgent = this.userAgent + "|F" + value;
+    }
+
     public void setSuffix(String suffix) {
         if (suffix.length() > maxSuffixLength) {
             suffix = suffix.substring(0, maxSuffixLength);

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/UserAgentFeatureFlags.java

@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation;
+
+/**
+ * An enum representing various flags whose opt-in context will appear in the user agent suffix.
+ * <p>
+ * IMPORTANT: When adding a new flag, please ensure the following:
+ * <ul>
+ *     <li>The value of the flag in the enum should follow the pattern FlagName(1 << (n - 1)) where n is the effective ordinal of the flag in the enum.</li>
+ *     <li>The value of the flag should be unique and not overlap with other flags.</li>
+ *     <li>Please keep the value consistent with that of other SDKs - (e.g. see - <a href="https://github.com/Azure/azure-cosmos-dotnet-v3/blob/master/Microsoft.Azure.Cosmos/src/Diagnostics/UserAgentFeatureFlags.cs">UserAgentFeatureFlags.cs</a>) for more details.</li>
+ * </ul>
+ * </p>
+ */
+public enum UserAgentFeatureFlags {
+    PerPartitionAutomaticFailover(1),
+    PerPartitionCircuitBreaker(1 << 1);
+
+    private final int value;
+
+    UserAgentFeatureFlags(int value) {
+        this.value = value;
+    }
+
+    public int getValue() {
+        return this.value;
+    }
+}