Fix CodeOrigin for @Trace annotation

dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/DebuggerContext.java

@@ -100,7 +100,7 @@ public interface ExceptionDebugger {
   public interface CodeOriginRecorder {
     String captureCodeOrigin(boolean entry);
 
-    String captureCodeOrigin(Method method, boolean entry);
+    String captureCodeOrigin(Method method, boolean entry, boolean instrument);
   }
 
   private static volatile ProbeResolver probeResolver;
@@ -405,10 +405,14 @@ public static String captureCodeOrigin(boolean entry) {
   }
 
   public static String captureCodeOrigin(Method method, boolean entry) {
+    return captureCodeOrigin(method, entry, true);
+  }
+
+  public static String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
     try {
       CodeOriginRecorder recorder = codeOriginRecorder;
       if (recorder != null) {
-        return recorder.captureCodeOrigin(method, entry);
+        return recorder.captureCodeOrigin(method, entry, instrument);
       }
     } catch (Exception ex) {
       LOGGER.debug("Error in captureCodeOrigin: ", ex);

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/codeorigin/DefaultCodeOriginRecorder.java

@@ -71,20 +71,25 @@ public String captureCodeOrigin(boolean entry) {
               element.getMethodName(),
               null,
               String.valueOf(element.getLineNumber()));
-      probe = createProbe(fingerprint, entry, where);
+      probe = createProbe(fingerprint, entry, where, true);
 
       LOG.debug("Creating probe for location {}", where);
     }
     return probe.getId();
   }
 
   @Override
-  public String captureCodeOrigin(Method method, boolean entry) {
+  public String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
     String fingerprint = method.toString();
     CodeOriginProbe probe = probesByFingerprint.get(fingerprint);
     if (probe == null) {
-      probe = createProbe(fingerprint, entry, Where.of(method));
+      probe = createProbe(fingerprint, entry, Where.of(method), instrument);
       LOG.debug("Creating probe for method {}", fingerprint);
+    } else if (!instrument) {
+      // direct call to fill code origin info without using probe instrumentation
+      // buildLocation should be called before in order to gather location info
+      probe.commit(
+          CapturedContext.EMPTY_CONTEXT, CapturedContext.EMPTY_CONTEXT, Collections.emptyList());
     }
     return probe.getId();
   }
@@ -105,11 +110,12 @@ public void registerLogProbe(CodeOriginProbe probe) {
                     .build());
   }
 
