Fix custom instrumentation type (#2776)

See #2774 Co-authored-by: github-actions[bot] <github-action[bot]@users.noreply.github.com>
microsoft · Dec 11, 2022 · e21d92f · e21d92f
1 parent 4d6b393
commit e21d92f
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 11 deletions.
diff --git a/...va/io/opentelemetry/javaagent/instrumentation/methods/ai/MethodInstrumentationModule.java b/...va/io/opentelemetry/javaagent/instrumentation/methods/ai/MethodInstrumentationModule.java
@@ -9,8 +9,8 @@
 
 package io.opentelemetry.javaagent.instrumentation.methods.ai;
 
+import static java.util.Arrays.asList;
 import static java.util.Collections.emptyList;
-import static java.util.Collections.singletonList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
@@ -49,7 +49,9 @@ public MethodInstrumentationModule() {
   public List<String> getAdditionalHelperClassNames() {
     return typeInstrumentations.isEmpty()
         ? emptyList()
-        : singletonList("io.opentelemetry.javaagent.instrumentation.methods.ai.MethodSingletons");
+        : asList(
+            "io.opentelemetry.javaagent.instrumentation.methods.ai.MethodSingletons",
+            "io.opentelemetry.javaagent.instrumentation.methods.ai.MethodSingletons$MethodSpanKindExtractor");
   }
 
   @Override

diff --git a/...src/main/java/io/opentelemetry/javaagent/instrumentation/methods/ai/MethodSingletons.java b/...src/main/java/io/opentelemetry/javaagent/instrumentation/methods/ai/MethodSingletons.java
@@ -10,6 +10,8 @@
 package io.opentelemetry.javaagent.instrumentation.methods.ai;
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
 import io.opentelemetry.instrumentation.api.instrumenter.code.CodeAttributesExtractor;
@@ -33,15 +35,29 @@ public final class MethodSingletons {
                 CodeSpanNameExtractor.create(codeAttributesGetter))
             .addAttributesExtractor(CodeAttributesExtractor.create(codeAttributesGetter))
             // START APPLICATION INSIGHTS MODIFICATIONS
-            // we emit SERVER spans instead of INTERNAL spans, so that it works well with
-            // Application Insights' customInstrumentations feature
-            .buildInstrumenter(SpanKindExtractor.alwaysServer());
+            .buildInstrumenter(new MethodSpanKindExtractor());
     // END APPLICATION INSIGHTS MODIFICATIONS
   }
 
   public static Instrumenter<ClassAndMethod, Void> instrumenter() {
     return INSTRUMENTER;
   }
 
+  // START APPLICATION INSIGHTS MODIFICATIONS
+  private static class MethodSpanKindExtractor implements SpanKindExtractor<ClassAndMethod> {
+
+    @Override
+    public SpanKind extract(ClassAndMethod classAndMethod) {
+      // we emit SERVER spans instead of INTERNAL spans when there is no parent, so that it works
+      // well with Application Insights' customInstrumentations feature
+      return willHaveParentSpan() ? SpanKind.INTERNAL : SpanKind.SERVER;
+    }
+
+    private static boolean willHaveParentSpan() {
+      return Span.current().getSpanContext().isValid();
+    }
+  }
+  // END APPLICATION INSIGHTS MODIFICATIONS
+
   private MethodSingletons() {}
 }
diff --git a/...entation/src/main/java/com/microsoft/applicationinsights/smoketestapp/TestController.java b/...entation/src/main/java/com/microsoft/applicationinsights/smoketestapp/TestController.java
@@ -18,14 +18,20 @@ public String root() {
     return "OK";
   }
 
