Only decorate spans without code origin information

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/CodeOriginProbe.java

@@ -5,7 +5,6 @@
 import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_TYPE;
 import static java.lang.String.format;
 import static java.util.Arrays.asList;
-import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 
 import com.datadog.debugger.agent.DebuggerAgent;
@@ -54,6 +53,13 @@ public void evaluate(
     }
 
     super.evaluate(context, status, methodLocation);
+    AgentSpan span = findSpan(AgentTracer.activeSpan());
+
+    if (span == null) {
+      LOGGER.debug("Could not find the span for probeId {}", id);
+      return;
+    }
+    applyCodeOriginTags(span);
   }
 
   @Override
@@ -68,63 +74,53 @@ public void commit(
       LOGGER.debug("Could not find the span for probeId {}", id);
       return;
     }
-    String snapshotId = null;
     DebuggerSink sink = DebuggerAgent.getSink();
     if (isDebuggerEnabled(span) && sink != null) {
       Snapshot snapshot = createSnapshot();
       if (fillSnapshot(entryContext, exitContext, caughtExceptions, snapshot)) {
-        snapshotId = snapshot.getId();
+        String snapshotId = snapshot.getId();
         LOGGER.debug("committing code origin probe id={}, snapshot id={}", id, snapshotId);
         commitSnapshot(snapshot, sink);
+
+        List<AgentSpan> agentSpans =
+            entrySpanProbe ? asList(span, span.getLocalRootSpan()) : singletonList(span);
+        for (AgentSpan agentSpan : agentSpans) {
+          if (agentSpan.getTag(format(DD_CODE_ORIGIN_FRAME, 0, "snapshot_id")) == null) {
+            agentSpan.setTag(format(DD_CODE_ORIGIN_FRAME, 0, "snapshot_id"), snapshotId);
+          }
+        }
       }
     }
-    applySpanOriginTags(span, snapshotId);
     if (sink != null) {
       sink.getProbeStatusSink().addEmitting(probeId);
     }
     span.getLocalRootSpan().setTag(getId(), (String) null); // clear possible span reference
   }
 
-  private List<StackTraceElement> findLocation() {
-    if (entrySpanProbe && stackTraceElements == null) {
-      ProbeLocation probeLocation = getLocation();
-      List<String> lines = probeLocation.getLines();
-      int line = lines == null ? -1 : Integer.parseInt(lines.get(0));
-      stackTraceElements =
-          singletonList(
-              new StackTraceElement(
-                  probeLocation.getType(),
-                  probeLocation.getMethod(),
-                  probeLocation.getFile(),
-                  line));
-    }
-    return stackTraceElements != null ? stackTraceElements : emptyList();
-  }
-
-  private void applySpanOriginTags(AgentSpan span, String snapshotId) {
+  private List<AgentSpan> applyCodeOriginTags(AgentSpan span) {
     List<StackTraceElement> entries =
         stackTraceElements != null ? stackTraceElements : getUserStackFrames();
     List<AgentSpan> agentSpans =
         entrySpanProbe ? asList(span, span.getLocalRootSpan()) : singletonList(span);
 
     for (AgentSpan s : agentSpans) {
-      s.setTag(DD_CODE_ORIGIN_TYPE, entrySpanProbe ? "entry" : "exit");
-
-      for (int i = 0; i < entries.size(); i++) {
-        StackTraceElement info = entries.get(i);
-        String fileName = info.getFileName();
-        s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "file"), fileName);
-        s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "method"), info.getMethodName());
-        s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "line"), info.getLineNumber());
-        s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "type"), info.getClassName());
-        if (i == 0 && entrySpanProbe) {
-          s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "signature"), where.getSignature());
-        }
-        if (i == 0 && snapshotId != null) {
-          s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "snapshot_id"), snapshotId);
+      if (s.getTag(DD_CODE_ORIGIN_TYPE) == null) {
+        s.setTag(DD_CODE_ORIGIN_TYPE, entrySpanProbe ? "entry" : "exit");
+
+        for (int i = 0; i < entries.size(); i++) {
+          StackTraceElement info = entries.get(i);
+          String fileName = info.getFileName();
+          s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "file"), fileName);
+          s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "method"), info.getMethodName());
+          s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "line"), info.getLineNumber());
+          s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "type"), info.getClassName());
+          if (i == 0 && entrySpanProbe) {
+            s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "signature"), where.getSignature());
+          }
         }
       }
     }
