Rate limited sampling

Author: trask

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

@@ -74,7 +74,7 @@ public void validate() {
     preview.validate();
   }
 
-  // TODO (trask) investigate options for mapping lowercase values to otel enum directly
+  @Deprecated
   public enum SpanKind {
     @JsonProperty("server")
     SERVER(io.opentelemetry.api.trace.SpanKind.SERVER),
@@ -94,6 +94,17 @@ public enum SpanKind {
     }
   }
 
+  public enum TelemetryKind {
+    @JsonProperty("request")
+    REQUEST,
+    @JsonProperty("dependency")
+    DEPENDENCY,
+    @JsonProperty("trace")
+    TRACE,
+    @JsonProperty("exception")
+    EXCEPTION
+  }
+
   public enum MatchType {
     @JsonProperty("strict")
     STRICT,
@@ -151,7 +162,11 @@ public static class Role {
 
   public static class Sampling {
 
-    public float percentage = 100;
+    // fixed percentage of requests
+    @Nullable public Double percentage;
+
+    // default is 5 requests per second
+    @Nullable public Double limitPerSecond;
   }
 
   public static class SamplingPreview {
@@ -175,6 +190,8 @@ public static class SamplingPreview {
     // Another (lesser) reason is because .NET SDK always propagates trace flags "00" (not
     // sampled)
     //
+    // future goal: make parentBased sampling the default if item count is received via tracestate
+    //
     // IMPORTANT if changing this default, we need to keep it at least on Azure Functions
     public boolean parentBased;
 
@@ -352,7 +369,7 @@ public static class PreviewConfiguration {
         new HashSet<>(asList("b3", "b3multi"));
 
     public void validate() {
-      for (Configuration.SamplingOverride samplingOverride : sampling.overrides) {
+      for (SamplingOverride samplingOverride : sampling.overrides) {
         samplingOverride.validate();
       }
       for (Configuration.InstrumentationKeyOverride instrumentationKeyOverride :
@@ -578,22 +595,34 @@ private static boolean isRuntimeAttached() {
   }
 
   public static class SamplingOverride {
-    // TODO (trask) consider making this required when moving out of preview
-    @Nullable public SpanKind spanKind;
+    @Deprecated @Nullable public SpanKind spanKind;
+
+    // TODO (trask) make this required when moving out of preview
+    //   for now the default is both "request" and "dependency" for backwards compatibility
+    @Nullable public TelemetryKind telemetryKind;
+
+    // TODO (trask) should this be named "standalone" (and meaning flipped)
+    // TODO (trask) need to add test for this
+    public boolean inRequest = true;
+
     // not using include/exclude, because you can still get exclude with this by adding a second
     // (exclude) override above it
     // (since only the first matching override is used)
     public List<SamplingOverrideAttribute> attributes = new ArrayList<>();
-    public Float percentage;
+    public Double percentage;
     public String id; // optional, used for debugging purposes only
 
+    public boolean isForRequestTelemetry() {
+      return telemetryKind == TelemetryKind.REQUEST
+          || (telemetryKind == null && spanKind != SpanKind.CLIENT);
+    }
+
+    public boolean isForDependencyTelemetry() {
+      return telemetryKind == TelemetryKind.DEPENDENCY
+          || (telemetryKind == null && spanKind != SpanKind.SERVER);
+    }
+
     public void validate() {
-      if (spanKind == null && attributes.isEmpty()) {
-        // TODO add doc and go link, similar to telemetry processors
-        throw new FriendlyException(
-            "A sampling override configuration is missing \"spanKind\" and has no attributes.",
-            "Please provide at least one of \"spanKind\" or \"attributes\" for the sampling override configuration.");
-      }
       if (percentage == null) {
         // TODO add doc and go link, similar to telemetry processors
         throw new FriendlyException(

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

@@ -83,6 +83,9 @@ public class ConfigurationBuilder {
   private static final String APPLICATIONINSIGHTS_SAMPLING_PERCENTAGE =
       "APPLICATIONINSIGHTS_SAMPLING_PERCENTAGE";
 
+  private static final String APPLICATIONINSIGHTS_SAMPLING_LIMIT_PER_SECOND =
+      "APPLICATIONINSIGHTS_SAMPLING_LIMIT_PER_SECOND";
+
   private static final String APPLICATIONINSIGHTS_INSTRUMENTATION_LOGGING_LEVEL =
       "APPLICATIONINSIGHTS_INSTRUMENTATION_LOGGING_LEVEL";
 
@@ -181,6 +184,15 @@ public static Configuration create(Path agentJarPath, @Nullable RpConfiguration
               + " and it is now enabled by default,"
               + " so no need to enable it under preview configuration");
     }
+    for (SamplingOverride override : config.preview.sampling.overrides) {
+      if (override.spanKind != null) {
+        configurationLogger.warn(
+            "Sampling overrides \"spanKind\" has been deprecated,"
+                + " and support for it will be removed in a future release, please transition from"
+                + " \"spanKind\" to \"telemetryKind\".");
+      }
+    }
+
     logWarningIfUsingInternalAttributes(config);
 
     overlayFromEnv(config, agentJarPath.getParent());
@@ -194,6 +206,10 @@ public static Configuration create(Path agentJarPath, @Nullable RpConfiguration
       overlayFromEnv(rpConfiguration);
       overlayRpConfiguration(config, rpConfiguration);
     }
+    // only fall back to default sampling configuration after all overlays have been performed
+    if (config.sampling.limitPerSecond == null && config.sampling.percentage == null) {
+      config.sampling.limitPerSecond = 5.0;
+    }
     // only set role instance to host name as a last resort
     if (config.role.instance == null) {
       String hostname = HostName.get();
@@ -457,6 +473,10 @@ static void overlayFromEnv(Configuration config, Path baseDir) throws IOExceptio
     config.sampling.percentage =
         overlayWithEnvVar(APPLICATIONINSIGHTS_SAMPLING_PERCENTAGE, config.sampling.percentage);
 
+    config.sampling.limitPerSecond =
+        overlayWithEnvVar(
+            APPLICATIONINSIGHTS_SAMPLING_LIMIT_PER_SECOND, config.sampling.limitPerSecond);
+
     config.proxy = overlayProxyFromEnv(config.proxy);
 
     config.selfDiagnostics.level =
@@ -466,10 +486,9 @@ static void overlayFromEnv(Configuration config, Path baseDir) throws IOExceptio
             APPLICATIONINSIGHTS_SELF_DIAGNOSTICS_FILE_PATH, config.selfDiagnostics.file.path);
 
     config.preview.metricIntervalSeconds =
-        (int)
-            overlayWithEnvVar(
-                APPLICATIONINSIGHTS_PREVIEW_METRIC_INTERVAL_SECONDS,
-                config.preview.metricIntervalSeconds);
+        overlayWithEnvVar(
+            APPLICATIONINSIGHTS_PREVIEW_METRIC_INTERVAL_SECONDS,
+            config.preview.metricIntervalSeconds);
 
     config.preview.instrumentation.springIntegration.enabled =
         overlayWithEnvVar(
@@ -576,6 +595,7 @@ static void overlayRpConfiguration(Configuration config, RpConfiguration rpConfi
     }
     if (rpConfiguration.sampling != null) {
       config.sampling.percentage = rpConfiguration.sampling.percentage;
+      config.sampling.limitPerSecond = rpConfiguration.sampling.limitPerSecond;
     }
     if (isTrimEmpty(config.role.name)) {
       // only use rp configuration role name as a fallback, similar to WEBSITE_SITE_NAME
@@ -626,13 +646,25 @@ public static String overlayWithEnvVar(String name, String defaultValue) {
     return defaultValue;
   }
 
-  static float overlayWithEnvVar(String name, float defaultValue) {
+  @Nullable
+  static Double overlayWithEnvVar(String name, @Nullable Double defaultValue) {
+    String value = getEnvVar(name);
+    if (value != null) {
+      configurationLogger.debug("using environment variable: {}", name);
+      // intentionally allowing NumberFormatException to bubble up as invalid configuration and
+      // prevent agent from starting
+      return Double.parseDouble(value);
+    }
+    return defaultValue;
+  }
+
+  static int overlayWithEnvVar(String name, int defaultValue) {
     String value = getEnvVar(name);
     if (value != null) {
       configurationLogger.debug("using environment variable: {}", name);
       // intentionally allowing NumberFormatException to bubble up as invalid configuration and
       // prevent agent from starting
-      return Float.parseFloat(value);
+      return Integer.parseInt(value);
     }
     return defaultValue;
   }
@@ -798,17 +830,21 @@ public static Configuration loadJsonConfigFile(Path configPath) throws IOExcepti
   }
 
   // this is for external callers, where logging is ok
-  public static float roundToNearest(float samplingPercentage) {
+  public static double roundToNearest(double samplingPercentage) {
     return roundToNearest(samplingPercentage, false);
   }
 
-  // visible for testing
-  private static float roundToNearest(float samplingPercentage, boolean doNotLogWarnMessages) {
+  @Nullable
+  private static Double roundToNearest(
+      @Nullable Double samplingPercentage, boolean doNotLogWarnMessages) {
+    if (samplingPercentage == null) {
+      return null;
+    }
     if (samplingPercentage == 0) {
-      return 0;
+      return 0.0;
     }
     double itemCount = 100 / samplingPercentage;
-    float rounded = 100.0f / Math.round(itemCount);
+    double rounded = 100.0 / Math.round(itemCount);
 
     if (Math.abs(samplingPercentage - rounded) >= 1) {
       // TODO include link to docs in this warning message

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

@@ -28,15 +28,22 @@
 import com.azure.monitor.opentelemetry.exporter.implementation.logging.OperationLogger;
 import com.azure.monitor.opentelemetry.exporter.implementation.models.TelemetryItem;
 import com.azure.monitor.opentelemetry.exporter.implementation.quickpulse.QuickPulse;
+import com.microsoft.applicationinsights.agent.internal.configuration.Configuration.SamplingOverride;
+import com.microsoft.applicationinsights.agent.internal.sampling.AiSampler;
+import com.microsoft.applicationinsights.agent.internal.sampling.SamplingOverrides;
 import com.microsoft.applicationinsights.agent.internal.telemetry.BatchItemProcessor;
 import com.microsoft.applicationinsights.agent.internal.telemetry.TelemetryClient;
 import com.microsoft.applicationinsights.agent.internal.telemetry.TelemetryObservers;
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.sdk.common.CompletableResultCode;
 import io.opentelemetry.sdk.logs.data.LogData;
 import io.opentelemetry.sdk.logs.data.Severity;
 import io.opentelemetry.sdk.logs.export.LogExporter;
+import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Consumer;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
@@ -52,15 +59,21 @@ public class AgentLogExporter implements LogExporter {
   // TODO (trask) could implement this in a filtering LogExporter instead
   private volatile Severity threshold;
 
+  private final SamplingOverrides logSamplingOverrides;
+  private final SamplingOverrides exceptionSamplingOverrides;
   private final LogDataMapper mapper;
   private final Consumer<TelemetryItem> telemetryItemConsumer;
 
   public AgentLogExporter(
       Severity threshold,
+      List<SamplingOverride> logSamplingOverrides,
+      List<SamplingOverride> exceptionSamplingOverrides,
       LogDataMapper mapper,
       @Nullable QuickPulse quickPulse,
       BatchItemProcessor batchItemProcessor) {
     this.threshold = threshold;
+    this.logSamplingOverrides = new SamplingOverrides(logSamplingOverrides);
+    this.exceptionSamplingOverrides = new SamplingOverrides(exceptionSamplingOverrides);
     this.mapper = mapper;
     telemetryItemConsumer =
         telemetryItem -> {
@@ -86,18 +99,44 @@ public CompletableResultCode export(Collection<LogData> logs) {
       return CompletableResultCode.ofFailure();
     }
     for (LogData log : logs) {
-      SpanContext spanContext = log.getSpanContext();
-      if (spanContext.isValid() && !spanContext.getTraceFlags().isSampled()) {
-        continue;
-      }
       logger.debug("exporting log: {}", log);
       try {
         int severity = log.getSeverity().getSeverityNumber();
         int threshold = this.threshold.getSeverityNumber();
         if (severity < threshold) {
           continue;
         }
-        mapper.map(log, telemetryItemConsumer);
+
+        String stack = log.getAttributes().get(SemanticAttributes.EXCEPTION_STACKTRACE);
+
+        SamplingOverrides samplingOverrides =
+            stack != null ? exceptionSamplingOverrides : logSamplingOverrides;
+
+        Double samplingPercentage = samplingOverrides.getOverridePercentage(log.getAttributes());
+
+        SpanContext spanContext = log.getSpanContext();
+
+        if (samplingPercentage != null && !shouldSample(spanContext, samplingPercentage)) {
+          continue;
+        }
+
+        if (samplingPercentage == null
+            && spanContext.isValid()
+            && !spanContext.getTraceFlags().isSampled()) {
+          // if there is no sampling override, and the log is part of an unsampled trace, then don't
+          // capture it
+          continue;
+        }
+
+        Long itemCount = null;
+        if (samplingPercentage != null) {
+          // samplingPercentage cannot be 0 here
+          itemCount = Math.round(100.0 / samplingPercentage);
+        }
+
+        TelemetryItem telemetryItem = mapper.map(log, stack, itemCount);
+        telemetryItemConsumer.accept(telemetryItem);
+
         exportingLogLogger.recordSuccess();
       } catch (Throwable t) {
         exportingLogLogger.recordFailure(t.getMessage(), t, EXPORTER_MAPPING_ERROR);
@@ -116,4 +155,22 @@ public CompletableResultCode flush() {
   public CompletableResultCode shutdown() {
     return CompletableResultCode.ofSuccess();
   }
+
+  @SuppressFBWarnings(
+      value = "SECPR", // Predictable pseudorandom number generator
+      justification = "Predictable random is ok for sampling decision")
+  private static boolean shouldSample(SpanContext spanContext, double percentage) {
+    if (percentage == 100) {
+      // optimization, no need to calculate score
+      return true;
+    }
+    if (percentage == 0) {
+      // optimization, no need to calculate score
+      return false;
+    }
+    if (spanContext.isValid()) {
+      return AiSampler.shouldRecordAndSample(spanContext.getTraceId(), percentage);
+    }
+    return ThreadLocalRandom.current().nextDouble() < percentage / 100;
+  }
 }

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

@@ -36,6 +36,7 @@
 import java.nio.file.Files;
 import java.nio.file.attribute.BasicFileAttributes;
 import java.nio.file.attribute.FileTime;
+import java.util.Objects;
 import java.util.concurrent.Executors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -100,24 +101,28 @@ public void run() {
         if (!newRpConfiguration.connectionString.equals(rpConfiguration.connectionString)) {
           logger.debug(
               "Connection string from the JSON config file is overriding the previously configured connection string.");
+          configuration.connectionString = newRpConfiguration.connectionString;
           telemetryClient.updateConnectionStrings(
-              newRpConfiguration.connectionString,
+              configuration.connectionString,
               configuration.internal.statsbeat.instrumentationKey,
               configuration.internal.statsbeat.endpoint);
           appIdSupplier.updateAppId();
         }
 
-        if (newRpConfiguration.sampling.percentage != rpConfiguration.sampling.percentage) {
-          logger.debug(
-              "Updating sampling percentage from {} to {}",
-              rpConfiguration.sampling.percentage,
-              newRpConfiguration.sampling.percentage);
-          float roundedSamplingPercentage =
-              ConfigurationBuilder.roundToNearest(newRpConfiguration.sampling.percentage);
-          DelegatingSampler.getInstance()
-              .setDelegate(Samplers.getSampler(roundedSamplingPercentage, configuration));
-          BytecodeUtilImpl.samplingPercentage = roundedSamplingPercentage;
-          rpConfiguration.sampling.percentage = newRpConfiguration.sampling.percentage;
+        if (!Objects.equals(
+                newRpConfiguration.sampling.percentage, rpConfiguration.sampling.percentage)
+            || !Objects.equals(
+                newRpConfiguration.sampling.limitPerSecond,
+                rpConfiguration.sampling.limitPerSecond)) {
+          logger.debug("Updating sampling percentage");
+          configuration.sampling.percentage = newRpConfiguration.sampling.percentage;
+          configuration.sampling.limitPerSecond = newRpConfiguration.sampling.limitPerSecond;
+          DelegatingSampler.getInstance().setDelegate(Samplers.getSampler(configuration));
+          if (configuration.sampling.percentage != null) {
+            BytecodeUtilImpl.samplingPercentage = configuration.sampling.percentage.floatValue();
+          } else {
+            BytecodeUtilImpl.samplingPercentage = 100;
+          }
         }
         rpConfiguration = newRpConfiguration;
       }