Add Micronaut 4 support for code origin for spans (#8039)

Make Code Origin for spans feature independent from dynamic instrumentation so it could be enabled while DI is not Add in Status Logger the code origin feature status
DataDog · Dec 4, 2024 · 5b58772 · 5b58772
1 parent 686f51a
commit 5b58772
Show file tree

Hide file tree

Showing 10 changed files with 111 additions and 35 deletions.
diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java
@@ -105,6 +105,9 @@ private enum AgentFeature {
     DEBUGGER(propertyNameToSystemPropertyName(DebuggerConfig.DEBUGGER_ENABLED), false),
     EXCEPTION_DEBUGGING(
         propertyNameToSystemPropertyName(DebuggerConfig.EXCEPTION_REPLAY_ENABLED), false),
+    SPAN_ORIGIN(
+        propertyNameToSystemPropertyName(TraceInstrumentationConfig.CODE_ORIGIN_FOR_SPANS_ENABLED),
+        false),
     DATA_JOBS(propertyNameToSystemPropertyName(GeneralConfig.DATA_JOBS_ENABLED), false),
     AGENTLESS_LOG_SUBMISSION(
         propertyNameToSystemPropertyName(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED), false);
@@ -152,6 +155,7 @@ public boolean isEnabledByDefault() {
   private static boolean telemetryEnabled = true;
   private static boolean debuggerEnabled = false;
   private static boolean exceptionDebuggingEnabled = false;
+  private static boolean spanOriginEnabled = false;
   private static boolean agentlessLogSubmissionEnabled = false;
 
   /**
@@ -263,6 +267,7 @@ public static void start(
     telemetryEnabled = isFeatureEnabled(AgentFeature.TELEMETRY);
     debuggerEnabled = isFeatureEnabled(AgentFeature.DEBUGGER);
     exceptionDebuggingEnabled = isFeatureEnabled(AgentFeature.EXCEPTION_DEBUGGING);
+    spanOriginEnabled = isFeatureEnabled(AgentFeature.SPAN_ORIGIN);
     agentlessLogSubmissionEnabled = isFeatureEnabled(AgentFeature.AGENTLESS_LOG_SUBMISSION);
 
     if (profilingEnabled) {
@@ -1073,7 +1078,7 @@ private static void shutdownProfilingAgent(final boolean sync) {
   }
 
   private static void maybeStartDebugger(Instrumentation inst, Class<?> scoClass, Object sco) {
-    if (!debuggerEnabled && !exceptionDebuggingEnabled) {
+    if (!debuggerEnabled && !exceptionDebuggingEnabled && !spanOriginEnabled) {
       return;
     }
     if (!remoteConfigEnabled) {

diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerAgent.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerAgent.java
@@ -100,6 +100,7 @@ public static synchronized void run(
       DebuggerContext.initExceptionDebugger(defaultExceptionDebugger);
     }
     if (config.isDebuggerCodeOriginEnabled()) {
+      LOGGER.info("Starting Code Origin for spans");
       DebuggerContext.initCodeOrigin(new DefaultCodeOriginRecorder(config, configurationUpdater));
     }
     if (config.isDebuggerInstrumentTheWorld()) {

diff --git a/...ent-debugger/src/main/java/com/datadog/debugger/codeorigin/DefaultCodeOriginRecorder.java b/...ent-debugger/src/main/java/com/datadog/debugger/codeorigin/DefaultCodeOriginRecorder.java
@@ -30,7 +30,7 @@ public class DefaultCodeOriginRecorder implements CodeOriginRecorder {
 
   private final ConfigurationUpdater configurationUpdater;
 
-  private final Map<String, CodeOriginProbe> fingerprints = new HashMap<>();
+  private final Map<String, CodeOriginProbe> probesByFingerprint = new HashMap<>();
 
   private final Map<String, CodeOriginProbe> probes = new ConcurrentHashMap<>();
 
@@ -45,36 +45,28 @@ public DefaultCodeOriginRecorder(Config config, ConfigurationUpdater configurati
   public String captureCodeOrigin(boolean entry) {
     StackTraceElement element = findPlaceInStack();
     String fingerprint = Fingerprinter.fingerprint(element);
-    CodeOriginProbe probe;
-
-    if (isAlreadyInstrumented(fingerprint)) {
-      probe = fingerprints.get(fingerprint);
-    } else {
-      probe =
-          createProbe(
-              fingerprint,
-              entry,
-              Where.of(
-                  element.getClassName(),
-                  element.getMethodName(),
-                  null,
-                  String.valueOf(element.getLineNumber())));
+    CodeOriginProbe probe = probesByFingerprint.get(fingerprint);
+    if (probe == null) {
+      Where where =
+          Where.of(
+              element.getClassName(),
+              element.getMethodName(),
+              null,
+              String.valueOf(element.getLineNumber()));
+      probe = createProbe(fingerprint, entry, where);
+      LOG.debug("Creating probe for location {}", where);
     }
-
     return probe.getId();
   }
 
   @Override
   public String captureCodeOrigin(Method method, boolean entry) {
-    CodeOriginProbe probe;
-
-    String fingerPrint = method.toString();
-    if (isAlreadyInstrumented(fingerPrint)) {
-      probe = fingerprints.get(fingerPrint);
-    } else {
-      probe = createProbe(fingerPrint, entry, Where.of(method));
+    String fingerprint = method.toString();
+    CodeOriginProbe probe = probesByFingerprint.get(fingerprint);
+    if (probe == null) {
+      probe = createProbe(fingerprint, entry, Where.of(method));
+      LOG.debug("Creating probe for method {}", fingerprint);
     }
-
     return probe.getId();
   }
 
@@ -106,22 +98,16 @@ private StackTraceElement findPlaceInStack() {
                 .orElse(null));
   }
 
-  public boolean isAlreadyInstrumented(String fingerprint) {
-    return fingerprints.containsKey(fingerprint);
-  }
-
   void addFingerprint(String fingerprint, CodeOriginProbe probe) {
-    fingerprints.putIfAbsent(fingerprint, probe);
+    probesByFingerprint.putIfAbsent(fingerprint, probe);
   }
 
-  public String installProbe(CodeOriginProbe probe) {
+  public void installProbe(CodeOriginProbe probe) {
     CodeOriginProbe installed = probes.putIfAbsent(probe.getId(), probe);
     if (installed == null) {
       AgentTaskScheduler.INSTANCE.execute(
           () -> configurationUpdater.accept(CODE_ORIGIN, getProbes()));
-      return probe.getId();
     }
-    return installed.getId();
   }
 
   public CodeOriginProbe getProbe(String probeId) {

diff --git a/dd-java-agent/instrumentation/micronaut/build.gradle b/dd-java-agent/instrumentation/micronaut/build.gradle
@@ -4,6 +4,16 @@
 
 apply from: "$rootDir/gradle/java.gradle"
 
+muzzle {
+  pass {
+    name = "micronaut-common"
+    group = "io.micronaut"
+    module = "micronaut-http-server-netty"
+    versions = "[2,)"
+  }
+}
+
 dependencies {
   compileOnly group: 'io.micronaut', name: 'micronaut-http-server-netty', version: '2.0.0'
+  implementation project(':dd-java-agent:instrumentation:span-origin')
 }
diff --git a/dd-java-agent/instrumentation/micronaut/http-server-netty-4.0/build.gradle b/dd-java-agent/instrumentation/micronaut/http-server-netty-4.0/build.gradle
@@ -30,11 +30,13 @@ addTestSuiteForDir('latestDepTest', 'test')
 
 dependencies {
   main_java17CompileOnly group: 'io.micronaut', name: 'micronaut-http-server-netty', version: '4.0.0'
+  implementation project(':dd-java-agent:instrumentation:micronaut')
 
   // Added to ensure cross compatibility:
   testImplementation project(':dd-java-agent:instrumentation:micronaut:http-server-netty-2.0')
   testImplementation project(':dd-java-agent:instrumentation:micronaut:http-server-netty-3.0')
   testImplementation project(':dd-java-agent:instrumentation:netty-4.1')
+  testImplementation project(':dd-java-agent:agent-debugger')
   testImplementation group: 'io.micronaut', name: 'micronaut-http-server-netty', version: '4.0.0', {
     exclude group: 'org.slf4j', module: 'slf4j-api'
     exclude group: 'ch.qos.logback', module: 'logback-classic'

diff --git a/...gent/instrumentation/micronaut/http-server-netty-4.0/src/test/groovy/MicronautTest.groovy b/...gent/instrumentation/micronaut/http-server-netty-4.0/src/test/groovy/MicronautTest.groovy
@@ -2,20 +2,49 @@ import datadog.trace.agent.test.asserts.TraceAssert
 import datadog.trace.agent.test.base.HttpServer
 import datadog.trace.agent.test.base.HttpServerTest
 import datadog.trace.api.DDSpanTypes
+import datadog.trace.bootstrap.debugger.DebuggerContext
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.instrumentation.micronaut.v4_0.MicronautDecorator
 import datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator
 import test.MicronautServer
 
+import java.lang.reflect.Method
+
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.NOT_FOUND
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.PATH_PARAM
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_ENCODED_BOTH
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS
+import static datadog.trace.api.config.TraceInstrumentationConfig.CODE_ORIGIN_FOR_SPANS_ENABLED
 
 class MicronautTest extends HttpServerTest<Object> {
 
+  def codeOriginRecorder
+
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig(CODE_ORIGIN_FOR_SPANS_ENABLED, "true")
+    codeOriginRecorder = new DebuggerContext.CodeOriginRecorder() {
+        def invoked = false
+        @Override
+        String captureCodeOrigin(boolean entry) {
+          invoked = true
+          return "done"
+        }
+
+        @Override
+        String captureCodeOrigin(Method method, boolean entry) {
+          invoked = true
+          return "done"
+        }
+      }
+    DebuggerContext.initCodeOrigin(codeOriginRecorder)
+  }
+
+
+
   @Override
   HttpServer server() {
     return new MicronautServer()
@@ -67,6 +96,9 @@ class MicronautTest extends HttpServerTest<Object> {
 
   @Override
   void handlerSpan(TraceAssert trace, ServerEndpoint endpoint = SUCCESS) {
+    if (endpoint != NOT_FOUND) {
+      assert codeOriginRecorder.invoked
+    }
     trace.span {
       serviceName expectedServiceName()
       operationName "micronaut-controller"

diff --git a/...main/java/datadog/trace/instrumentation/micronaut/MicronautCodeOriginInstrumentation.java b/...main/java/datadog/trace/instrumentation/micronaut/MicronautCodeOriginInstrumentation.java
@@ -0,0 +1,36 @@
+package datadog.trace.instrumentation.micronaut;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.instrumentation.codeorigin.CodeOriginInstrumentation;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
+@AutoService(InstrumenterModule.class)
+public class MicronautCodeOriginInstrumentation extends CodeOriginInstrumentation {
+
+  public static final String IO_MICRONAUT_HTTP_ANNOTATION = "io.micronaut.http.annotation.";
+
+  public MicronautCodeOriginInstrumentation() {
+    super("micronaut", "micronaut-span-origin");
+  }
+
+  @Override
+  public String muzzleDirective() {
+    return "micronaut-common";
+  }
+
+  @Override
+  protected Set<String> getAnnotations() {
+    return new HashSet<>(
+        Arrays.asList(
+            IO_MICRONAUT_HTTP_ANNOTATION + "Get",
+            IO_MICRONAUT_HTTP_ANNOTATION + "Post",
+            IO_MICRONAUT_HTTP_ANNOTATION + "Put",
+            IO_MICRONAUT_HTTP_ANNOTATION + "Delete",
+            IO_MICRONAUT_HTTP_ANNOTATION + "Patch",
+            IO_MICRONAUT_HTTP_ANNOTATION + "Head",
+            IO_MICRONAUT_HTTP_ANNOTATION + "Options"));
+  }
+}
diff --git a/...gin/src/main/java/datadog/trace/instrumentation/codeorigin/CodeOriginInstrumentation.java b/...gin/src/main/java/datadog/trace/instrumentation/codeorigin/CodeOriginInstrumentation.java
@@ -17,8 +17,8 @@ public abstract class CodeOriginInstrumentation extends Tracing implements ForTy
   private final OneOf<NamedElement> matcher;
 
   @SuppressForbidden
-  public CodeOriginInstrumentation(String instrumentationName) {
-    super(instrumentationName);
+  public CodeOriginInstrumentation(String instrumentationName, String... additionalNames) {
+    super(instrumentationName, additionalNames);
     this.matcher = NameMatchers.namedOneOf(getAnnotations());
   }
 

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/StatusLogger.java b/dd-trace-core/src/main/java/datadog/trace/core/StatusLogger.java
@@ -114,6 +114,8 @@ public void toJson(JsonWriter writer, Config config) throws IOException {
     writer.value(config.isDebuggerEnabled());
     writer.name("debugger_exception_enabled");
     writer.value(config.isDebuggerExceptionEnabled());
+    writer.name("debugger_span_origin_enabled");
+    writer.value(config.isDebuggerCodeOriginEnabled());
     writer.name("appsec_enabled");
     writer.value(config.getAppSecActivation().toString());
     writer.name("appsec_rules_file_path");

diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java
@@ -4326,6 +4326,8 @@ public String toString() {
         + debuggerSymbolIncludes
         + ", debuggerExceptionEnabled="
         + debuggerExceptionEnabled
+        + ", debuggerCodeOriginEnabled="
+        + debuggerCodeOriginEnabled
         + ", awsPropagationEnabled="
         + awsPropagationEnabled
         + ", sqsPropagationEnabled="