Update Statsbeat (#1859)

* Send endpoint and host as part of network statsbeat

* Add unit test for getHost

* Update statsbeat smoke test

* Remove a fixme

* Send network statbeat per url

* Reuse createMetricTelemetry in tests

* Send correct host when redirect happens

* Fix a warning

* Add unit test

* Refactor

* Update unit tests

* Fix spotless

* Add two Statsbeat configs

* Fix spotless

* Rename Statsbeat intervals

* Send attach and azure metadata service on start

* Fix spotless

* Send instrumentation as a FeatureStatsbeat

* Update tests

* Fix spotless

* Fix statsbeat smoke test

* Track non-essential Statsbeat disablement

* Fix smoke test failures

* Fix smoke test failures

* Don't send statsbeat on redirect in tests

* Fix a typo

* Address feedback

* Remove unused method

* Use uppercase enums

* Init telemetry client when disabled is on

* Inject StatsbeatModule to RedirectPolicy

* Use initial delay

* Remove null check and use Mockito.doNothing()

* Move disabled to preview

* Change initiaDelay to 5 seconds and revert appServers

* Address feedback

* Track network counter per ikey

* Update tests

* Handle redirect

* Fix spotlessApply

* Fix nullpointexception at runtime

* Fix spotlessApply

* Null check outside sychronize

* Rename

* Refactor

* Fix a warning methodcanbestatic

* fix test failure

* Fix spotlessApply

* suggestions in progress

* Retrieve endpoint from redirect policy cache

* Fix tests

* Comment out code

* Remove unused methods

* Fix runtime nullpointerexception

* Skip statsbeat ikey in the ikey cache map

* Handle redirect

* Remove null check

* Remove the ikey once it's sent on redirect

* Fix style violations

* Remove arg duplication

* More injection, less global

* Comment

* Remove unneeded conditions

* Remove redirct for now

* Remove StatsbeatModule.get()

* Remove used method

* Remove setinstance

* Refactor

* Add a todo

* Replace regex

* Refactor

Co-authored-by: Trask Stalnaker <[email protected]>

Author: heyams

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/configuration/Configuration.java

@@ -162,12 +162,16 @@ public static class Heartbeat {
   }
 
   public static class Statsbeat {
+    // disabledAll is used internally as an emergency kill-switch to turn off Statsbeat completely
+    // when something goes wrong.
+    public boolean disabledAll = false;
+
     public String instrumentationKey =
         "c4a29126-a7cb-47e5-b348-11414998b11e"; // workspace-aistatsbeat
     public String endpoint =
         DefaultEndpoints.INGESTION_ENDPOINT; // this supports the government cloud
-    public long intervalSeconds = MINUTES.toSeconds(15); // default to 15 minutes
-    public long featureIntervalSeconds = DAYS.toSeconds(1); // default to daily
+    public long shortIntervalSeconds = MINUTES.toSeconds(15); // default to 15 minutes
+    public long longIntervalSeconds = DAYS.toSeconds(1); // default to daily
   }
 
   public static class Proxy {
@@ -201,6 +205,7 @@ public static class PreviewConfiguration {
     public ProfilerConfiguration profiler = new ProfilerConfiguration();
     public GcEventConfiguration gcEvents = new GcEventConfiguration();
     public AadAuthentication authentication = new AadAuthentication();
+    public PreviewStatsbeat statsbeat = new PreviewStatsbeat();
   }
 
   public static class InheritedAttribute {
@@ -279,6 +284,12 @@ public static class PreviewInstrumentation {
         new DisabledByDefaultInstrumentation();
   }
 
+  public static class PreviewStatsbeat {
+    // disabled is used by customer to turn off non-essential Statsbeat, e.g. disk persistence
+    // operation status, optional network statsbeat, other endpoints except Breeze, etc.
+    public boolean disabled = false;
+  }
+
   public static class EnabledByDefaultInstrumentation {
     public boolean enabled = true;
   }

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/exporter/Exporter.java

@@ -35,7 +35,6 @@
 import com.microsoft.applicationinsights.agent.internal.exporter.models.TelemetryExceptionData;
 import com.microsoft.applicationinsights.agent.internal.exporter.models.TelemetryExceptionDetails;
 import com.microsoft.applicationinsights.agent.internal.exporter.models.TelemetryItem;
-import com.microsoft.applicationinsights.agent.internal.statsbeat.StatsbeatModule;
 import com.microsoft.applicationinsights.agent.internal.telemetry.FormattedDuration;
 import com.microsoft.applicationinsights.agent.internal.telemetry.FormattedTime;
 import com.microsoft.applicationinsights.agent.internal.telemetry.TelemetryClient;
@@ -192,7 +191,10 @@ public CompletableResultCode shutdown() {
   private void internalExport(SpanData span) {
     SpanKind kind = span.getKind();
     String instrumentationName = span.getInstrumentationLibraryInfo().getName();
-    StatsbeatModule.get().getNetworkStatsbeat().addInstrumentation(instrumentationName);
+    telemetryClient
+        .getStatsbeatModule()
+        .getInstrumentationStatsbeat()
+        .addInstrumentation(instrumentationName);
     if (kind == SpanKind.INTERNAL) {
       Boolean isLog = span.getAttributes().get(AI_LOG_KEY);
       if (isLog != null && isLog) {

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/httpclient/LazyHttpClient.java

@@ -39,6 +39,7 @@
 import com.azure.identity.VisualStudioCodeCredential;
 import com.azure.identity.VisualStudioCodeCredentialBuilder;
 import com.microsoft.applicationinsights.agent.internal.configuration.Configuration;
+import io.opentelemetry.instrumentation.api.caching.Cache;
 import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.List;
@@ -120,12 +121,13 @@ ProxyOptions.Type.HTTP, new InetSocketAddress(proxyHost, proxyPortNumber)))
     return new NettyAsyncHttpClientBuilder().connectionProvider(connectionProvider).build();
   }
 
+  // pass non-null ikeyRedirectCache if you want to use ikey-specific redirect policy
   public static HttpPipeline newHttpPipeLine(
       @Nullable Configuration.AadAuthentication aadConfiguration,
-      boolean followInstrumentationKeyForRedirect) {
+      @Nullable Cache<String, String> ikeyRedirectCache) {
     List<HttpPipelinePolicy> policies = new ArrayList<>();
-    // Redirect policy to to handle v2.1/track redirects (and other redirects too, e.g. profiler)
-    policies.add(new RedirectPolicy(followInstrumentationKeyForRedirect));
+    // Redirect policy to handle v2.1/track redirects (and other redirects too, e.g. profiler)
+    policies.add(new RedirectPolicy(ikeyRedirectCache));
     if (aadConfiguration != null && aadConfiguration.enabled) {
       policies.add(getAuthenticationPolicy(aadConfiguration));
     }

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/httpclient/RedirectPolicy.java

@@ -36,8 +36,6 @@
 
 // This is mostly a copy from Azure Monitor Open Telemetry Exporter SDK AzureMonitorRedirectPolicy
 public final class RedirectPolicy implements HttpPipelinePolicy {
-  private final boolean followInstrumentationKeyForRedirect;
-  // use this only when followInstrumentationKeyForRedirect is true and instrumentation key is null
   private static final int PERMANENT_REDIRECT_STATUS_CODE = 308;
   private static final int TEMP_REDIRECT_STATUS_CODE = 307;
   // Based on Stamp specific redirects design doc
@@ -47,11 +45,12 @@ public final class RedirectPolicy implements HttpPipelinePolicy {
 
   private final Cache<URL, String> redirectMappings =
       Cache.newBuilder().setMaximumSize(100).build();
-  private final Cache<String, String> instrumentationKeyMappings =
-      Cache.newBuilder().setMaximumSize(100).build();
 
-  public RedirectPolicy(boolean followInstrumentationKeyForRedirect) {
-    this.followInstrumentationKeyForRedirect = followInstrumentationKeyForRedirect;
+  @Nullable private final Cache<String, String> ikeyRedirectCache;
+
+  // pass non-null ikeyRedirectCache if you want to use ikey-specific redirect policy
+  public RedirectPolicy(@Nullable Cache<String, String> ikeyRedirectCache) {
+    this.ikeyRedirectCache = ikeyRedirectCache;
   }
 
   @Override
@@ -92,24 +91,27 @@ private Mono<HttpResponse> attemptRetry(
   }
 
   private void cacheRedirectUrl(String redirectUrl, String instrumentationKey, URL originalUrl) {
-    if (!followInstrumentationKeyForRedirect) {
+    if (ikeyRedirectCache == null) {
       redirectMappings.put(originalUrl, redirectUrl);
       return;
     }
-    if (instrumentationKey != null) {
-      instrumentationKeyMappings.put(instrumentationKey, redirectUrl);
+    if (instrumentationKey == null) {
+      throw new IllegalArgumentException(
+          "instrumentationKey must be non-null when using ikey redirect policy");
     }
+    ikeyRedirectCache.put(instrumentationKey, redirectUrl);
   }
 
   @Nullable
   private String getCachedRedirectUrl(String instrumentationKey, URL originalUrl) {
-    if (!followInstrumentationKeyForRedirect) {
+    if (ikeyRedirectCache == null) {
       return redirectMappings.get(originalUrl);
     }
-    if (instrumentationKey != null) {
-      return instrumentationKeyMappings.get(instrumentationKey);
+    if (instrumentationKey == null) {
+      throw new IllegalArgumentException(
+          "instrumentationKey must be non-null when using ikey redirect policy");
     }
-    return null;
+    return ikeyRedirectCache.get(instrumentationKey);
   }
 
   /**

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/init/AiComponentInstaller.java

@@ -57,6 +57,7 @@
 import com.microsoft.applicationinsights.profiler.config.ServiceProfilerServiceConfig;
 import io.opentelemetry.instrumentation.api.aisdk.AiAppId;
 import io.opentelemetry.instrumentation.api.aisdk.AiLazyConfiguration;
+import io.opentelemetry.instrumentation.api.caching.Cache;
 import io.opentelemetry.instrumentation.api.config.Config;
 import io.opentelemetry.javaagent.extension.AgentListener;
 import io.opentelemetry.sdk.common.CompletableResultCode;
@@ -166,8 +167,16 @@ private void start(Instrumentation instrumentation) {
             .map(MetricFilter::new)
             .collect(Collectors.toList());
 
+    Cache<String, String> ikeyEndpointMap = Cache.newBuilder().setMaximumSize(100).build();
+    StatsbeatModule statsbeatModule = new StatsbeatModule(ikeyEndpointMap);
+    // TODO (heya) apply Builder design pattern to TelemetryClient
     TelemetryClient telemetryClient =
-        new TelemetryClient(config.customDimensions, metricFilters, config.preview.authentication);
+        new TelemetryClient(
+            config.customDimensions,
+            metricFilters,
+            ikeyEndpointMap,
+            statsbeatModule,
+            config.preview.authentication);
     TelemetryClientInitializer.initialize(telemetryClient, config);
     TelemetryClient.setActive(telemetryClient);
 
@@ -211,7 +220,7 @@ private void start(Instrumentation instrumentation) {
     }
 
     // initialize StatsbeatModule
-    StatsbeatModule.get().start(telemetryClient, config);
+    statsbeatModule.start(telemetryClient, config);
   }
 
   private static GcEventMonitor.GcEventMonitorConfiguration formGcEventMonitorConfiguration(

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/profiler/ProfilerServiceInitializer.java

@@ -75,7 +75,7 @@ public static synchronized void initialize(
       GcEventMonitor.GcEventMonitorConfiguration gcEventMonitorConfiguration) {
 
     HttpPipeline httpPipeline =
-        LazyHttpClient.newHttpPipeLine(telemetryClient.getAadAuthentication(), false);
+        LazyHttpClient.newHttpPipeLine(telemetryClient.getAadAuthentication(), null);
 
     initialize(
         appIdSupplier,

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/quickpulse/QuickPulse.java

@@ -69,7 +69,7 @@ private void initializeSync(CountDownLatch latch, TelemetryClient telemetryClien
           initialized = true;
           String quickPulseId = UUID.randomUUID().toString().replace("-", "");
           HttpPipeline httpPipeline =
-              LazyHttpClient.newHttpPipeLine(telemetryClient.getAadAuthentication(), false);
+              LazyHttpClient.newHttpPipeLine(telemetryClient.getAadAuthentication(), null);
           ArrayBlockingQueue<HttpRequest> sendQueue = new ArrayBlockingQueue<>(256, true);
 
           QuickPulseDataSender quickPulseDataSender =

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

@@ -68,7 +68,7 @@ void scheduleWithFixedDelay(long interval) {
     // Querying Azure Metadata Service is required for every 15 mins since VM id will get updated
     // frequently.
     // Starting and restarting a VM will generate a new VM id each time.
-    scheduledExecutor.scheduleWithFixedDelay(this, interval, interval, TimeUnit.SECONDS);
+    scheduledExecutor.scheduleWithFixedDelay(this, 60, interval, TimeUnit.SECONDS);
   }
 
   // only used by tests

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

@@ -48,7 +48,8 @@ enum Feature {
   RABBITMQ_DISABLED(18),
   SPRING_INTEGRATION_DISABLED(19),
   LEGACY_PROPAGATION_DISABLED(20),
-  GRIZZLY_DISABLED(21); // preview instrumentation
+  GRIZZLY_DISABLED(21), // preview instrumentation
+  STATSBEAT_DISABLED(22); // disable non-essential statsbeat
 
   private static final Map<String, Feature> javaVendorFeatureMap;
 

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

@@ -25,18 +25,25 @@
 import com.microsoft.applicationinsights.agent.internal.exporter.models.TelemetryItem;
 import com.microsoft.applicationinsights.agent.internal.telemetry.TelemetryClient;
 import com.microsoft.applicationinsights.agent.internal.telemetry.TelemetryUtil;
-import java.util.HashSet;
+import java.util.Collections;
+import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
-class FeatureStatsbeat extends BaseStatsbeat {
+public class FeatureStatsbeat extends BaseStatsbeat {
 
   private static final String FEATURE_METRIC_NAME = "Feature";
+  private static final String INSTRUMENTATION_METRIC_NAME = "Instrumentation";
 
-  private final Set<Feature> featureList = new HashSet<>(64);
+  private final Set<Feature> featureList = Collections.newSetFromMap(new ConcurrentHashMap<>());
+  private final Set<String> instrumentationList =
+      Collections.newSetFromMap(new ConcurrentHashMap<>());
+  private final FeatureType type;
 
-  FeatureStatsbeat(CustomDimensions customDimensions) {
+  FeatureStatsbeat(CustomDimensions customDimensions, FeatureType type) {
     // track java distribution
     super(customDimensions);
+    this.type = type;
     String javaVendor = System.getProperty("java.vendor");
     featureList.add(Feature.fromJavaVendor(javaVendor));
   }
@@ -46,13 +53,42 @@ long getFeature() {
     return Feature.encode(featureList);
   }
 
+  /**
+   * Returns a long that represents a list of instrumentations. Each bitfield maps to an
+   * instrumentation.
+   */
+  long getInstrumentation() {
+    return Instrumentations.encode(instrumentationList);
+  }
+
+  // this is used by Exporter
+  public void addInstrumentation(String instrumentation) {
+    instrumentationList.add(instrumentation);
+  }
+
   @Override
   protected void send(TelemetryClient telemetryClient) {
-    TelemetryItem statsbeatTelemetry =
-        createStatsbeatTelemetry(telemetryClient, FEATURE_METRIC_NAME, 0);
-    TelemetryUtil.getProperties(statsbeatTelemetry.getData().getBaseData())
-        .put("feature", String.valueOf(getFeature()));
-    telemetryClient.trackStatsbeatAsync(statsbeatTelemetry);
+    String metricName;
+    long encodedLong;
+    String featureType;
+
+    if (type == FeatureType.FEATURE) {
+      metricName = FEATURE_METRIC_NAME;
+      encodedLong = getFeature();
+      featureType = "feature";
+    } else {
+      metricName = INSTRUMENTATION_METRIC_NAME;
+      encodedLong = getInstrumentation();
+      featureType = "instrumentation";
+    }
+
+    TelemetryItem telemetryItem = createStatsbeatTelemetry(telemetryClient, metricName, 0);
+    Map<String, String> properties =
+        TelemetryUtil.getProperties(telemetryItem.getData().getBaseData());
+    properties.put("feature", String.valueOf(encodedLong));
+    properties.put("type", featureType);
+
+    telemetryClient.trackStatsbeatAsync(telemetryItem);
   }
 
   void trackConfigurationOptions(Configuration config) {
@@ -102,5 +138,10 @@ void trackConfigurationOptions(Configuration config) {
     if (!config.preview.instrumentation.springIntegration.enabled) {
       featureList.add(Feature.SPRING_INTEGRATION_DISABLED);
     }
+
+    // Statsbeat
+    if (config.preview.statsbeat.disabled) {
+      featureList.add(Feature.STATSBEAT_DISABLED);
+    }
   }
 }