Allow proxy for client telemetry

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/ctl/AsyncCtlWorkload.java

@@ -33,7 +33,6 @@
 import com.codahale.metrics.jvm.MemoryUsageGaugeSet;
 import io.micrometer.core.instrument.MeterRegistry;
 import org.apache.commons.lang3.RandomStringUtils;
-import org.apache.commons.lang3.StringUtils;
 import org.mpierce.metrics.reservoir.hdrhistogram.HdrHistogramResetOnSnapshotReservoir;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -43,7 +42,6 @@
 import reactor.core.scheduler.Schedulers;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;

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

@@ -97,7 +97,7 @@ public static String getServiceEndpoint(CosmosAsyncClient cosmosAsyncClient) {
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
     public static boolean isClientTelemetryEnabled(CosmosAsyncClient cosmosAsyncClient) {
-        return cosmosAsyncClient.isClientTelemetryEnabled();
+        return cosmosAsyncClient.getClientTelemetryConfig().isClientTelemetryEnabled();
     }
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)

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

@@ -0,0 +1,145 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos;
+
+import com.azure.core.http.ProxyOptions;
+
+import java.time.Duration;
+
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument;
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
+
+/**
+ * Represents the client telemetry config.
+ */
+public class ClientTelemetryConnectionConfig {
+    private static final Duration MIN_NETWORK_REQUEST_TIMEOUT = Duration.ofSeconds(60);
+    private static final Duration DEFAULT_NETWORK_REQUEST_TIMEOUT = Duration.ofSeconds(60);
+    private static final Duration DEFAULT_IDLE_CONNECTION_TIMEOUT = Duration.ofSeconds(60);
+    private static final int DEFAULT_MAX_CONNECTION_POOL_SIZE = 1000;
+
+    private Duration networkRequestTimeout;
+    private int maxConnectionPoolSize;
+    private Duration idleConnectionTimeout;
+    private ProxyOptions proxy;
+
+    public ClientTelemetryConnectionConfig() {
+        this.idleConnectionTimeout = DEFAULT_IDLE_CONNECTION_TIMEOUT;
+        this.maxConnectionPoolSize = DEFAULT_MAX_CONNECTION_POOL_SIZE;
+        this.networkRequestTimeout = DEFAULT_NETWORK_REQUEST_TIMEOUT;
+    }
+
+    public static ClientTelemetryConnectionConfig getDefaultConfig() {
+        return new ClientTelemetryConnectionConfig();
+    }
+
+    /**
+     * Gets the network request timeout interval (time to wait for response from network peer).
+     * The default is 60 seconds.
+     *
+     * @return the network request timeout duration.
+     */
+    public Duration getNetworkRequestTimeout() {
+        return this.networkRequestTimeout;
+    }
+
+    /**
+     * Sets the network request timeout interval (time to wait for response from network peer).
+     * The default is 60 seconds.
+     *
+     * @param networkRequestTimeout the network request timeout duration.
+     * @return the {@link ClientTelemetryConnectionConfig}.
+     */
+    ClientTelemetryConnectionConfig setNetworkRequestTimeout(Duration networkRequestTimeout) {
+        checkNotNull(networkRequestTimeout, "NetworkRequestTimeout can not be null");
+        checkArgument(networkRequestTimeout.toMillis() >= MIN_NETWORK_REQUEST_TIMEOUT.toMillis(),
+                "NetworkRequestTimeout can not be less than %s millis", MIN_NETWORK_REQUEST_TIMEOUT.toMillis());
+        this.networkRequestTimeout = networkRequestTimeout;
+        return this;
+    }
+
+    /**
+     * Gets the value of the connection pool size the client is using.
+     *
+     * @return connection pool size.
+     */
+    public int getMaxConnectionPoolSize() {
+        return this.maxConnectionPoolSize;
+    }
+
+    /**
+     * Sets the value of the connection pool size, the default
+     * is 1000.
+     *
+     * @param maxConnectionPoolSize The value of the connection pool size.
+     * @return the {@link ClientTelemetryConnectionConfig}.
+     */
+    ClientTelemetryConnectionConfig setMaxConnectionPoolSize(int maxConnectionPoolSize) {
+        this.maxConnectionPoolSize = maxConnectionPoolSize;
+        return this;
+    }
+
+    /**
+     * Gets the value of the timeout for an idle connection, the default is 60
+     * seconds.
+     *
+     * @return Idle connection timeout duration.
+     */
+    public Duration getIdleConnectionTimeout() {
+        return this.idleConnectionTimeout;
+    }
+
+    /**
+     * sets the value of the timeout for an idle connection. After that time,
+     * the connection will be automatically closed.
+     *
+     * @param idleConnectionTimeout the duration for an idle connection.
+     * @return the {@link ClientTelemetryConnectionConfig}.
+     */
+    ClientTelemetryConnectionConfig setIdleConnectionTimeout(Duration idleConnectionTimeout) {
+        this.idleConnectionTimeout = idleConnectionTimeout;
+        return this;
+    }
+
+    /**
+     * Gets the proxy options which contain the InetSocketAddress of proxy server.
+     *
+     * @return the proxy options.
+     */
+    public ProxyOptions getProxy() {
+        return this.proxy;
+    }
+
+    /**
+     * Sets the proxy options.
+     *
+     * Currently, only support Http proxy type with just the routing address. Username and password will be ignored.
+     *
+     * @param proxy The proxy options.
+     * @return the {@link ClientTelemetryConnectionConfig}.
+     */
+
+    public ClientTelemetryConnectionConfig setProxy(ProxyOptions proxy) {
+        if (proxy.getType() != ProxyOptions.Type.HTTP) {
+            throw new IllegalArgumentException("Only http proxy type is supported.");
+        }
+
+        this.proxy = proxy;
+        return this;
+    }
+
+    @Override
+    public String toString() {
+        String proxyType = proxy != null ? proxy.getType().toString() : null;
+        String proxyAddress = proxy != null ? proxy.getAddress().toString() : null;
+
+        return "ClientTelemetryConfig{" +
+                ", maxConnectionPoolSize=" + maxConnectionPoolSize +
+                ", idleConnectionTimeout=" + idleConnectionTimeout +
+                ", networkRequestTimeout=" + networkRequestTimeout +
+                ", proxyType=" + proxyType +
+                ", inetSocketProxyAddress=" + proxyAddress +
+                '}';
+    }
+}

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

