Fix the instrumentation to instrument at the exit line (#8639)

Author: evanchooly

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CodeOriginInstrumentor.java

@@ -4,13 +4,17 @@
 import static com.datadog.debugger.instrumentation.ASMHelper.ldc;
 import static com.datadog.debugger.instrumentation.Types.DEBUGGER_CONTEXT_TYPE;
 import static com.datadog.debugger.instrumentation.Types.STRING_TYPE;
+import static java.lang.Integer.parseInt;
 
 import com.datadog.debugger.instrumentation.InstrumentationResult.Status;
+import com.datadog.debugger.probe.CodeOriginProbe;
 import com.datadog.debugger.probe.ProbeDefinition;
 import datadog.trace.bootstrap.debugger.ProbeId;
 import java.util.List;
 import org.objectweb.asm.Type;
+import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.InsnList;
+import org.objectweb.asm.tree.LabelNode;
 
 public class CodeOriginInstrumentor extends Instrumentor {
   public CodeOriginInstrumentor(
@@ -26,11 +30,21 @@ public Status instrument() {
     InsnList insnList = new InsnList();
 
     ldc(insnList, probeIds.get(0).getEncodedId());
-
     invokeStatic(insnList, DEBUGGER_CONTEXT_TYPE, "codeOrigin", Type.VOID_TYPE, STRING_TYPE);
-
-    methodNode.instructions.insert(methodEnterLabel, insnList);
+    methodNode.instructions.insert(findInsertionPoint(), insnList);
 
     return InstrumentationResult.Status.INSTALLED;
   }
+
+  private AbstractInsnNode findInsertionPoint() {
+    CodeOriginProbe probe = (CodeOriginProbe) definition;
+    List<String> lines = probe.getLocation().getLines();
+    if (!probe.entrySpanProbe() && lines != null && !lines.isEmpty()) {
+      LabelNode lineLabel = classFileLines.getLineLabel(parseInt(lines.get(0)));
+      if (lineLabel != null) {
+        return lineLabel.getNext();
+      }
+    }
+    return methodEnterLabel;
+  }
 }

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

@@ -82,14 +82,14 @@ public boolean entrySpanProbe() {
   public void buildLocation(MethodInfo methodInfo) {
     String type = where.getTypeName();
     String method = where.getMethodName();
-    List<String> lines = null;
+    List<String> lines = where.getLines() != null ? asList(where.getLines()) : null;
 
     String file = where.getSourceFile();
 
     if (methodInfo != null) {
       type = methodInfo.getTypeName();
       method = methodInfo.getMethodName();
-      if (methodInfo.getMethodStart() != -1) {
+      if (entrySpanProbe || where.getLines() == null) {
         lines = singletonList(String.valueOf(methodInfo.getMethodStart()));
       }
       if (file == null) {

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

@@ -14,6 +14,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static utils.InstrumentationTestHelper.compileAndLoadClass;
+import static utils.InstrumentationTestHelper.getLineForLineProbe;
 import static utils.TestHelper.setFieldInConfig;
 
 import com.datadog.debugger.agent.CapturingTestBase;
@@ -130,12 +131,25 @@ 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"), true),
-        new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(className, "exit", "()", "62"), true),
+        new CodeOriginProbe(
+            CODE_ORIGIN_ID1,
+            true,
+            Where.of(
+                className, "entry", "()", "" + getLineForLineProbe(className, CODE_ORIGIN_ID1)),
+            true),
+        new CodeOriginProbe(
+            CODE_ORIGIN_ID2,
+            false,
+            Where.of(className, "exit", "()", "" + getLineForLineProbe(className, CODE_ORIGIN_ID2)),
+            true),
         new CodeOriginProbe(
             CODE_ORIGIN_DOUBLE_ENTRY_ID,
             true,
-            Where.of(className, "doubleEntry", "()", "66"),
+            Where.of(
+                className,
+                "doubleEntry",
+                "()",
+                "" + getLineForLineProbe(className, CODE_ORIGIN_DOUBLE_ENTRY_ID)),
             true));
     final Class<?> testClass = compileAndLoadClass(className);
     checkResults(testClass, "fullTrace", 0);

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

@@ -49,21 +49,21 @@ private static AgentSpan newSpan(String name) {
 
   public static void entry() throws NoSuchMethodException {
     // just to fill out the method body
-    boolean dummyCode = true;
+    boolean dummyCode = true; // code origin 1
     if (!dummyCode) {
       dummyCode = false;
     }
     doubleEntry();
   }
 
   private static void exit() {
-    int x = 47 / 3;
+    int x = 47 / 3; // code origin 2
   }
 
   public static void doubleEntry() throws NoSuchMethodException {
     // just to fill out the method body
     boolean dummyCode = true;
-    if (!dummyCode) {
+    if (!dummyCode) { // double entry code origin
       dummyCode = false;
     }
   }