Cosmos Metrics: allow enabling only metrics of certain categories

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java

@@ -19,6 +19,7 @@
 import com.azure.cosmos.implementation.TracerProvider;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetry;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetryMetrics;
+import com.azure.cosmos.implementation.clienttelemetry.MetricCategory;
 import com.azure.cosmos.implementation.clienttelemetry.TagName;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics;
 import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal;
@@ -84,6 +85,7 @@ public final class CosmosAsyncClient implements Closeable {
     private final Tag clientCorrelationTag;
     private final String accountTagValue;
     private final EnumSet<TagName> metricTagNames;
+    private final EnumSet<MetricCategory> metricCategories;
     private final boolean clientMetricsEnabled;
     private final boolean isSendClientTelemetryToServiceEnabled;
     private final MeterRegistry clientMetricRegistrySnapshot;
@@ -139,6 +141,8 @@ public final class CosmosAsyncClient implements Closeable {
             .getClientCorrelationId(effectiveTelemetryConfig);
         this.metricTagNames = telemetryConfigAccessor
             .getMetricTagNames(effectiveTelemetryConfig);
+        this.metricCategories = telemetryConfigAccessor
+            .getMetricCategories(effectiveTelemetryConfig);
 
         List<Permission> permissionList = new ArrayList<>();
         if (this.permissions != null) {
@@ -168,6 +172,7 @@ public final class CosmosAsyncClient implements Closeable {
                                        .withClientTelemetryConfig(this.clientTelemetryConfig)
                                        .withClientCorrelationId(this.clientCorrelationId)
                                        .withMetricTagNames(this.metricTagNames)
+                                       .withMetricCategories(this.metricCategories)
                                        .build();
 
         String effectiveClientCorrelationId = this.asyncDocumentClient.getClientCorrelationId();
@@ -728,6 +733,11 @@ public EnumSet<TagName> getMetricTagNames(CosmosAsyncClient client) {
                     return client.metricTagNames;
                 }
 
+                @Override
+                public EnumSet<MetricCategory> getMetricCategories(CosmosAsyncClient client) {
+                    return client.metricCategories;
+                }
+
                 @Override
                 public boolean isClientTelemetryMetricsEnabled(CosmosAsyncClient client) {
                     return client.clientMetricsEnabled;

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

@@ -11,6 +11,7 @@
 import com.azure.cosmos.implementation.caches.RxClientCollectionCache;
 import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetry;
+import com.azure.cosmos.implementation.clienttelemetry.MetricCategory;
 import com.azure.cosmos.implementation.clienttelemetry.TagName;
 import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo;
 import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal;
@@ -99,6 +100,7 @@ class Builder {
         CosmosClientTelemetryConfig clientTelemetryConfig;
         private String clientCorrelationId = null;
         private EnumSet<TagName> metricTagNames = EnumSet.allOf(TagName.class);
+        private EnumSet<MetricCategory> metricCategories = MetricCategory.DEFAULT_CATEGORIES;
 
         public Builder withServiceEndpoint(String serviceEndpoint) {
             try {
@@ -131,6 +133,12 @@ public Builder withMetricTagNames(EnumSet<TagName> tagNames) {
             return this;
         }
 
+        public Builder withMetricCategories(EnumSet<MetricCategory> categories) {
+            this.metricCategories = categories;
+
+            return this;
+        }
+
         /**
          * New method withMasterKeyOrResourceToken will take either master key or resource token
          * and perform authentication for accessing resource.

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

@@ -19,6 +19,7 @@
 import com.azure.cosmos.ThroughputControlGroupConfig;
 import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import com.azure.cosmos.implementation.batch.PartitionScopeThresholds;
+import com.azure.cosmos.implementation.clienttelemetry.MetricCategory;
 import com.azure.cosmos.implementation.clienttelemetry.TagName;
 import com.azure.cosmos.implementation.patch.PatchOperation;
 import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
@@ -1040,6 +1041,7 @@ public interface CosmosAsyncClientAccessor {
             Tag getClientCorrelationTag(CosmosAsyncClient client);
             String getAccountTagValue(CosmosAsyncClient client);
             EnumSet<TagName> getMetricTagNames(CosmosAsyncClient client);
+            EnumSet<MetricCategory> getMetricCategories(CosmosAsyncClient client);
             boolean isClientTelemetryMetricsEnabled(CosmosAsyncClient client);
             boolean isSendClientTelemetryToServiceEnabled(CosmosAsyncClient client);
             List<String> getPreferredRegions(CosmosAsyncClient client);
@@ -1124,6 +1126,7 @@ public interface CosmosClientTelemetryConfigAccessor {
             int getMaxConnectionPoolSize(CosmosClientTelemetryConfig config);
             Duration getIdleHttpConnectionTimeout(CosmosClientTelemetryConfig config);
             ProxyOptions getProxy(CosmosClientTelemetryConfig config);
+            EnumSet<MetricCategory> getMetricCategories(CosmosClientTelemetryConfig config);
             EnumSet<TagName> getMetricTagNames(CosmosClientTelemetryConfig config);
             String getClientCorrelationId(CosmosClientTelemetryConfig config);
             MeterRegistry getClientMetricRegistry(CosmosClientTelemetryConfig config);

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/clienttelemetry/ClientTelemetry.java

@@ -163,6 +163,10 @@ public ClientTelemetryInfo getClientTelemetryInfo() {
         return clientTelemetryInfo;
     }
 
+    public CosmosClientTelemetryConfig getClientTelemetryConfig() {
+        return clientTelemetryConfig;
+    }
+
     public static String getMachineId(DiagnosticsClientContext.DiagnosticsClientConfig diagnosticsClientConfig) {
         AzureVMMetadata metadataSnapshot = azureVmMetaDataSingleton.get();