Add span kind to sampling overrides (#1960)

* Option 1

* Revert "Option 1"

This reverts commit cf5b048.

* Add span kind to sampling overrides

* Support key-only matches

* Fix

* Bump version

* Fix sporadic test failure

* Fix test

* Add todo

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

@@ -37,6 +37,7 @@
 import java.util.Map;
 import java.util.regex.Pattern;
 import java.util.regex.PatternSyntaxException;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 // an assumption is made throughout this file that user will not explicitly use `null` value in json
 // file
@@ -62,6 +63,26 @@ private static boolean isEmpty(String str) {
     return str == null || str.trim().isEmpty();
   }
 
+  // TODO (trask) investigate options for mapping lowercase values to otel enum directly
+  public enum SpanKind {
+    @JsonProperty("server")
+    SERVER(io.opentelemetry.api.trace.SpanKind.SERVER),
+    @JsonProperty("client")
+    CLIENT(io.opentelemetry.api.trace.SpanKind.CLIENT),
+    @JsonProperty("consumer")
+    CONSUMER(io.opentelemetry.api.trace.SpanKind.CONSUMER),
+    @JsonProperty("producer")
+    PRODUCER(io.opentelemetry.api.trace.SpanKind.PRODUCER),
+    @JsonProperty("internal")
+    INTERNAL(io.opentelemetry.api.trace.SpanKind.INTERNAL);
+
+    public final io.opentelemetry.api.trace.SpanKind otelSpanKind;
+
+    SpanKind(io.opentelemetry.api.trace.SpanKind otelSpanKind) {
+      this.otelSpanKind = otelSpanKind;
+    }
+  }
+
   public enum MatchType {
     @JsonProperty("strict")
     STRICT,
@@ -373,6 +394,8 @@ private static String getDefaultPath() {
   }
 
   public static class SamplingOverride {
+    // TODO (trask) consider making this required when moving out of preview
+    @Nullable public SpanKind spanKind;
     // 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)
@@ -381,11 +404,11 @@ public static class SamplingOverride {
     public String id; // optional, used for debugging purposes only
 
     public void validate() {
-      if (attributes.isEmpty()) {
+      if (spanKind == null && attributes.isEmpty()) {
         // TODO add doc and go link, similar to telemetry processors
         throw new FriendlyException(
-            "A sampling override configuration has no attributes.",
-            "Please provide one or more attributes for the sampling override configuration.");
+            "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
@@ -407,8 +430,8 @@ public void validate() {
 
   public static class SamplingOverrideAttribute {
     public String key;
-    public String value;
-    public MatchType matchType;
+    @Nullable public String value;
+    @Nullable public MatchType matchType;
 
     private void validate() {
       if (isEmpty(key)) {
@@ -417,9 +440,9 @@ private void validate() {
             "A sampling override configuration has an attribute section that is missing a \"key\".",
             "Please provide a \"key\" under the attribute section of the sampling override configuration.");
       }
-      if (matchType == null) {
+      if (matchType == null && value != null) {
         throw new FriendlyException(
-            "A sampling override configuration has an attribute section that is missing a \"matchType\".",
+            "A sampling override configuration has an attribute section with a \"value\" that is missing a \"matchType\".",
             "Please provide a \"matchType\" under the attribute section of the sampling override configuration.");
       }
       if (matchType == MatchType.REGEXP) {

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/sampling/AiSampler.java

@@ -87,7 +87,7 @@ public SamplingResult shouldSample(
       Attributes attributes,
       List<LinkData> parentLinks) {
 
-    MatcherGroup override = samplingOverrides.getOverride(attributes);
+    MatcherGroup override = samplingOverrides.getOverride(spanKind, attributes);
 
     if (override != null) {
       return getSamplingResult(

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/sampling/SamplingOverrides.java

@@ -28,6 +28,7 @@
 import com.microsoft.applicationinsights.agent.internal.telemetry.TelemetryUtil;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.TraceState;
 import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
@@ -37,7 +38,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.regex.Pattern;
-import org.checkerframework.checker.nullness.qual.Nullable;
+import javax.annotation.Nullable;
 
 // TODO find a better name for this class (and MatcherGroup too)
 class SamplingOverrides {
@@ -52,10 +53,10 @@ class SamplingOverrides {
   }
 
   @Nullable
-  MatcherGroup getOverride(Attributes attributes) {
+  MatcherGroup getOverride(SpanKind spanKind, Attributes attributes) {
     LazyHttpUrl lazyHttpUrl = new LazyHttpUrl(attributes);
     for (MatcherGroup matcherGroups : matcherGroups) {
-      if (matcherGroups.matches(attributes, lazyHttpUrl)) {
+      if (matcherGroups.matches(spanKind, attributes, lazyHttpUrl)) {
         return matcherGroups;
       }
     }
@@ -143,11 +144,13 @@ public TraceState getUpdatedTraceState(TraceState parentTraceState) {
   }
 
   static class MatcherGroup {
+    @Nullable private final SpanKind spanKind;
     private final List<TempPredicate> predicates;
     private final double percentage;
     private final SamplingResult recordAndSampleAndOverwriteTraceState;
 
     private MatcherGroup(SamplingOverride override) {
+      spanKind = override.spanKind != null ? override.spanKind.otelSpanKind : null;
       predicates = new ArrayList<>();
       for (SamplingOverrideAttribute attribute : override.attributes) {
         predicates.add(toPredicate(attribute));
@@ -165,7 +168,10 @@ SamplingResult getRecordAndSampleAndOverwriteTraceState() {
       return recordAndSampleAndOverwriteTraceState;
     }
 
-    private boolean matches(Attributes attributes, LazyHttpUrl lazyHttpUrl) {
+    private boolean matches(SpanKind spanKind, Attributes attributes, LazyHttpUrl lazyHttpUrl) {
+      if (this.spanKind != null && !this.spanKind.equals(spanKind)) {
+        return false;
+      }
       for (TempPredicate predicate : predicates) {
         if (!predicate.test(attributes, lazyHttpUrl)) {
           return false;
@@ -179,6 +185,8 @@ private static TempPredicate toPredicate(SamplingOverrideAttribute attribute) {
         return new StrictMatcher(attribute.key, attribute.value);
       } else if (attribute.matchType == MatchType.REGEXP) {
         return new RegexpMatcher(attribute.key, attribute.value);
+      } else if (attribute.matchType == null) {
+        return new KeyOnlyMatcher(attribute.key);
       } else {
         throw new IllegalStateException("Unexpected match type: " + attribute.matchType);
       }
@@ -223,6 +231,23 @@ public boolean test(Attributes attributes, LazyHttpUrl lazyHttpUrl) {
     }
   }
 
+  private static class KeyOnlyMatcher implements TempPredicate {
+    private final AttributeKey<String> key;
+
+    private KeyOnlyMatcher(String key) {
+      this.key = AttributeKey.stringKey(key);
+    }
+
+    @Override
+    public boolean test(Attributes attributes, LazyHttpUrl lazyHttpUrl) {
+      String val = attributes.get(key);
+      if (val == null && key.getKey().equals(SemanticAttributes.HTTP_URL.getKey())) {
+        val = lazyHttpUrl.get();
+      }
+      return val != null;
+    }
+  }
+
   // this is temporary until semantic attributes stabilize and we make breaking change
   private static class LazyHttpUrl {
     private final Attributes attributes;

agent/agent-tooling/src/test/java/com/microsoft/applicationinsights/agent/internal/heartbeat/HeartbeatTests.java

@@ -86,7 +86,7 @@ void heartBeatPayloadContainsDataByDefault() throws InterruptedException {
     provider.initialize(TelemetryClient.createForTest());
 
     // some of the initialization above happens in a separate thread
-    Thread.sleep(100);
+    Thread.sleep(500);
 
     // then
     MetricsData t = (MetricsData) provider.gatherData().getData().getBaseData();