Skip to content

Commit 99a587f

Browse files
xinlian12annie-macannie-mac
authored
Allow proxy for client telemetry (Azure#29022)
* update comments * allowClientTelemetryRequestToGoThroughProxy * resolve feedback * clear * update changelog * fix tests * fix spotbug issue * fix tests * resolve comments * resolve comments * add the tests change back * resolve comments Co-authored-by: annie-mac <annie-mac@XBX-2505-B09E.redmond.corp.microsoft.com> Co-authored-by: annie-mac <annie-mac@yindeng2019.fareast.corp.microsoft.com>
1 parent 6bdfe75 commit 99a587f

34 files changed

+906
-340
lines changed

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@
3333
import com.codahale.metrics.jvm.MemoryUsageGaugeSet;
3434
import io.micrometer.core.instrument.MeterRegistry;
3535
import org.apache.commons.lang3.RandomStringUtils;
36-
import org.apache.commons.lang3.StringUtils;
3736
import org.mpierce.metrics.reservoir.hdrhistogram.HdrHistogramResetOnSnapshotReservoir;
3837
import org.slf4j.Logger;
3938
import org.slf4j.LoggerFactory;
@@ -43,7 +42,6 @@
4342
import reactor.core.scheduler.Schedulers;
4443

4544
import java.util.ArrayList;
46-
import java.util.Arrays;
4745
import java.util.HashMap;
4846
import java.util.List;
4947
import java.util.Map;

sdk/cosmos/azure-cosmos-benchmark/src/test/java/com/azure/cosmos/benchmark/Utils.java

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,17 @@
44
package com.azure.cosmos.benchmark;
55

66
import com.azure.cosmos.DirectConnectionConfig;
7-
import com.azure.cosmos.implementation.ConnectionPolicy;
8-
import com.azure.cosmos.models.FeedResponse;
97
import com.azure.cosmos.ThrottlingRetryOptions;
10-
import com.azure.cosmos.models.SqlQuerySpec;
118
import com.azure.cosmos.implementation.AsyncDocumentClient;
9+
import com.azure.cosmos.implementation.ClientTelemetryConfig;
10+
import com.azure.cosmos.implementation.ConnectionPolicy;
1211
import com.azure.cosmos.implementation.Database;
1312
import com.azure.cosmos.implementation.DatabaseForTest;
1413
import com.azure.cosmos.implementation.DocumentCollection;
1514
import com.azure.cosmos.implementation.ResourceResponse;
1615
import com.azure.cosmos.implementation.TestConfigurations;
16+
import com.azure.cosmos.models.FeedResponse;
17+
import com.azure.cosmos.models.SqlQuerySpec;
1718
import reactor.core.publisher.Flux;
1819
import reactor.core.publisher.Mono;
1920

@@ -26,11 +27,13 @@ public static AsyncDocumentClient housekeepingClient() {
2627
options.setMaxRetryAttemptsOnThrottledRequests(100);
2728
options.setMaxRetryWaitTime(Duration.ofSeconds(60));
2829
connectionPolicy.setThrottlingRetryOptions(options);
29-
return new AsyncDocumentClient.Builder().withServiceEndpoint(TestConfigurations.HOST)
30-
.withMasterKeyOrResourceToken(TestConfigurations.MASTER_KEY)
31-
.withConnectionPolicy(connectionPolicy)
32-
.withContentResponseOnWriteEnabled(true)
33-
.build();
30+
return new AsyncDocumentClient.Builder()
31+
.withServiceEndpoint(TestConfigurations.HOST)
32+
.withMasterKeyOrResourceToken(TestConfigurations.MASTER_KEY)
33+
.withConnectionPolicy(connectionPolicy)
34+
.withContentResponseOnWriteEnabled(true)
35+
.withClientTelemetryConfig(ClientTelemetryConfig.getDefaultConfig())
36+
.build();
3437
}
3538

3639
public static String getCollectionLink(Database db, DocumentCollection collection) {

sdk/cosmos/azure-cosmos-benchmark/src/test/java/com/azure/cosmos/benchmark/linkedin/data/ClientTelemetrySampleTest.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ private void enableClientTelemetry() {
3131
String clientTelemetryStagingEndpoint = "https://tools-staging.cosmos.azure.com/api/clienttelemetry/trace/";
3232
String clientTelemetryTestEndpoint = "https://juno-test.documents-dev.windows-int.net/api/clienttelemetry/trace/";
3333
System.setProperty("COSMOS.CLIENT_TELEMETRY_ENDPOINT", clientTelemetryStagingEndpoint);
34+
System.setProperty("COSMOS.CLIENT_TELEMETRY_PROXY_OPTIONS_CONFIG", "{\"type\":\"HTTP\", \"host\": \"localhost\", \"port\": 8080}");
3435

3536
CosmosClient cosmosClient = new CosmosClientBuilder()
3637
.endpoint(TestConfigurations.HOST)

sdk/cosmos/azure-cosmos/CHANGELOG.md

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
#### Other Changes
1212

1313
### 4.30.0 (2022-05-20)
14-
1514
#### Bugs Fixed
1615
* Fixed bubbling of Errors in case of any `java.lang.Error` - See [PR 28620](https://github.com/Azure/azure-sdk-for-java/pull/28620)
1716
* Fixed an issue with creating new Throughput control client item when `enableThroughputControlGroup` is being called multiple times with the same throughput control group. - See [PR 28905](https://github.com/Azure/azure-sdk-for-java/pull/28905)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ public static String getServiceEndpoint(CosmosAsyncClient cosmosAsyncClient) {
9797

9898
@Warning(value = INTERNAL_USE_ONLY_WARNING)
9999
public static boolean isClientTelemetryEnabled(CosmosAsyncClient cosmosAsyncClient) {
100-
return cosmosAsyncClient.isClientTelemetryEnabled();
100+
return cosmosAsyncClient.getClientTelemetryConfig().isClientTelemetryEnabled();
101101
}
102102

103103
@Warning(value = INTERNAL_USE_ONLY_WARNING)

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

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import com.azure.core.util.tracing.Tracer;
1010
import com.azure.cosmos.implementation.ApiType;
1111
import com.azure.cosmos.implementation.AsyncDocumentClient;
12+
import com.azure.cosmos.implementation.ClientTelemetryConfig;
1213
import com.azure.cosmos.implementation.Configs;
1314
import com.azure.cosmos.implementation.ConnectionPolicy;
1415
import com.azure.cosmos.implementation.Database;
@@ -66,7 +67,7 @@ public final class CosmosAsyncClient implements Closeable {
6667
private final TokenCredential tokenCredential;
6768
private final boolean sessionCapturingOverride;
6869
private final boolean enableTransportClientSharing;
69-
private final boolean clientTelemetryEnabled;
70+
private final ClientTelemetryConfig clientTelemetryConfig;
7071
private final TracerProvider tracerProvider;
7172
private final boolean contentResponseOnWriteEnabled;
7273
private static final Tracer TRACER;
@@ -94,7 +95,7 @@ public final class CosmosAsyncClient implements Closeable {
9495
this.tokenCredential = builder.getTokenCredential();
9596
this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled();
9697
this.enableTransportClientSharing = builder.isConnectionSharingAcrossClientsEnabled();
97-
this.clientTelemetryEnabled = builder.isClientTelemetryEnabled();
98+
this.clientTelemetryConfig = builder.getClientTelemetryConfig();
9899
this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled();
99100
this.tracerProvider = new TracerProvider(TRACER);
100101
this.apiType = builder.apiType();
@@ -123,7 +124,8 @@ public final class CosmosAsyncClient implements Closeable {
123124
.withTokenCredential(this.tokenCredential)
124125
.withState(builder.metadataCaches())
125126
.withPermissionFeed(permissionList)
126-
.withApiType(apiType)
127+
.withApiType(this.apiType)
128+
.withClientTelemetryConfig(this.clientTelemetryConfig)
127129
.build();
128130
}
129131

@@ -232,8 +234,13 @@ boolean isContentResponseOnWriteEnabled() {
232234
return contentResponseOnWriteEnabled;
233235
}
234236

235-
boolean isClientTelemetryEnabled() {
236-
return clientTelemetryEnabled;
237+
/***
238+
* Get the client telemetry config.
239+
*
240+
* @return the {@link ClientTelemetryConfig}.
241+
*/
242+
ClientTelemetryConfig getClientTelemetryConfig() {
243+
return this.clientTelemetryConfig;
237244
}
238245

239246
/**

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

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,25 +9,26 @@
99
import com.azure.core.credential.AzureKeyCredential;
1010
import com.azure.core.credential.TokenCredential;
1111
import com.azure.cosmos.implementation.ApiType;
12+
import com.azure.cosmos.implementation.ClientTelemetryConfig;
1213
import com.azure.cosmos.implementation.Configs;
1314
import com.azure.cosmos.implementation.ConnectionPolicy;
14-
import com.azure.cosmos.models.CosmosAuthorizationTokenResolver;
1515
import com.azure.cosmos.implementation.CosmosClientMetadataCachesSnapshot;
1616
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
1717
import com.azure.cosmos.implementation.guava25.base.Preconditions;
1818
import com.azure.cosmos.implementation.routing.LocationHelper;
19+
import com.azure.cosmos.models.CosmosAuthorizationTokenResolver;
1920
import com.azure.cosmos.models.CosmosPermissionProperties;
2021
import com.azure.cosmos.util.Beta;
2122

22-
import static com.azure.cosmos.implementation.ImplementationBridgeHelpers.CosmosClientBuilderHelper;
23-
2423
import java.net.URI;
2524
import java.net.URISyntaxException;
2625
import java.util.Collections;
2726
import java.util.List;
2827
import java.util.Locale;
2928
import java.util.Objects;
3029

30+
import static com.azure.cosmos.implementation.ImplementationBridgeHelpers.CosmosClientBuilderHelper;
31+
3132
/**
3233
* Helper class to build CosmosAsyncClient {@link CosmosAsyncClient} and CosmosClient {@link CosmosClient}
3334
* instances as logical representation of the Azure Cosmos database service.
@@ -116,7 +117,7 @@ public class CosmosClientBuilder implements
116117
private boolean endpointDiscoveryEnabled = true;
117118
private boolean multipleWriteRegionsEnabled = true;
118119
private boolean readRequestsFallbackEnabled = true;
119-
private boolean clientTelemetryEnabled = false;
120+
private final ClientTelemetryConfig clientTelemetryConfig;
120121
private ApiType apiType = null;
121122

122123
/**
@@ -128,6 +129,7 @@ public CosmosClientBuilder() {
128129
// Some default values
129130
this.userAgentSuffix = "";
130131
this.throttlingRetryOptions = new ThrottlingRetryOptions();
132+
this.clientTelemetryConfig = ClientTelemetryConfig.getDefaultConfig();
131133
}
132134

133135
CosmosClientBuilder metadataCaches(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot) {
@@ -640,7 +642,7 @@ public CosmosClientBuilder multipleWriteRegionsEnabled(boolean multipleWriteRegi
640642
* @return current CosmosClientBuilder
641643
*/
642644
public CosmosClientBuilder clientTelemetryEnabled(boolean clientTelemetryEnabled) {
643-
this.clientTelemetryEnabled = clientTelemetryEnabled;
645+
this.clientTelemetryConfig.setClientTelemetryEnabled(clientTelemetryEnabled);
644646
return this;
645647
}
646648

@@ -742,7 +744,7 @@ boolean isMultipleWriteRegionsEnabled() {
742744
* @return flag to enable client telemetry.
743745
*/
744746
boolean isClientTelemetryEnabled() {
745-
return clientTelemetryEnabled;
747+
return this.clientTelemetryConfig.isClientTelemetryEnabled();
746748
}
747749

748750
/**
@@ -761,6 +763,10 @@ boolean isReadRequestsFallbackEnabled() {
761763
return readRequestsFallbackEnabled;
762764
}
763765

766+
ClientTelemetryConfig getClientTelemetryConfig() {
767+
return this.clientTelemetryConfig;
768+
}
769+
764770
/**
765771
* Builds a cosmos async client with the provided properties
766772
*
@@ -803,7 +809,6 @@ private void buildConnectionPolicy() {
803809
this.connectionPolicy.setEndpointDiscoveryEnabled(this.endpointDiscoveryEnabled);
804810
this.connectionPolicy.setMultipleWriteRegionsEnabled(this.multipleWriteRegionsEnabled);
805811
this.connectionPolicy.setReadRequestsFallbackEnabled(this.readRequestsFallbackEnabled);
806-
this.connectionPolicy.setClientTelemetryEnabled(this.clientTelemetryEnabled);
807812
}
808813

809814
private void validateConfig() {

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ class Builder {
9292
boolean contentResponseOnWriteEnabled;
9393
private CosmosClientMetadataCachesSnapshot state;
9494
private ApiType apiType;
95+
ClientTelemetryConfig clientTelemetryConfig;
9596

9697
public Builder withServiceEndpoint(String serviceEndpoint) {
9798
try {
@@ -210,6 +211,18 @@ public Builder withTokenCredential(TokenCredential tokenCredential) {
210211
return this;
211212
}
212213

214+
/***
215+
* Set the client telemetry config.
216+
*
217+
* @param clientTelemetryConfig the {@link ClientTelemetryConfig}.
218+
*
219+
* @return the current builder.
220+
*/
221+
public Builder withClientTelemetryConfig(ClientTelemetryConfig clientTelemetryConfig) {
222+
this.clientTelemetryConfig = clientTelemetryConfig;
223+
return this;
224+
}
225+
213226
private void ifThrowIllegalArgException(boolean value, String error) {
214227
if (value) {
215228
throw new IllegalArgumentException(error);
@@ -240,7 +253,8 @@ public AsyncDocumentClient build() {
240253
transportClientSharing,
241254
contentResponseOnWriteEnabled,
242255
state,
243-
apiType);
256+
apiType,
257+
clientTelemetryConfig);
244258

245259
client.init(state, null);
246260
return client;
Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.azure.cosmos.implementation;
5+
6+
import com.azure.core.http.ProxyOptions;
7+
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
8+
import com.fasterxml.jackson.annotation.JsonProperty;
9+
import com.fasterxml.jackson.core.JsonProcessingException;
10+
import com.fasterxml.jackson.databind.JsonMappingException;
11+
import org.slf4j.Logger;
12+
import org.slf4j.LoggerFactory;
13+
14+
import java.net.InetSocketAddress;
15+
import java.time.Duration;
16+
17+
public class ClientTelemetryConfig {
18+
private static Logger logger = LoggerFactory.getLogger(ClientTelemetryConfig.class);
19+
private static boolean DEFAULT_CLIENT_TELEMETRY_ENABLED = false;
20+
private static final Duration DEFAULT_NETWORK_REQUEST_TIMEOUT = Duration.ofSeconds(60);
21+
private static final Duration DEFAULT_IDLE_CONNECTION_TIMEOUT = Duration.ofSeconds(60);
22+
private static final int DEFAULT_MAX_CONNECTION_POOL_SIZE = 1000;
23+
24+
private boolean clientTelemetryEnabled;
25+
private final Duration httpNetworkRequestTimeout;
26+
private final int maxConnectionPoolSize;
27+
private final Duration idleHttpConnectionTimeout;
28+
private final ProxyOptions proxy;
29+
30+
public ClientTelemetryConfig() {
31+
this.clientTelemetryEnabled = DEFAULT_CLIENT_TELEMETRY_ENABLED;
32+
this.httpNetworkRequestTimeout = DEFAULT_NETWORK_REQUEST_TIMEOUT;
33+
this.maxConnectionPoolSize = DEFAULT_MAX_CONNECTION_POOL_SIZE;
34+
this.idleHttpConnectionTimeout = DEFAULT_IDLE_CONNECTION_TIMEOUT;
35+
this.proxy = this.getProxyOptions();
36+
}
37+
38+
public static ClientTelemetryConfig getDefaultConfig() {
39+
return new ClientTelemetryConfig();
40+
}
41+
42+
public void setClientTelemetryEnabled(boolean clientTelemetryEnabled) {
43+
this.clientTelemetryEnabled = clientTelemetryEnabled;
44+
}
45+
46+
public boolean isClientTelemetryEnabled() {
47+
return this.clientTelemetryEnabled;
48+
}
49+
50+
public Duration getHttpNetworkRequestTimeout() {
51+
return this.httpNetworkRequestTimeout;
52+
}
53+
54+
public int getMaxConnectionPoolSize() {
55+
return this.maxConnectionPoolSize;
56+
}
57+
58+
public Duration getIdleHttpConnectionTimeout() {
59+
return this.idleHttpConnectionTimeout;
60+
}
61+
62+
public ProxyOptions getProxy() {
63+
return this.proxy;
64+
}
65+
66+
private ProxyOptions getProxyOptions() {
67+
String config = Configs.getClientTelemetryProxyOptionsConfig();
68+
69+
if (StringUtils.isNotEmpty(config)) {
70+
try {
71+
JsonProxyOptionsConfig proxyOptionsConfig = Utils.getSimpleObjectMapper().readValue(config, JsonProxyOptionsConfig.class);
72+
ProxyOptions.Type type = ProxyOptions.Type.valueOf(proxyOptionsConfig.type);
73+
74+
if (type != ProxyOptions.Type.HTTP) {
75+
throw new IllegalArgumentException("Only http proxy type is supported.");
76+
}
77+
78+
if (logger.isDebugEnabled()) {
79+
logger.debug(
80+
"Enable proxy with type {}, host {}, port {}",
81+
type,
82+
proxyOptionsConfig.host,
83+
proxyOptionsConfig.port);
84+
}
85+
86+
return new ProxyOptions(type, new InetSocketAddress(proxyOptionsConfig.host, proxyOptionsConfig.port));
87+
} catch (JsonProcessingException e) {
88+
logger.error("Failed to parse client telemetry proxy option config", e);
89+
}
90+
}
91+
92+
return null;
93+
}
94+
95+
private static class JsonProxyOptionsConfig {
96+
@JsonProperty
97+
private String host;
98+
@JsonProperty
99+
private int port;
100+
@JsonProperty
101+
private String type;
102+
103+
private JsonProxyOptionsConfig() {}
104+
private JsonProxyOptionsConfig(String host, int port, String type) {
105+
this.host = host;
106+
this.port = port;
107+
this.type = type;
108+
}
109+
}
110+
}

0 commit comments

Comments
 (0)