Rate limited sampling (#2456)

* Rate limited sampling

* Comment

* Comment

* Feedback

* Private

* Link

* Fix merge

* Feedback

* Fix

* Fix

* Fix

* Add test

* Spotless

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,18 @@ public enum SpanKind {
     }
   }
 
+  public enum SamplingTelemetryKind {
+    // restricted to telemetry kinds that are supported by SamplingOverrides
+    @JsonProperty("request")
+    REQUEST,
+    @JsonProperty("dependency")
+    DEPENDENCY,
+    @JsonProperty("trace")
+    TRACE,
+    @JsonProperty("exception")
+    EXCEPTION
+  }
+
   public enum MatchType {
     @JsonProperty("strict")
     STRICT,
@@ -151,7 +163,12 @@ 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 (set in ConfigurationBuilder if neither percentage nor
+    // limitPerSecond was configured)
+    @Nullable public Double limitPerSecond;
   }
 
   public static class SamplingPreview {
@@ -175,6 +192,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 +371,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 +597,36 @@ 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 SamplingTelemetryKind telemetryKind;
+
+    // TODO (trask) add test for this
+    // this is primarily useful for batch jobs
+    public boolean includeStandaloneTelemetry;
+
     // 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 == SamplingTelemetryKind.REQUEST
+          // this part is for backwards compatibility:
+          || (telemetryKind == null && spanKind != SpanKind.CLIENT);
+    }
+
+    public boolean isForDependencyTelemetry() {
+      return telemetryKind == SamplingTelemetryKind.DEPENDENCY
+          // this part is for backwards compatibility:
+          || (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

@@ -84,6 +84,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";
 
@@ -190,6 +193,15 @@ private static void logConfigurationWarnings(Configuration config) {
               + " 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);
   }
 
@@ -206,6 +218,10 @@ private static void overlayConfiguration(
       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();
@@ -478,6 +494,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 =
@@ -487,10 +507,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(
@@ -597,6 +616,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
@@ -646,13 +666,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("applying environment variable: {}={}", name, value);
       // intentionally allowing NumberFormatException to bubble up as invalid configuration and
       // prevent agent from starting
-      return Float.parseFloat(value);
+      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 Integer.parseInt(value);
     }
     return defaultValue;
   }
@@ -827,17 +859,21 @@ static String getJsonEncodingExceptionMessage(String message, JsonOrigin jsonOri
   }
 
   // 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,46 @@ 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;
+
+        SpanContext spanContext = log.getSpanContext();
+
+        boolean standaloneLog = !spanContext.isValid();
+        Double samplingPercentage =
+            samplingOverrides.getOverridePercentage(standaloneLog, log.getAttributes());
+
+        if (samplingPercentage != null && !shouldSample(spanContext, samplingPercentage)) {
+          continue;
+        }
+
+        if (samplingPercentage == null
+            && !standaloneLog
+            && !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 +157,23 @@ 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);
+    }
+    // this is a standalone log (not part of a trace), so randomly sample at the given percentage
+    return ThreadLocalRandom.current().nextDouble() < percentage / 100;
+  }
 }