-  private CodeOriginProbe createProbe(String fingerPrint, boolean entry, Where where) {
+  private CodeOriginProbe createProbe(
+      String fingerPrint, boolean entry, Where where, boolean instrument) {
     CodeOriginProbe probe;
     AgentSpan span = AgentTracer.activeSpan();
 
-    probe = new CodeOriginProbe(ProbeId.newId(), entry, where);
+    probe = new CodeOriginProbe(ProbeId.newId(), entry, where, instrument);
     addFingerprint(fingerPrint, probe);
     CodeOriginProbe installed = probes.putIfAbsent(probe.getId(), probe);
 

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

@@ -27,19 +27,23 @@
 public class CodeOriginProbe extends ProbeDefinition {
   private static final Logger LOGGER = LoggerFactory.getLogger(CodeOriginProbe.class);
 
+  private final boolean instrument;
   private final boolean entrySpanProbe;
-
   private String signature;
 
-  public CodeOriginProbe(ProbeId probeId, boolean entry, Where where) {
+  public CodeOriginProbe(ProbeId probeId, boolean entry, Where where, boolean instrument) {
     super(LANGUAGE, probeId, (Tag[]) null, where, MethodLocation.ENTRY);
+    this.instrument = instrument;
     this.entrySpanProbe = entry;
   }
 
   @Override
   public Status instrument(
       MethodInfo methodInfo, List<DiagnosticMessage> diagnostics, List<ProbeId> probeIds) {
-    return new CodeOriginInstrumentor(this, methodInfo, diagnostics, probeIds).instrument();
+    if (instrument) {
+      return new CodeOriginInstrumentor(this, methodInfo, diagnostics, probeIds).instrument();
+    }
+    return Status.INSTALLED;
   }
 
   @Override
@@ -55,6 +59,10 @@ public void commit(
     List<AgentSpan> agentSpans =
         entrySpanProbe ? asList(span, span.getLocalRootSpan()) : singletonList(span);
 
+    if (location == null) {
+      LOGGER.debug("Code origin probe {} has no location", id);
+      return;
+    }
     for (AgentSpan s : agentSpans) {
       if (s.getTag(DD_CODE_ORIGIN_TYPE) == null) {
         s.setTag(DD_CODE_ORIGIN_TYPE, entrySpanProbe ? "entry" : "exit");

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

@@ -105,8 +105,8 @@ public void withDebug1() throws Exception {
     final String className = "com.datadog.debugger.CodeOrigin02";
     installProbes();
     final Class<?> testClass = compileAndLoadClass(className);
-    codeOriginRecorder.captureCodeOrigin(testClass.getMethod("entry"), true);
-    codeOriginRecorder.captureCodeOrigin(testClass.getMethod("exit"), false);
+    codeOriginRecorder.captureCodeOrigin(testClass.getMethod("entry"), true, true);
+    codeOriginRecorder.captureCodeOrigin(testClass.getMethod("exit"), false, true);
     checkResults(testClass, "fullTrace", 0);
     checkResults(testClass, "debug_1", 2);
   }
@@ -117,8 +117,8 @@ public void withLogProbe() throws Exception {
     installProbes(
         createProbeBuilder(PROBE_ID, CLASS_NAME, "entry", "()").captureSnapshot(true).build());
     final Class<?> testClass = compileAndLoadClass(CLASS_NAME);
-    codeOriginRecorder.captureCodeOrigin(testClass.getMethod("entry"), true);
-    codeOriginRecorder.captureCodeOrigin(testClass.getMethod("exit"), false);
+    codeOriginRecorder.captureCodeOrigin(testClass.getMethod("entry"), true, true);
+    codeOriginRecorder.captureCodeOrigin(testClass.getMethod("exit"), false, true);
     checkResults(testClass, "debug_1", 3);
   }
 
@@ -127,10 +127,13 @@ public void doubleEntry() throws IOException, URISyntaxException {
     final String className = "com.datadog.debugger.CodeOrigin05";
 
     installProbes(
-        new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(className, "entry", "()", "53")),
-        new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(className, "exit", "()", "62")),
+        new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(className, "entry", "()", "53"), true),
+        new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(className, "exit", "()", "62"), true),
         new CodeOriginProbe(
-            CODE_ORIGIN_DOUBLE_ENTRY_ID, true, Where.of(className, "doubleEntry", "()", "66")));
+            CODE_ORIGIN_DOUBLE_ENTRY_ID,
+            true,
+            Where.of(className, "doubleEntry", "()", "66"),
+            true));
     final Class<?> testClass = compileAndLoadClass(className);
     checkResults(testClass, "fullTrace", 0);
     List<? extends MutableSpan> trace = traceInterceptor.getTrace();
@@ -143,7 +146,7 @@ public void doubleEntry() throws IOException, URISyntaxException {
   public void stackDepth() throws IOException, URISyntaxException {
     final String CLASS_NAME = "com.datadog.debugger.CodeOrigin04";
     installProbes(
-        new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(CLASS_NAME, "exit", "()", "39")));
+        new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(CLASS_NAME, "exit", "()", "39"), true));
 
     Class<?> testClass = compileAndLoadClass("com.datadog.debugger.CodeOrigin04");
     countFrames(testClass, 10);
@@ -187,7 +190,8 @@ public void testCaptureCodeOriginWithExplicitInfo()
     installProbes();
     CodeOriginProbe probe =
         codeOriginRecorder.getProbe(
-            codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true));
+            codeOriginRecorder.captureCodeOrigin(
+                testClass.getMethod("main", int.class), true, true));
     assertNotNull(probe, "The probe should have been created.");
     assertTrue(probe.entrySpanProbe(), "Should be an entry probe.");
   }
@@ -199,18 +203,18 @@ public void testDuplicateInstrumentations()
     final Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     installProbes();
     String probe1 =
-        codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true);
+        codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true, true);
     String probe2 =
-        codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true);
+        codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true, true);
     assertEquals(probe1, probe2);
   }
 
   @NotNull
   private CodeOriginProbe[] codeOriginProbes(String type) {
     CodeOriginProbe entry =
-        new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(type, "entry", "()", "53"));
+        new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(type, "entry", "()", "53"), true);
     CodeOriginProbe exit =
-        new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(type, "exit", "()", "60"));
+        new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(type, "exit", "()", "60"), true);
     return new CodeOriginProbe[] {entry, exit};
   }
 

dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/GrpcCodeOriginTest.groovy

@@ -260,7 +260,7 @@ abstract class GrpcCodeOriginTest extends VersionedNamingTestBase {
         }
 
         @Override
-        String captureCodeOrigin(Method method, boolean entry) {
+        String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
           invoked = true
           return "done"
         }

dd-java-agent/instrumentation/micronaut/http-server-netty-4.0/src/test/groovy/MicronautTest.groovy

