Clean up the initial code origin set up instrumentation

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

@@ -116,7 +116,7 @@ public interface ExceptionDebugger {
   public interface CodeOriginRecorder {
     String captureCodeOrigin(boolean entry);
 
-    String captureCodeOrigin(Method method, boolean entry, boolean instrument);
+    String captureCodeOrigin(Method method, boolean entry);
   }
 
   private static volatile ProductConfigUpdater productConfigUpdater;
@@ -466,32 +466,28 @@ public static void commit(
     }
   }
 
-  public static String captureCodeOrigin(boolean entry) {
+  public static void marker() {}
+
+  public static void captureCodeOrigin(boolean entry) {
     try {
       CodeOriginRecorder recorder = codeOriginRecorder;
       if (recorder != null) {
-        return recorder.captureCodeOrigin(entry);
+        recorder.captureCodeOrigin(entry);
       }
     } catch (Exception ex) {
       LOGGER.debug("Error in captureCodeOrigin: ", ex);
     }
-    return null;
-  }
-
-  public static String captureCodeOrigin(Method method, boolean entry) {
-    return captureCodeOrigin(method, entry, true);
   }
 
-  public static String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
+  public static void captureCodeOrigin(Method method, boolean entry) {
     try {
       CodeOriginRecorder recorder = codeOriginRecorder;
       if (recorder != null) {
-        return recorder.captureCodeOrigin(method, entry, instrument);
+        recorder.captureCodeOrigin(method, entry);
       }
     } catch (Exception ex) {
       LOGGER.debug("Error in captureCodeOrigin: ", ex);
     }
-    return null;
   }
 
   public static void handleException(Throwable t, AgentSpan span) {

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

@@ -45,7 +45,7 @@ public class DefaultCodeOriginRecorder implements CodeOriginRecorder {
 
   private final int maxUserFrames;
 
-  private AgentTaskScheduler scheduler = AgentTaskScheduler.INSTANCE;
+  private AgentTaskScheduler scheduler;
 
   public DefaultCodeOriginRecorder(Config config, ConfigurationUpdater configurationUpdater) {
     this(config, configurationUpdater, AgentTaskScheduler.INSTANCE);
@@ -70,51 +70,44 @@ public String captureCodeOrigin(boolean entry) {
               element.getMethodName(),
               null,
               String.valueOf(element.getLineNumber()));
-      probe = createProbe(fingerprint, entry, where, true);
+      probe = createProbe(fingerprint, entry, where);
 
       LOG.debug("Creating probe for location {}", where);
     }
     return probe.getId();
   }
 
   @Override
-  public String captureCodeOrigin(Method method, boolean entry, boolean instrument) {
+  public String captureCodeOrigin(Method method, boolean entry) {
     String fingerprint = method.toString();
     CodeOriginProbe probe = probesByFingerprint.get(fingerprint);
     if (probe == null) {
-      probe = createProbe(fingerprint, entry, Where.of(method), instrument);
+      probe = createProbe(fingerprint, entry, Where.of(method));
       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();
   }
 
   public void registerLogProbe(CodeOriginProbe probe) {
-    LogProbe logProbe =
-        logProbes.computeIfAbsent(
-            probe.getId(),
-            key ->
-                new Builder()
-                    .language(probe.getLanguage())
-                    .probeId(ProbeId.newId())
-                    .where(probe.getWhere())
-                    .evaluateAt(probe.getEvaluateAt())
-                    .captureSnapshot(true)
-                    .tags("session_id:*")
-                    .snapshotProcessor(new CodeOriginSnapshotConsumer(probe.entrySpanProbe()))
-                    .build());
+    logProbes.computeIfAbsent(
+        probe.getId(),
+        key ->
+            new Builder()
+                .language(probe.getLanguage())
+                .probeId(ProbeId.newId())
+                .where(probe.getWhere())
+                .evaluateAt(probe.getEvaluateAt())
+                .captureSnapshot(true)
+                .tags("session_id:*")
+                .snapshotProcessor(new CodeOriginSnapshotConsumer(probe.entrySpanProbe()))
+                .build());
   }
 
-  private CodeOriginProbe createProbe(
-      String fingerPrint, boolean entry, Where where, boolean instrument) {
+  private CodeOriginProbe createProbe(String fingerPrint, boolean entry, Where where) {
     CodeOriginProbe probe;
     AgentSpan span = AgentTracer.activeSpan();
 
-    probe = new CodeOriginProbe(ProbeId.newId(), entry, where, instrument);
+    probe = new CodeOriginProbe(ProbeId.newId(), entry, where);
     addFingerprint(fingerPrint, probe);
     CodeOriginProbe installed = probes.putIfAbsent(probe.getId(), probe);
 

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

@@ -1,20 +1,29 @@
 package com.datadog.debugger.instrumentation;
 
 import static com.datadog.debugger.instrumentation.Types.REFLECTIVE_FIELD_VALUE_RESOLVER_TYPE;
+import static java.lang.String.format;
 import static org.objectweb.asm.Type.getMethodDescriptor;
 import static org.objectweb.asm.Type.getObjectType;
 
 import com.datadog.debugger.agent.Generated;
 import datadog.trace.util.Strings;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.io.StringWriter;
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
+import java.util.StringJoiner;
 import java.util.stream.Collectors;
+import org.objectweb.asm.ClassWriter;
 import org.objectweb.asm.Opcodes;
-import org.objectweb.asm.Type;
 import org.objectweb.asm.signature.SignatureReader;
 import org.objectweb.asm.signature.SignatureVisitor;
 import org.objectweb.asm.tree.AbstractInsnNode;
@@ -24,10 +33,14 @@
 import org.objectweb.asm.tree.InsnList;
 import org.objectweb.asm.tree.InsnNode;
 import org.objectweb.asm.tree.LdcInsnNode;
+import org.objectweb.asm.tree.LineNumberNode;
 import org.objectweb.asm.tree.LocalVariableNode;
 import org.objectweb.asm.tree.MethodInsnNode;
 import org.objectweb.asm.tree.MethodNode;
 import org.objectweb.asm.tree.TypeInsnNode;
+import org.objectweb.asm.util.Printer;
+import org.objectweb.asm.util.Textifier;
+import org.objectweb.asm.util.TraceClassVisitor;
 
 /** Helper class for bytecode generation */
 public class ASMHelper {
@@ -36,6 +49,45 @@ public class ASMHelper {
   public static final Type STRING_TYPE = new Type(Types.STRING_TYPE);
   public static final Type LONG_TYPE = new Type(org.objectweb.asm.Type.LONG_TYPE);
 
+  public static void dumpMethod(ClassNode classNode, String method, String suffix) {
+    String content = extractMethod(classNode, method);
+    File output;
+    int count = 0;
+    do {
+      output = new File(format("build/%s-%d-%s.txt", method, count++, suffix));
+    } while (output.exists());
+    output.getParentFile().mkdirs();
+    try (PrintWriter writer = new PrintWriter(output)) {
+      writer.println(content);
+      String absolutePath = output.getAbsolutePath();
+      absolutePath = absolutePath.substring(0, absolutePath.lastIndexOf('.'));
+      absolutePath += ".class";
+      ClassWriter classWriter = new ClassWriter(0);
+      classNode.accept(classWriter);
+      try (FileOutputStream stream = new FileOutputStream(absolutePath)) {
+        stream.write(classWriter.toByteArray());
+      }
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  @SuppressForbidden
+  static String extractMethod(ClassNode classNode, String method) {
+    StringJoiner joiner = new StringJoiner("\n");
+    joiner.add("Class: " + classNode.name);
+    StringWriter writer = new StringWriter();
+    classNode.accept(new TraceClassVisitor(null, new Textifier(), new PrintWriter(writer)));
+    List<String> strings = Arrays.asList(writer.toString().split("\n"));
+    for (int i = 0; i < strings.size(); i++) {
+      if (strings.get(i).matches(format(".*(private|public).* %s\\(.*", method))) {
+        while (!strings.get(i).equals(""))
+          joiner.add(String.format("[%3d] %s", i, strings.get(i++)));
+      }
+    }
+    return joiner.toString();
+  }
+
   public static void invokeInterface(
       InsnList insnList,
       org.objectweb.asm.Type owner,
@@ -322,6 +374,18 @@ public static boolean isStoreCompatibleType(
     return previousSort == currentSort;
   }
 
+  public static String toString(AbstractInsnNode node) {
+    String opcode = node.getOpcode() >= 0 ? Printer.OPCODES[node.getOpcode()] : node.toString();
+    if (node instanceof LineNumberNode) {
+      return String.format("LineNumber: %s", ((LineNumberNode) node).line);
+    } else if (node instanceof MethodInsnNode) {
+      MethodInsnNode method = (MethodInsnNode) node;
+      return String.format("%s: [%s] %s", opcode, method.name, method.desc);
+    } else {
+      return opcode;
+    }
+  }
+
   private static int widenIntType(int sort) {
     switch (sort) {
       case org.objectweb.asm.Type.BOOLEAN:

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

@@ -5,35 +5,81 @@
 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 static java.lang.String.format;
 
 import com.datadog.debugger.instrumentation.InstrumentationResult.Status;
 import com.datadog.debugger.probe.CodeOriginProbe;
 import com.datadog.debugger.probe.ProbeDefinition;
+import datadog.trace.bootstrap.debugger.DebuggerContext;
 import datadog.trace.bootstrap.debugger.ProbeId;
 import java.util.List;
+import java.util.ListIterator;
 import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.InsnList;
 import org.objectweb.asm.tree.LabelNode;
+import org.objectweb.asm.tree.MethodInsnNode;
 
 public class CodeOriginInstrumentor extends Instrumentor {
+  private static final String CAPTURE;
+  private static final String MARKER;
+
+  static {
+    String className = DebuggerContext.class.getName().replace('.', '/');
+    CAPTURE = format("%s#%s", className, "captureCodeOrigin");
+    MARKER = format("%s#%s", className, "marker");
+  }
+
   public CodeOriginInstrumentor(
-      ProbeDefinition definition,
-      MethodInfo methodInfo,
-      List<DiagnosticMessage> diagnostics,
-      List<ProbeId> probeIds) {
-    super(definition, methodInfo, diagnostics, probeIds);
+      ProbeDefinition definition, MethodInfo methodInfo, List<ProbeId> probeIds) {
+    super(definition, methodInfo, null, probeIds);
   }
 
   @Override
   public Status instrument() {
-    InsnList insnList = new InsnList();
+    AbstractInsnNode insertionPoint = stripSetup();
+    methodNode.instructions.insert(
+        insertionPoint != null ? insertionPoint : findInsertionPoint(), codeOriginCall());
 
+    return Status.INSTALLED;
+  }
+
+  private static String buildDescription(AbstractInsnNode node) {
+    if (!(node instanceof MethodInsnNode)) return "";
+    MethodInsnNode method = (MethodInsnNode) node;
+    return format("%s#%s", method.owner, method.name);
+  }
+
+  private InsnList codeOriginCall() {
+    InsnList insnList = new InsnList();
     ldc(insnList, probeIds.get(0).getEncodedId());
     invokeStatic(insnList, DEBUGGER_CONTEXT_TYPE, "codeOrigin", Type.VOID_TYPE, STRING_TYPE);
-    methodNode.instructions.insert(findInsertionPoint(), insnList);
+    return insnList;
+  }
 
-    return InstrumentationResult.Status.INSTALLED;
+  private AbstractInsnNode stripSetup() {
+    try {
+      ListIterator<AbstractInsnNode> iterator = methodNode.instructions.iterator();
+      AbstractInsnNode insertionPoint = null;
+      InsnList list = new InsnList();
+      while (iterator.hasNext()) {
+        AbstractInsnNode next = iterator.next();
+        if (buildDescription(next).equals(MARKER)) {
+          insertionPoint = next.getPrevious();
+          while (iterator.hasNext() && !buildDescription(next = iterator.next()).equals(CAPTURE)) {}
+
+          if (!iterator.hasNext()) {
+            return null;
+          }
+        } else {
+          list.add(next);
+        }
+      }
+      methodNode.instructions = list;
+      return insertionPoint;
+    } catch (Exception e) {
+      return null;
+    }
   }
 
   private AbstractInsnNode findInsertionPoint() {

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

@@ -26,23 +26,18 @@
 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, boolean instrument) {
+  public CodeOriginProbe(ProbeId probeId, boolean entry, Where where) {
     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) {
-    if (instrument) {
-      return new CodeOriginInstrumentor(this, methodInfo, diagnostics, probeIds).instrument();
-    }
-    return Status.INSTALLED;
+    return new CodeOriginInstrumentor(this, methodInfo, probeIds).instrument();
   }
 
   @Override

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

@@ -0,0 +1,7 @@
+package com.datadog.debugger.origin;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Target;
+
+@Target(ElementType.METHOD)
+public @interface CodeOrigin {}