+    return agentSpans;
   }
 
   public boolean entrySpanProbe() {

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/origin/CodeOriginTest.java

@@ -5,6 +5,7 @@
 import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_PREFIX;
 import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_TYPE;
 import static java.lang.String.format;
+import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
 import static java.util.Arrays.asList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -31,8 +32,11 @@
 import datadog.trace.bootstrap.debugger.Limits;
 import datadog.trace.bootstrap.debugger.ProbeId;
 import datadog.trace.core.CoreTracer;
+import java.io.File;
 import java.io.IOException;
 import java.net.URISyntaxException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
@@ -42,14 +46,19 @@
 import java.util.stream.Collectors;
 import org.jetbrains.annotations.NotNull;
 import org.joor.Reflect;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 public class CodeOriginTest extends CapturingTestBase {
 
   private static final ProbeId CODE_ORIGIN_ID1 = new ProbeId("code origin 1", 0);
+
   private static final ProbeId CODE_ORIGIN_ID2 = new ProbeId("code origin 2", 0);
 
+  private static final ProbeId CODE_ORIGIN_DOUBLE_ENTRY_ID =
+      new ProbeId("double entry code origin", 0);
+
   private static final int MAX_FRAMES = 20;
 
   private DefaultCodeOriginRecorder codeOriginRecorder;
@@ -84,6 +93,27 @@ public void before() {
     setFieldInConfig(InstrumenterConfig.get(), "codeOriginEnabled", true);
   }
 
+  @AfterAll
+  public static void collectClasses() throws IOException {
+    File[] files =
+        Paths.get(System.getProperty("java.io.tmpdir"), "debugger/com/datadog/debugger")
+            .toFile()
+            .listFiles(
+                file -> {
+                  String name = file.getName();
+                  return name.startsWith("CodeOrigin") && name.endsWith(".class");
+                });
+
+    if (new File("build").exists()) {
+      File buildDir = new File("build/debugger");
+      buildDir.mkdirs();
+      for (File file : files) {
+        Files.copy(
+            file.toPath(), Paths.get(buildDir.getAbsolutePath(), file.getName()), REPLACE_EXISTING);
+      }
+    }
+  }
+
   @Test
   public void basicInstrumentation() throws Exception {
     final String className = "com.datadog.debugger.CodeOrigin01";
@@ -115,6 +145,30 @@ public void withLogProbe() throws IOException, URISyntaxException {
     checkResults(testClass, "debug_1", true);
   }
 
+  @Test
+  public void doubleEntry() throws IOException, URISyntaxException {
+    final String className = "com.datadog.debugger.CodeOrigin05";
+
+    List<LogProbe> probes =
+        asList(
+            new CodeOriginProbe(
+                CODE_ORIGIN_ID1, true, Where.of(className, "entry", "()", "53"), MAX_FRAMES),
+            new CodeOriginProbe(
+                CODE_ORIGIN_ID2, false, Where.of(className, "exit", "()", "62"), MAX_FRAMES),
+            new CodeOriginProbe(
+                CODE_ORIGIN_DOUBLE_ENTRY_ID,
+                true,
+                Where.of(className, "doubleEntry", "()", "66"),
+                MAX_FRAMES));
+    installProbes(probes);
+    final Class<?> testClass = compileAndLoadClass(className);
+    checkResults(testClass, "fullTrace", false);
+    List<? extends MutableSpan> trace = traceInterceptor.getTrace();
+    MutableSpan span = trace.get(0);
+    // this should be entry but until we get the ordering resolved, it's this.
+    assertEquals("doubleEntry", span.getTag(format(DD_CODE_ORIGIN_FRAME, 0, "method")));
+  }
+
   @Test
   public void stackDepth() throws IOException, URISyntaxException {
     final String CLASS_NAME = "com.datadog.debugger.CodeOrigin04";
@@ -192,7 +246,6 @@ private List<LogProbe> codeOriginProbes(String type) {
   }
 
   private void checkResults(Class<?> testClass, String parameter, boolean includeSnapshot) {
-    Reflect.onClass(testClass).call("main", parameter).get();
     int result = Reflect.onClass(testClass).call("main", parameter).get();
     assertEquals(0, result);
     List<? extends MutableSpan> spans = traceInterceptor.getTrace();

dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CodeOrigin05.java

@@ -0,0 +1,72 @@
+package com.datadog.debugger;
+
+import datadog.trace.bootstrap.debugger.spanorigin.CodeOriginInfo;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI;
+import datadog.trace.bootstrap.instrumentation.api.ScopeSource;
+import datadog.trace.core.DDSpan;
+
+public class CodeOrigin05 {
+  private int intField = 42;
+
+  private static TracerAPI tracerAPI = AgentTracer.get();
+
+  public static int main(String arg) throws ReflectiveOperationException {
+    AgentSpan span = newSpan("main");
+    AgentScope scope = tracerAPI.activateSpan(span, ScopeSource.MANUAL);
+    if (arg.equals("debug_1")) {
+      ((DDSpan) span.getLocalRootSpan()).setTag("_dd.p.debug", "1");
+    } else if (arg.equals("debug_0")) {
+      ((DDSpan) span.getLocalRootSpan()).setTag("_dd.p.debug", "0");
+    }
+
+    fullTrace();
+
+    span.finish();
+    scope.close();
+
+    return 0;
+  }
+
+  private static void fullTrace() throws NoSuchMethodException {
+    AgentSpan span = newSpan("entry");
+    AgentScope scope = tracerAPI.activateSpan(span, ScopeSource.MANUAL);
+    entry();
+    span.finish();
+    scope.close();
+
+    span = newSpan("exit");
+    scope = tracerAPI.activateSpan(span, ScopeSource.MANUAL);
+    exit();
+    span.finish();
+    scope.close();
+  }
+
+  private static AgentSpan newSpan(String name) {
+    return tracerAPI.buildSpan("code origin tests", name).start();
+  }
+
+  public static void entry() throws NoSuchMethodException {
+    // just to fill out the method body
+    boolean dummyCode = true;
+    if (!dummyCode) {
+      dummyCode = false;
+    }
+    doubleEntry();
+  }
+
+  private static void exit() {
+    int x = 47 / 3;
+  }
+
+  public static void doubleEntry() throws NoSuchMethodException {
+    // just to fill out the method body
+    boolean dummyCode = true;
+    if (!dummyCode) {
+      dummyCode = false;
+    }
+  }
+
+}