@@ -35,7 +35,7 @@ class MicronautTest extends HttpServerTest<Object> {
         }
 
         @Override
-        String captureCodeOrigin(Method method, boolean entry) {
+        String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
           invoked = true
           return "done"
         }

dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/TraceAdvice.java

@@ -3,6 +3,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.instrumentation.trace_annotation.TraceDecorator.DECORATE;
 
+import datadog.trace.bootstrap.debugger.DebuggerContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import java.lang.invoke.MethodType;
 import java.lang.reflect.Method;
@@ -13,7 +14,9 @@ public class TraceAdvice {
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static AgentScope onEnter(@Advice.Origin final Method method) {
-    return activateSpan(DECORATE.startMethodSpan(method));
+    AgentScope agentScope = activateSpan(DECORATE.startMethodSpan(method));
+    DebuggerContext.captureCodeOrigin(method, true, false);
+    return agentScope;
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/TraceDecorator.java

@@ -4,7 +4,6 @@
 
 import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.api.Trace;
-import datadog.trace.bootstrap.debugger.spanorigin.CodeOriginInfo;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.decorator.AsyncResultDecorator;
@@ -91,7 +90,6 @@ public AgentSpan startMethodSpan(Method method) {
     afterStart(span);
     span.setResourceName(resourceName);
 
-    CodeOriginInfo.entry(method);
     if (measured || InstrumenterConfig.get().isMethodMeasured(method)) {
       span.setMeasured(true);
     }

dd-smoke-tests/debugger-integration-tests/src/test/java/datadog/smoketest/CodeOriginIntegrationTest.java

@@ -0,0 +1,60 @@
+package datadog.smoketest;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import datadog.trace.api.DDTags;
+import datadog.trace.test.agent.decoder.DecodedSpan;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+public class CodeOriginIntegrationTest extends ServerAppDebuggerIntegrationTest {
+
+  private static final String DD_CODE_ORIGIN_FRAMES_0_FILE =
+      String.format(DDTags.DD_CODE_ORIGIN_FRAME, 0, "file");
+  private static final String DD_CODE_ORIGIN_FRAMES_0_METHOD =
+      String.format(DDTags.DD_CODE_ORIGIN_FRAME, 0, "method");
+  private static final String DD_CODE_ORIGIN_FRAMES_0_SIGNATURE =
+      String.format(DDTags.DD_CODE_ORIGIN_FRAME, 0, "signature");
+  private static final String DD_CODE_ORIGIN_FRAMES_0_LINE =
+      String.format(DDTags.DD_CODE_ORIGIN_FRAME, 0, "line");
+
+  @Override
+  protected ProcessBuilder createProcessBuilder(Path logFilePath, String... params) {
+    List<String> commandParams = getDebuggerCommandParams();
+    commandParams.add("-Ddd.trace.enabled=true"); // explicitly enable tracer
+    commandParams.add("-Ddd.code.origin.for.spans.enabled=true");
+    return ProcessBuilderHelper.createProcessBuilder(
+        commandParams, logFilePath, getAppClass(), params);
+  }
+
+  @Test
+  @DisplayName("testCodeOriginTraceAnnotation")
+  void testCodeOriginTraceAnnotation() throws Exception {
+    execute(appUrl, TRACED_METHOD_NAME);
+    waitForInstrumentation(appUrl);
+    execute(appUrl, TRACED_METHOD_NAME);
+    AtomicBoolean codeOrigin = new AtomicBoolean();
+    registerTraceListener(
+        decodedTrace -> {
+          for (DecodedSpan span : decodedTrace.getSpans()) {
+            if (isTracedFullMethodSpan(span)) {
+              if (span.getMeta().containsKey(DDTags.DD_CODE_ORIGIN_TYPE)) {
+                assertEquals("entry", span.getMeta().get(DDTags.DD_CODE_ORIGIN_TYPE));
+                assertEquals(
+                    "ServerDebuggerTestApplication.java",
+                    span.getMeta().get(DD_CODE_ORIGIN_FRAMES_0_FILE));
+                assertEquals("runTracedMethod", span.getMeta().get(DD_CODE_ORIGIN_FRAMES_0_METHOD));
+                assertEquals(
+                    "(java.lang.String)", span.getMeta().get(DD_CODE_ORIGIN_FRAMES_0_SIGNATURE));
+                assertEquals("133", span.getMeta().get(DD_CODE_ORIGIN_FRAMES_0_LINE));
+                codeOrigin.set(true);
+              }
+            }
+          }
+        });
+    processRequests(codeOrigin::get);
+  }
+}