@@ -9,6 +9,7 @@
 import com.azure.core.util.tracing.Tracer;
 import com.azure.cosmos.implementation.ApiType;
 import com.azure.cosmos.implementation.AsyncDocumentClient;
+import com.azure.cosmos.implementation.ClientTelemetryConfig;
 import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.ConnectionPolicy;
 import com.azure.cosmos.implementation.Database;
@@ -66,7 +67,7 @@ public final class CosmosAsyncClient implements Closeable {
     private final TokenCredential tokenCredential;
     private final boolean sessionCapturingOverride;
     private final boolean enableTransportClientSharing;
-    private final boolean clientTelemetryEnabled;
+    private final ClientTelemetryConfig clientTelemetryConfig;
     private final TracerProvider tracerProvider;
     private final boolean contentResponseOnWriteEnabled;
     private static final Tracer TRACER;
@@ -94,7 +95,7 @@ public final class CosmosAsyncClient implements Closeable {
         this.tokenCredential = builder.getTokenCredential();
         this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled();
         this.enableTransportClientSharing = builder.isConnectionSharingAcrossClientsEnabled();
-        this.clientTelemetryEnabled = builder.isClientTelemetryEnabled();
+        this.clientTelemetryConfig = builder.getClientTelemetryConfig();
         this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled();
         this.tracerProvider = new TracerProvider(TRACER);
         this.apiType = builder.apiType();
@@ -123,7 +124,8 @@ public final class CosmosAsyncClient implements Closeable {
                                        .withTokenCredential(this.tokenCredential)
                                        .withState(builder.metadataCaches())
                                        .withPermissionFeed(permissionList)
-                                       .withApiType(apiType)
+                                       .withApiType(this.apiType)
+                                       .withClientTelemetryConfig(this.clientTelemetryConfig)
                                        .build();
     }
 
@@ -232,8 +234,13 @@ boolean isContentResponseOnWriteEnabled() {
         return contentResponseOnWriteEnabled;
     }
 
-    boolean isClientTelemetryEnabled() {
-        return clientTelemetryEnabled;
+    /***
+     * Get the client telemetry config.
+     *
+     * @return the {@link ClientTelemetryConfig}.
+     */
+    ClientTelemetryConfig getClientTelemetryConfig() {
+        return this.clientTelemetryConfig;
     }
 
     /**

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

@@ -9,6 +9,7 @@
 import com.azure.core.credential.AzureKeyCredential;
 import com.azure.core.credential.TokenCredential;
 import com.azure.cosmos.implementation.ApiType;
+import com.azure.cosmos.implementation.ClientTelemetryConfig;
 import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.ConnectionPolicy;
 import com.azure.cosmos.models.CosmosAuthorizationTokenResolver;
@@ -20,6 +21,7 @@
 import com.azure.cosmos.util.Beta;
 
 import static com.azure.cosmos.implementation.ImplementationBridgeHelpers.CosmosClientBuilderHelper;
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
 
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -117,6 +119,8 @@ public class CosmosClientBuilder implements
     private boolean multipleWriteRegionsEnabled = true;
     private boolean readRequestsFallbackEnabled = true;
     private boolean clientTelemetryEnabled = false;
+    private ClientTelemetryConnectionConfig clientTelemetryConnectionConfig;
+    private ClientTelemetryConfig clientTelemetryConfig;
     private ApiType apiType = null;
 
     /**
@@ -128,6 +132,7 @@ public CosmosClientBuilder() {
         //  Some default values
         this.userAgentSuffix = "";
         this.throttlingRetryOptions = new ThrottlingRetryOptions();
+        this.clientTelemetryConnectionConfig = ClientTelemetryConnectionConfig.getDefaultConfig();
     }
 
     CosmosClientBuilder metadataCaches(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot) {
@@ -663,6 +668,19 @@ public CosmosClientBuilder readRequestsFallbackEnabled(boolean readRequestsFallb
         return this;
     }
 
+    /***
+     * Set the client telemetry connection config.
+     *
+     * @param clientTelemetryConnectionConfig the {@link ClientTelemetryConnectionConfig}.
+     * @return the current CosmosClientBuilder.
+     */
+    public CosmosClientBuilder clientTelemetryConnectionConfig(ClientTelemetryConnectionConfig clientTelemetryConnectionConfig) {
+        checkNotNull(clientTelemetryConnectionConfig, "Argument 'clientTelemetryConnectionConfig' can not be null");
+        this.clientTelemetryConnectionConfig = clientTelemetryConnectionConfig;
+
+        return this;
+    }
+
     /**
      * Gets the GATEWAY connection configuration to be used.
      *
@@ -742,7 +760,7 @@ boolean isMultipleWriteRegionsEnabled() {
      * @return flag to enable client telemetry.
      */
     boolean isClientTelemetryEnabled() {
-        return clientTelemetryEnabled;
+        return this.clientTelemetryEnabled;
     }
 
     /**
@@ -761,6 +779,10 @@ boolean isReadRequestsFallbackEnabled() {
         return readRequestsFallbackEnabled;
     }
 
+    ClientTelemetryConfig getClientTelemetryConfig() {
+        return this.clientTelemetryConfig;
+    }
+
     /**
      * Builds a cosmos async client with the provided properties
      *
@@ -770,6 +792,7 @@ public CosmosAsyncClient buildAsyncClient() {
 
         validateConfig();
         buildConnectionPolicy();
+        buildClientTelemetryConfig();
         return new CosmosAsyncClient(this);
     }
 
@@ -782,6 +805,7 @@ public CosmosClient buildClient() {
 
         validateConfig();
         buildConnectionPolicy();
+        buildClientTelemetryConfig();
         return new CosmosClient(this);
     }
 
@@ -803,7 +827,16 @@ private void buildConnectionPolicy() {
         this.connectionPolicy.setEndpointDiscoveryEnabled(this.endpointDiscoveryEnabled);
         this.connectionPolicy.setMultipleWriteRegionsEnabled(this.multipleWriteRegionsEnabled);
         this.connectionPolicy.setReadRequestsFallbackEnabled(this.readRequestsFallbackEnabled);
-        this.connectionPolicy.setClientTelemetryEnabled(this.clientTelemetryEnabled);
+    }
+
+    private void buildClientTelemetryConfig() {
+        // There are two ways customer can enable the client telemetry:
+        // 1. Enable through CosmosClientBuilder
+        // 2. Enabled by system property
+
+        this.clientTelemetryConfig = new ClientTelemetryConfig(
+                Configs.isClientTelemetryEnabled(this.clientTelemetryEnabled),
+                this.clientTelemetryConnectionConfig);
     }
 
     private void validateConfig() {

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

@@ -92,6 +92,7 @@ class Builder {
         boolean contentResponseOnWriteEnabled;
         private CosmosClientMetadataCachesSnapshot state;
         private ApiType apiType;
+        ClientTelemetryConfig clientTelemetryConfig;
 
         public Builder withServiceEndpoint(String serviceEndpoint) {
             try {
@@ -210,6 +211,18 @@ public Builder withTokenCredential(TokenCredential tokenCredential) {
             return this;
         }
 
+        /***
+         * Set the client telemetry config.
+         *
+         * @param clientTelemetryConfig the {@link ClientTelemetryConfig}.
+         *
+         * @return the current builder.
+         */
+        public Builder withClientTelemetryConfig(ClientTelemetryConfig clientTelemetryConfig) {
+            this.clientTelemetryConfig = clientTelemetryConfig;
+            return this;
+        }
+
         private void ifThrowIllegalArgException(boolean value, String error) {
             if (value) {
                 throw new IllegalArgumentException(error);
@@ -240,7 +253,8 @@ public AsyncDocumentClient build() {
                 transportClientSharing,
                 contentResponseOnWriteEnabled,
                 state,
-                apiType);
+                apiType,
+                clientTelemetryConfig);
 
             client.init(state, null);
             return client;