Skip to content

Commit 7e90a14

Browse files
heyamsheyams
authored
Turn off Statsbeat when AMPLS is used (#1994)
* shutdown statsbeat when sending statsbeat request and server returns UnknownHostException * Fix failing tests * Fix tests * Remove * Address comments
1 parent 8172def commit 7e90a14

File tree

7 files changed

+92
-41
lines changed

7 files changed

+92
-41
lines changed

‎agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/statsbeat/AzureMetadataService.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,11 @@ void scheduleWithFixedDelay(long interval) {
7171
scheduledExecutor.scheduleWithFixedDelay(this, 60, interval, TimeUnit.SECONDS);
7272
}
7373

74+
void shutdown() {
75+
logger.info("Shutting down Azure Metadata Service.");
76+
scheduledExecutor.shutdown();
77+
}
78+
7479
// only used by tests
7580
void updateMetadata(String response) throws IOException {
7681
updateMetadata(mapper.readValue(response, MetadataInstanceResponse.class));

‎agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/statsbeat/StatsbeatModule.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,12 @@ public class StatsbeatModule {
4141
ThreadPoolUtils.createDaemonThreadFactory(BaseStatsbeat.class));
4242

4343
private final CustomDimensions customDimensions;
44-
4544
private final NetworkStatsbeat networkStatsbeat;
4645
private final AttachStatsbeat attachStatsbeat;
4746
private final FeatureStatsbeat featureStatsbeat;
4847
private final FeatureStatsbeat instrumentationStatsbeat;
4948
private final NonessentialStatsbeat nonessentialStatsbeat;
49+
private final AzureMetadataService azureMetadataService;
5050

5151
private final AtomicBoolean started = new AtomicBoolean();
5252

@@ -57,6 +57,7 @@ public StatsbeatModule(Cache<String, String> ikeyEndpointMap) {
5757
featureStatsbeat = new FeatureStatsbeat(customDimensions, FeatureType.FEATURE);
5858
instrumentationStatsbeat = new FeatureStatsbeat(customDimensions, FeatureType.INSTRUMENTATION);
5959
nonessentialStatsbeat = new NonessentialStatsbeat(customDimensions);
60+
azureMetadataService = new AzureMetadataService(attachStatsbeat, customDimensions);
6061
}
6162

6263
public void start(TelemetryClient telemetryClient, Configuration config) {
@@ -99,9 +100,7 @@ public void start(TelemetryClient telemetryClient, Configuration config) {
99100
// only turn on AzureMetadataService when the resource provider is VM or UNKNOWN.
100101
if (rp == ResourceProvider.RP_VM || rp == ResourceProvider.UNKNOWN) {
101102
// will only reach here the first time, after instance has been instantiated
102-
AzureMetadataService metadataService =
103-
new AzureMetadataService(attachStatsbeat, customDimensions);
104-
metadataService.scheduleWithFixedDelay(longIntervalSeconds);
103+
azureMetadataService.scheduleWithFixedDelay(longIntervalSeconds);
105104
}
106105

107106
featureStatsbeat.trackConfigurationOptions(config);
@@ -117,6 +116,12 @@ public void start(TelemetryClient telemetryClient, Configuration config) {
117116
}
118117
}
119118

119+
public void shutdown() {
120+
logger.debug("Shutting down Statsbeat scheduler.");
121+
scheduledExecutor.shutdown();
122+
azureMetadataService.shutdown();
123+
}
124+
120125
public NetworkStatsbeat getNetworkStatsbeat() {
121126
return networkStatsbeat;
122127
}

‎agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/telemetry/TelemetryChannel.java‎

Lines changed: 48 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -40,12 +40,13 @@
4040
import com.microsoft.applicationinsights.agent.internal.httpclient.LazyHttpClient;
4141
import com.microsoft.applicationinsights.agent.internal.httpclient.RedirectPolicy;
4242
import com.microsoft.applicationinsights.agent.internal.localstorage.LocalFileWriter;
43-
import com.microsoft.applicationinsights.agent.internal.statsbeat.NetworkStatsbeat;
43+
import com.microsoft.applicationinsights.agent.internal.statsbeat.StatsbeatModule;
4444
import io.opentelemetry.instrumentation.api.cache.Cache;
4545
import io.opentelemetry.sdk.common.CompletableResultCode;
4646
import java.io.IOException;
4747
import java.io.StringWriter;
4848
import java.net.URL;
49+
import java.net.UnknownHostException;
4950
import java.nio.ByteBuffer;
5051
import java.util.ArrayList;
5152
import java.util.HashMap;
@@ -91,17 +92,19 @@ private static ObjectMapper createObjectMapper() {
9192
private final HttpPipeline pipeline;
9293
private final URL endpointUrl;
9394
@Nullable private final LocalFileWriter localFileWriter;
94-
// this is null for the statsbeat channel
95-
@Nullable private final NetworkStatsbeat networkStatsbeat;
95+
private final StatsbeatModule statsbeatModule;
96+
private final boolean isStatsbeat;
9697

9798
public static TelemetryChannel create(
9899
URL endpointUrl,
99100
LocalFileWriter localFileWriter,
100101
Cache<String, String> ikeyEndpointMap,
101-
@Nullable NetworkStatsbeat networkStatsbeat,
102+
StatsbeatModule statsbeatModule,
103+
boolean isStatsbeat,
102104
@Nullable Configuration.AadAuthentication aadAuthentication) {
103105
HttpPipeline httpPipeline = LazyHttpClient.newHttpPipeLine(aadAuthentication, ikeyEndpointMap);
104-
return new TelemetryChannel(httpPipeline, endpointUrl, localFileWriter, networkStatsbeat);
106+
return new TelemetryChannel(
107+
httpPipeline, endpointUrl, localFileWriter, statsbeatModule, isStatsbeat);
105108
}
106109

107110
public CompletableResultCode sendRawBytes(ByteBuffer buffer, String instrumentationKey) {
@@ -113,11 +116,13 @@ public TelemetryChannel(
113116
HttpPipeline pipeline,
114117
URL endpointUrl,
115118
LocalFileWriter localFileWriter,
116-
@Nullable NetworkStatsbeat networkStatsbeat) {
119+
StatsbeatModule statsbeatModule,
120+
boolean isStatsbeat) {
117121
this.pipeline = pipeline;
118122
this.endpointUrl = endpointUrl;
119123
this.localFileWriter = localFileWriter;
120-
this.networkStatsbeat = networkStatsbeat;
124+
this.statsbeatModule = statsbeatModule;
125+
this.isStatsbeat = isStatsbeat;
121126
}
122127

123128
public CompletableResultCode send(List<TelemetryItem> telemetryItems) {
@@ -233,13 +238,16 @@ private CompletableResultCode internalSend(
233238
parseResponseCode(
234239
response.getStatusCode(), instrumentationKey, byteBuffers, persisted);
235240
LazyHttpClient.consumeResponseBody(response);
236-
// networkStatsbeat is null when it's sending a Statsbeat request.
237-
if (networkStatsbeat != null) {
241+
if (!isStatsbeat) {
238242
if (response.getStatusCode() == 200) {
239-
networkStatsbeat.incrementRequestSuccessCount(
240-
System.currentTimeMillis() - startTime, instrumentationKey);
243+
statsbeatModule
244+
.getNetworkStatsbeat()
245+
.incrementRequestSuccessCount(
246+
System.currentTimeMillis() - startTime, instrumentationKey);
241247
} else {
242-
networkStatsbeat.incrementRequestFailureCount(instrumentationKey);
248+
statsbeatModule
249+
.getNetworkStatsbeat()
250+
.incrementRequestFailureCount(instrumentationKey);
243251
}
244252
}
245253
if (!persisted) {
@@ -253,21 +261,32 @@ private CompletableResultCode internalSend(
253261
}
254262
},
255263
error -> {
256-
if (!NetworkFriendlyExceptions.logSpecialOneTimeFriendlyException(
257-
error, endpointUrl.toString(), friendlyExceptionThrown, logger)) {
258-
operationLogger.recordFailure(
259-
"Error sending telemetry items: " + error.getMessage(), error);
260-
}
264+
// AMPLS
265+
if (isStatsbeat && error instanceof UnknownHostException) {
266+
// when sending a Statsbeat request and server returns an UnknownHostException, it's
267+
// likely that
268+
// it's using a virtual network. In that case, we use the kill-switch to turn off
269+
// Statsbeat.
270+
statsbeatModule.shutdown();
271+
} else {
272+
if (!NetworkFriendlyExceptions.logSpecialOneTimeFriendlyException(
273+
error, endpointUrl.toString(), friendlyExceptionThrown, logger)) {
274+
operationLogger.recordFailure(
275+
"Error sending telemetry items: " + error.getMessage(), error);
276+
}
261277

262-
// networkStatsbeat is null when it's sending a Statsbeat request.
263-
if (networkStatsbeat != null) {
264-
networkStatsbeat.incrementRequestFailureCount(instrumentationKey);
265-
}
266-
// no need to write to disk again when failing to send raw bytes from the persisted
267-
// file
268-
if (!persisted) {
269-
writeToDiskOnFailure(byteBuffers, instrumentationKey);
278+
if (!isStatsbeat) {
279+
statsbeatModule
280+
.getNetworkStatsbeat()
281+
.incrementRequestFailureCount(instrumentationKey);
282+
}
283+
// no need to write to disk again when failing to send raw bytes from the persisted
284+
// file
285+
if (!persisted) {
286+
writeToDiskOnFailure(byteBuffers, instrumentationKey);
287+
}
270288
}
289+
271290
if (!persisted) {
272291
// persisted byte buffers don't come from the pool so shouldn't go back to the pool
273292
byteBufferPool.offer(byteBuffers);
@@ -303,10 +322,9 @@ private void parseResponseCode(
303322
case 439: // Breeze-specific: THROTTLED OVER EXTENDED TIME
304323
// TODO handle throttling
305324
// TODO (heya) track throttling count via Statsbeat
306-
// networkStatsbeat is null when it's sending a Statsbeat request.
307325
// instrumentationKey is null when sending persisted file's raw bytes.
308-
if (networkStatsbeat != null) {
309-
networkStatsbeat.incrementThrottlingCount(instrumentationKey);
326+
if (!isStatsbeat) {
327+
statsbeatModule.getNetworkStatsbeat().incrementThrottlingCount(instrumentationKey);
310328
}
311329
break;
312330
case 200: // SUCCESS
@@ -318,10 +336,9 @@ private void parseResponseCode(
318336
case 0: // client-side exception
319337
// TODO exponential backoff and retry to a limit
320338
// TODO (heya) track failure count via Statsbeat
321-
// networkStatsbeat is null when it's sending a Statsbeat request.
322339
// instrumentationKey is null when sending persisted file's raw bytes.
323-
if (networkStatsbeat != null) {
324-
networkStatsbeat.incrementRetryCount(instrumentationKey);
340+
if (!isStatsbeat) {
341+
statsbeatModule.getNetworkStatsbeat().incrementRetryCount(instrumentationKey);
325342
}
326343
break;
327344
default:

‎agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/telemetry/TelemetryClient.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,8 @@ public BatchSpanProcessor getChannelBatcher() {
221221
endpointProvider.getIngestionEndpointUrl(),
222222
localFileWriter,
223223
ikeyEndpointMap,
224-
statsbeatModule.getNetworkStatsbeat(),
224+
statsbeatModule,
225+
false,
225226
aadAuthentication);
226227

227228
if (!readOnlyFileSystem) {
@@ -254,7 +255,8 @@ public BatchSpanProcessor getStatsbeatChannelBatcher() {
254255
endpointProvider.getStatsbeatEndpointUrl(),
255256
localFileWriter,
256257
ikeyEndpointMap,
257-
null,
258+
statsbeatModule,
259+
true,
258260
null);
259261

260262
if (!readOnlyFileSystem) {

‎agent/agent-tooling/src/test/java/com/microsoft/applicationinsights/agent/internal/localstorage/IntegrationTests.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,8 @@ public void setup() throws Exception {
100100
pipelineBuilder.build(),
101101
new URL("http://foo.bar"),
102102
new LocalFileWriter(localFileCache, tempFolder, null),
103-
null);
103+
null,
104+
false);
104105
}
105106

106107
@Test

‎agent/agent-tooling/src/test/java/com/microsoft/applicationinsights/agent/internal/localstorage/LocalFileLoaderTests.java‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@
3737
import com.fasterxml.jackson.databind.JsonNode;
3838
import com.fasterxml.jackson.databind.ObjectMapper;
3939
import com.microsoft.applicationinsights.agent.internal.MockHttpResponse;
40+
import com.microsoft.applicationinsights.agent.internal.statsbeat.NetworkStatsbeat;
41+
import com.microsoft.applicationinsights.agent.internal.statsbeat.StatsbeatModule;
4042
import com.microsoft.applicationinsights.agent.internal.telemetry.TelemetryChannel;
4143
import io.opentelemetry.sdk.common.CompletableResultCode;
4244
import java.io.ByteArrayInputStream;
@@ -54,6 +56,7 @@
5456
import org.junit.jupiter.api.AfterEach;
5557
import org.junit.jupiter.api.Test;
5658
import org.junit.jupiter.api.io.TempDir;
59+
import org.mockito.Mockito;
5760
import reactor.core.publisher.Mono;
5861

5962
public class LocalFileLoaderTests {
@@ -246,9 +249,16 @@ public void testDeleteFilePermanentlyOnSuccess() throws Exception {
246249
LocalFileWriter localFileWriter = new LocalFileWriter(localFileCache, tempFolder, null);
247250
LocalFileLoader localFileLoader = new LocalFileLoader(localFileCache, tempFolder, null);
248251

252+
StatsbeatModule mockedStatsbeatModule = Mockito.mock(StatsbeatModule.class);
253+
when(mockedStatsbeatModule.getNetworkStatsbeat())
254+
.thenReturn(Mockito.mock(NetworkStatsbeat.class));
249255
TelemetryChannel telemetryChannel =
250256
new TelemetryChannel(
251-
pipelineBuilder.build(), new URL("http://foo.bar"), localFileWriter, null);
257+
pipelineBuilder.build(),
258+
new URL("http://foo.bar"),
259+
localFileWriter,
260+
mockedStatsbeatModule,
261+
false);
252262

253263
// persist 10 files to disk
254264
for (int i = 0; i < 10; i++) {
@@ -298,7 +308,11 @@ public void testDeleteFilePermanentlyOnFailure() throws Exception {
298308

299309
TelemetryChannel telemetryChannel =
300310
new TelemetryChannel(
301-
pipelineBuilder.build(), new URL("http://foo.bar"), localFileWriter, null);
311+
pipelineBuilder.build(),
312+
new URL("http://foo.bar"),
313+
localFileWriter,
314+
Mockito.mock(StatsbeatModule.class),
315+
false);
302316

303317
// persist 10 files to disk
304318
for (int i = 0; i < 10; i++) {

‎agent/agent-tooling/src/test/java/com/microsoft/applicationinsights/agent/internal/telemetry/TelemetryChannelTest.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
package com.microsoft.applicationinsights.agent.internal.telemetry;
2323

2424
import static org.assertj.core.api.Assertions.assertThat;
25+
import static org.mockito.Mockito.when;
2526

2627
import com.azure.core.http.HttpClient;
2728
import com.azure.core.http.HttpHeaders;
@@ -36,6 +37,8 @@
3637
import com.microsoft.applicationinsights.agent.internal.httpclient.RedirectPolicy;
3738
import com.microsoft.applicationinsights.agent.internal.localstorage.LocalFileCache;
3839
import com.microsoft.applicationinsights.agent.internal.localstorage.LocalFileWriter;
40+
import com.microsoft.applicationinsights.agent.internal.statsbeat.NetworkStatsbeat;
41+
import com.microsoft.applicationinsights.agent.internal.statsbeat.StatsbeatModule;
3942
import io.opentelemetry.instrumentation.api.cache.Cache;
4043
import io.opentelemetry.sdk.common.CompletableResultCode;
4144
import java.io.ByteArrayInputStream;
@@ -57,6 +60,7 @@
5760
import org.junit.jupiter.api.BeforeEach;
5861
import org.junit.jupiter.api.Test;
5962
import org.junit.jupiter.api.io.TempDir;
63+
import org.mockito.Mockito;
6064
import reactor.core.publisher.Flux;
6165
import reactor.core.publisher.Mono;
6266

@@ -78,11 +82,14 @@ private TelemetryChannel getTelemetryChannel() throws MalformedURLException {
7882
.policies(policies.toArray(new HttpPipelinePolicy[0]))
7983
.httpClient(recordingHttpClient);
8084
LocalFileCache localFileCache = new LocalFileCache(tempFolder);
85+
StatsbeatModule mockedStatsModule = Mockito.mock(StatsbeatModule.class);
86+
when(mockedStatsModule.getNetworkStatsbeat()).thenReturn(Mockito.mock(NetworkStatsbeat.class));
8187
return new TelemetryChannel(
8288
pipelineBuilder.build(),
8389
new URL(END_POINT_URL),
8490
new LocalFileWriter(localFileCache, tempFolder, null),
85-
null);
91+
mockedStatsModule,
92+
false);
8693
}
8794

8895
@Nullable

0 commit comments

Comments
 (0)