Support sampling overrides on http request headers (#2308)

microsoft · Jun 1, 2022 · 024e347 · 024e347
1 parent e1b89d4
commit 024e347
Show file tree

Hide file tree

Showing 5 changed files with 81 additions and 9 deletions.
diff --git a/...in/java/com/microsoft/applicationinsights/agent/internal/configuration/Configuration.java b/...in/java/com/microsoft/applicationinsights/agent/internal/configuration/Configuration.java
@@ -344,7 +344,7 @@ public void validate() {
 
   public static class InheritedAttribute {
     public String key;
-    public SpanAttributeType type;
+    public AttributeType type;
 
     public AttributeKey<?> getAttributeKey() {
       switch (type) {
@@ -374,7 +374,7 @@ public static class HttpHeadersConfiguration {
     public List<String> responseHeaders = new ArrayList<>();
   }
 
-  public enum SpanAttributeType {
+  public enum AttributeType {
     @JsonProperty("string")
     STRING,
     @JsonProperty("boolean")

diff --git a/...ain/java/com/microsoft/applicationinsights/agent/internal/sampling/SamplingOverrides.java b/...ain/java/com/microsoft/applicationinsights/agent/internal/sampling/SamplingOverrides.java
@@ -182,15 +182,29 @@ private boolean matches(SpanKind spanKind, Attributes attributes, LazyHttpUrl la
 
     private static TempPredicate toPredicate(SamplingOverrideAttribute attribute) {
       if (attribute.matchType == MatchType.STRICT) {
-        return new StrictMatcher(attribute.key, attribute.value);
+        if (isHttpHeaderAttribute(attribute)) {
+          return new StrictArrayContainsMatcher(attribute.key, attribute.value);
+        } else {
+          return new StrictMatcher(attribute.key, attribute.value);
+        }
       } else if (attribute.matchType == MatchType.REGEXP) {
-        return new RegexpMatcher(attribute.key, attribute.value);
+        if (isHttpHeaderAttribute(attribute)) {
+          return new RegexpArrayContainsMatcher(attribute.key, attribute.value);
+        } else {
+          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);
       }
     }
+
+    private static boolean isHttpHeaderAttribute(SamplingOverrideAttribute attribute) {
+      // note that response headers are not typically available for sampling
+      return attribute.key.startsWith("http.request.header.")
+          || attribute.key.startsWith("http.response.header.");
+    }
   }
 
   private static class StrictMatcher implements TempPredicate {
@@ -212,6 +226,22 @@ public boolean test(Attributes attributes, LazyHttpUrl lazyHttpUrl) {
     }
   }
 
+  private static class StrictArrayContainsMatcher implements TempPredicate {
+    private final AttributeKey<List<String>> key;
+    private final String value;
+
+    private StrictArrayContainsMatcher(String key, String value) {
+      this.key = AttributeKey.stringArrayKey(key);
+      this.value = value;
+    }
+
+    @Override
+    public boolean test(Attributes attributes, LazyHttpUrl lazyHttpUrl) {
+      List<String> val = attributes.get(key);
+      return val != null && val.contains(value);
+    }
+  }
+
   private static class RegexpMatcher implements TempPredicate {
     private final AttributeKey<String> key;
     private final Pattern value;
@@ -231,6 +261,30 @@ public boolean test(Attributes attributes, LazyHttpUrl lazyHttpUrl) {
     }
   }
 
+  private static class RegexpArrayContainsMatcher implements TempPredicate {
+    private final AttributeKey<List<String>> key;
+    private final Pattern value;
+
+    private RegexpArrayContainsMatcher(String key, String value) {
+      this.key = AttributeKey.stringArrayKey(key);
+      this.value = Pattern.compile(value);
+    }
+
+    @Override
+    public boolean test(Attributes attributes, LazyHttpUrl lazyHttpUrl) {
+      List<String> val = attributes.get(key);
+      if (val == null) {
+        return false;
+      }
+      for (String v : val) {
+        if (value.matcher(v).matches()) {
+          return true;
+        }
+      }
+      return false;
+    }
+  }
+
   private static class KeyOnlyMatcher implements TempPredicate {
     private final AttributeKey<String> key;
 

diff --git a/...crosoft/applicationinsights/agent/internal/init/InheritedAttributesSpanProcessorTest.java b/...crosoft/applicationinsights/agent/internal/init/InheritedAttributesSpanProcessorTest.java
@@ -87,7 +87,7 @@ public void shouldNotInheritAttribute() {
   public void shouldInheritAttribute() {
     Configuration.InheritedAttribute inheritedAttribute = new Configuration.InheritedAttribute();
     inheritedAttribute.key = "one";
-    inheritedAttribute.type = Configuration.SpanAttributeType.STRING;
+    inheritedAttribute.type = Configuration.AttributeType.STRING;
 
     Tracer tracer = newTracer(Collections.singletonList(inheritedAttribute));
     Span span =
@@ -127,7 +127,7 @@ public void shouldInheritAttribute() {
   public void shouldNotInheritAttributeWithSameNameButDifferentType() {
     Configuration.InheritedAttribute inheritedAttribute = new Configuration.InheritedAttribute();
     inheritedAttribute.key = "one";
-    inheritedAttribute.type = Configuration.SpanAttributeType.STRING;
+    inheritedAttribute.type = Configuration.AttributeType.STRING;
 
     Tracer tracer = newTracer(Collections.singletonList(inheritedAttribute));
     Span span = tracer.spanBuilder("parent").setNoParent().setAttribute(oneLongKey, 1L).startSpan();

diff --git a/...pHeaders/src/main/java/com/microsoft/applicationinsights/smoketestapp/TestController.java b/...pHeaders/src/main/java/com/microsoft/applicationinsights/smoketestapp/TestController.java
@@ -45,10 +45,16 @@ public String serverHeaders(HttpServletResponse response) {
 
   @GetMapping("/clientHeaders")
   public String clientHeaders() throws IOException {
+    callMockCodes200(false);
+    callMockCodes200(true);
+    return "OK!";
+  }
+
+  private void callMockCodes200(boolean nope) throws IOException {
     URL obj = new URL("https://mock.codes/200");
 
     HttpURLConnection connection = (HttpURLConnection) obj.openConnection();
-    connection.setRequestProperty("abc", "testing123");
+    connection.setRequestProperty("abc", nope ? "nope" : "testing123");
     // calling getContentType() first, since this triggered a bug previously in the instrumentation
     // previously
     connection.getContentType();
@@ -57,7 +63,5 @@ public String clientHeaders() throws IOException {
     byte[] buffer = new byte[1024];
     while (content.read(buffer) != -1) {}
     content.close();
-
-    return "OK!";
   }
 }
diff --git a/smoke-tests/servers/global-resources/httpheaders_applicationinsights.json b/smoke-tests/servers/global-resources/httpheaders_applicationinsights.json
@@ -20,6 +20,20 @@
       "responseHeaders": [
         "date"
       ]
+    },
+    "sampling": {
+      "overrides": [
+        {
+          "attributes": [
+            {
+              "key": "http.request.header.abc",
+              "value": "nope",
+              "matchType": "strict"
+            }
+          ],
+          "percentage": 0
+        }
+      ]
     }
   }
 }