-  @GetMapping("/test")
-  public String test() {
+  @GetMapping("/server-span")
+  public String serverSpan() {
     // spin off separate thread to simulate a top-level (request) instrumentation captured on the
     // run() method
     new Thread(this::run).start();
     return "OK!";
   }
 
+  @GetMapping("/internal-span")
+  public String internalSpan() {
+    run();
+    return "OK!";
+  }
+
   private void run() {
     logger.info("hello");
   }

diff --git a/...smokeTest/java/com/microsoft/applicationinsights/smoketest/CustomInstrumentationTest.java b/...smokeTest/java/com/microsoft/applicationinsights/smoketest/CustomInstrumentationTest.java
@@ -30,11 +30,51 @@ public abstract class CustomInstrumentationTest {
   @RegisterExtension static final SmokeTestExtension testing = SmokeTestExtension.create();
 
   @Test
-  @TargetUri("/test")
-  void test() throws Exception {
+  @TargetUri("/internal-span")
+  void internalSpan() throws Exception {
+    Telemetry telemetry = testing.getTelemetry(1);
+
+    String operationId = telemetry.rdEnvelope.getTags().get("ai.operation.id");
+    List<Envelope> mdList = testing.mockedIngestion.waitForMessageItemsInRequest(1, operationId);
+
+    Envelope mdEnvelope = mdList.get(0);
+
+    assertThat(telemetry.rdEnvelope.getSampleRate()).isNull();
+    assertThat(telemetry.rddEnvelope1.getSampleRate()).isNull();
+    assertThat(mdEnvelope.getSampleRate()).isNull();
+
+    MessageData md = (MessageData) ((Data<?>) mdEnvelope.getData()).getBaseData();
+
+    assertThat(telemetry.rd.getName()).isEqualTo("GET /internal-span");
+    assertThat(telemetry.rd.getResponseCode()).isEqualTo("200");
+    assertThat(telemetry.rd.getProperties())
+        .containsExactly(entry("_MS.ProcessedByMetricExtractors", "True"));
+    assertThat(telemetry.rd.getSuccess()).isTrue();
+
+    assertThat(telemetry.rdd1.getName()).isEqualTo("TestController.run");
+    assertThat(telemetry.rdd1.getProperties()).isEmpty();
+    assertThat(telemetry.rdd1.getSuccess()).isTrue();
+
+    assertThat(md.getMessage()).isEqualTo("hello");
+    assertThat(md.getSeverityLevel()).isEqualTo(SeverityLevel.INFORMATION);
+    assertThat(md.getProperties()).containsEntry("SourceType", "Logger");
+    assertThat(md.getProperties()).containsEntry("LoggerName", "smoketestapp");
+    assertThat(md.getProperties()).containsKey("ThreadName");
+    assertThat(md.getProperties()).hasSize(3);
+
+    SmokeTestExtension.assertParentChild(
+        telemetry.rd, telemetry.rdEnvelope, telemetry.rddEnvelope1, "GET /internal-span");
+
+    SmokeTestExtension.assertParentChild(
+        telemetry.rdd1, telemetry.rddEnvelope1, mdEnvelope, "GET /internal-span");
+  }
+
+  @Test
+  @TargetUri("/server-span")
+  void serverSpan() throws Exception {
     List<Envelope> rdList = testing.mockedIngestion.waitForItems("RequestData", 2);
 
-    Envelope rdEnvelope1 = getRequestEnvelope(rdList, "GET /test");
+    Envelope rdEnvelope1 = getRequestEnvelope(rdList, "GET /server-span");
     Envelope rdEnvelope2 = getRequestEnvelope(rdList, "TestController.run");
 
     String operationId = rdEnvelope2.getTags().get("ai.operation.id");
@@ -50,7 +90,7 @@ void test() throws Exception {
     RequestData rd2 = (RequestData) ((Data<?>) rdEnvelope2.getData()).getBaseData();
     MessageData md = (MessageData) ((Data<?>) mdEnvelope.getData()).getBaseData();
 
-    assertThat(rd1.getName()).isEqualTo("GET /test");
+    assertThat(rd1.getName()).isEqualTo("GET /server-span");
     assertThat(rd1.getResponseCode()).isEqualTo("200");
     assertThat(rd1.getProperties())
         .containsExactly(entry("_MS.ProcessedByMetricExtractors", "True"));
@@ -67,6 +107,8 @@ void test() throws Exception {
     assertThat(md.getProperties()).containsEntry("LoggerName", "smoketestapp");
     assertThat(md.getProperties()).containsKey("ThreadName");
     assertThat(md.getProperties()).hasSize(3);
+
+    SmokeTestExtension.assertParentChild(rd2, rdEnvelope2, mdEnvelope, "TestController.run");
   }
 
   private static Envelope getRequestEnvelope(List<Envelope> envelopes, String name) {