Integrate managed identity support for benchmarks. (Azure#44857)

* Add managed identity support for account to which benchmark specific workload hits.

* Add managed identity support for account to which benchmark specific workload hits.

* Add managed identity support for account to which benchmark specific workload hits.

* Add managed identity support for account to which benchmark specific workload hits.

* Add managed identity support for account to which benchmark specific workload hits.

* Add managed identity support for account to which benchmark specific workload hits.

Author: jeet1995

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsynReadWithMultipleClients.java

@@ -197,8 +197,7 @@ private void createClients() {
                         .preferredRegions(this.configuration.getPreferredRegionsList())
                         .consistencyLevel(configuration.getConsistencyLevel())
                         .connectionSharingAcrossClientsEnabled(true)
-                        .contentResponseOnWriteEnabled(configuration.isContentResponseOnWriteEnabled())
-                        .clientTelemetryEnabled(configuration.isClientTelemetryEnabled());
+                        .contentResponseOnWriteEnabled(configuration.isContentResponseOnWriteEnabled());
 
                     if (this.configuration.getConnectionMode().equals(ConnectionMode.DIRECT)) {
                         cosmosClientBuilder = cosmosClientBuilder.directMode(DirectConnectionConfig.getDefaultConfig());

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncBenchmark.java

@@ -31,6 +31,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Function;
 
 public class Configuration {
     public static final String SUCCESS_COUNTER_METER_NAME = "#Successful Operations";
@@ -146,6 +147,9 @@ public class Configuration {
     @Parameter(names = "isPartitionLevelCircuitBreakerEnabled", description = "A flag to denote whether partition level circuit breaker is enabled.")
     private String isPartitionLevelCircuitBreakerEnabled = String.valueOf(true);
 
+    @Parameter(names = "-isManagedIdentityRequired", description = "A flag to denote whether benchmark-specific CosmosClient instance should use Managed Identity to authenticate.")
+    private String isManagedIdentityRequired = String.valueOf(false);
+
     @Parameter(names = "-operation", description = "Type of Workload:\n"
         + "\tReadThroughput- run a READ workload that prints only throughput *\n"
         + "\tReadThroughputWithMultipleClients - run a READ workload that prints throughput and latency for multiple client read.*\n"
@@ -179,6 +183,10 @@ public class Configuration {
     @Parameter(names = "-numberOfOperations", description = "Total NUMBER Of Documents To Insert")
     private int numberOfOperations = 100000;
 
+    public Boolean isManagedIdentityRequired() {
+        return Boolean.parseBoolean(this.isManagedIdentityRequired);
+    }
+
     static class DurationConverter implements IStringConverter<Duration> {
         @Override
         public Duration convert(String value) {
@@ -229,15 +237,6 @@ public Duration convert(String value) {
     @Parameter(names = "-accountNameInGraphiteReporter", description = "if set, account name with be appended in graphite reporter")
     private boolean accountNameInGraphiteReporter = false;
 
-    @Parameter(names = "-clientTelemetryEnabled", description = "Switch to enable client telemetry")
-    private String clientTelemetryEnabled = String.valueOf(false);
-
-    @Parameter(names = "-clientTelemetrySchedulingInSeconds", description = "Client telemetry scheduling intervals in seconds")
-    private int clientTelemetrySchedulingInSeconds = 10 * 60;
-
-    @Parameter(names = "-clientTelemetryEndpoint", description = "Client Telemetry Juno endpoint")
-    private String clientTelemetryEndpoint;
-
     @Parameter(names = "-pointLatencyThresholdMs", description = "Latency threshold for point operations")
     private int pointLatencyThresholdMs = -1;
 
@@ -588,22 +587,10 @@ public boolean isEncryptionEnabled() {
         return encryptionEnabled;
     }
 
-    public boolean isClientTelemetryEnabled() {
-        return Boolean.parseBoolean(clientTelemetryEnabled);
-    }
-
     public boolean isDefaultLog4jLoggerEnabled() {
         return Boolean.parseBoolean(defaultLog4jLoggerEnabled);
     }
 
-    public String getClientTelemetryEndpoint() {
-        return clientTelemetryEndpoint;
-    }
-
-    public int getClientTelemetrySchedulingInSeconds() {
-        return clientTelemetrySchedulingInSeconds;
-    }
-
     public Integer getTupleSize() {
         return tupleSize;
     }
@@ -851,4 +838,55 @@ public String[] tagsAsPrefix() {
 
         return this.graphiteMeterRegistry;
     }
+
+    public static String getAadLoginUri() {
+        return getOptionalConfigProperty(
+                "AAD_LOGIN_ENDPOINT",
+                "https://login.microsoftonline.com/",
+                v -> v);
+    }
+
+    public static String getAadManagedIdentityId() {
+        return getOptionalConfigProperty("AAD_MANAGED_IDENTITY_ID", null, v -> v);
+    }
+
+    public static String getAadTenantId() {
+        return getOptionalConfigProperty("AAD_TENANT_ID", null, v -> v);
+    }
+
+    private static <T> T getOptionalConfigProperty(String name, T defaultValue, Function<String, T> conversion) {
+        String textValue = getConfigPropertyOrNull(name);
+
+        if (textValue == null) {
+            return defaultValue;
+        }
+
+        T returnValue = conversion.apply(textValue);
+        return returnValue != null ? returnValue : defaultValue;
+    }
+
+    private static String getConfigPropertyOrNull(String name) {
+        String systemPropertyName = "COSMOS." + name;
+        String environmentVariableName = "COSMOS_" + name;
+        String fromSystemProperty = emptyToNull(System.getProperty(systemPropertyName));
+        if (fromSystemProperty != null) {
+            return fromSystemProperty;
+        }
+
+        return emptyToNull(System.getenv().get(environmentVariableName));
+    }
+
+    /**
+     * Returns the given string if it is nonempty; {@code null} otherwise.
+     *
+     * @param string the string to test and possibly return
+     * @return {@code string} itself if it is nonempty; {@code null} if it is empty or null
+     */
+    private static String emptyToNull(String string) {
+        if (string == null || string.isEmpty()) {
+            return null;
+        }
+
+        return string;
+    }
 }

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/Configuration.java

@@ -38,11 +38,6 @@ public static void main(String[] args) throws Exception {
                 return;
             }
 
-            if (cfg.isClientTelemetryEnabled()) {
-                System.setProperty("COSMOS.CLIENT_TELEMETRY_ENDPOINT", cfg.getClientTelemetryEndpoint());
-                System.setProperty("COSMOS.CLIENT_TELEMETRY_SCHEDULING_IN_SECONDS", String.valueOf(cfg.getClientTelemetrySchedulingInSeconds()));
-            }
-
             validateConfiguration(cfg);
 
             if (cfg.isSync()) {

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/Main.java

@@ -59,7 +59,7 @@ class ReadMyWriteWorkflow extends AsyncBenchmark<Document> {
     @Override
     protected void init() {
         // TODO: move read my writes to use v4 APIs
-        this.client =  CosmosBridgeInternal.getAsyncDocumentClient(cosmosClient);
+        this.client =  CosmosBridgeInternal.getAsyncDocumentClient(benchmarkWorkloadClient);
         Database database = DocDBUtils.getDatabase(client, configuration.getDatabaseId());
         this.collection = DocDBUtils.getCollection(client, database.getSelfLink(), configuration.getCollectionId());
         this.nameCollectionLink = String.format("dbs/%s/colls/%s", database.getId(), collection.getId());
@@ -282,7 +282,7 @@ private Flux<Document> xPartitionQuery(SqlQuerySpec query) {
         options.setMaxDegreeOfParallelism(-1);
 
         QueryFeedOperationState state = new QueryFeedOperationState(
-            cosmosClient,
+                benchmarkWorkloadClient,
             "xPartitionQuery",
             configuration.getDatabaseId(),
             configuration.getCollectionId(),
@@ -313,7 +313,7 @@ private Flux<Document> singlePartitionQuery(Document d) {
                                                                    d.getString(QUERY_FIELD_NAME)));
 
         QueryFeedOperationState state = new QueryFeedOperationState(
-            cosmosClient,
+                benchmarkWorkloadClient,
             "singlePartitionQuery",
             configuration.getDatabaseId(),
             configuration.getCollectionId(),