Update the CodeOriginProbe fingerprint to not rely on a stack walk (#…

…8016)

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

@@ -3,6 +3,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.debugger.util.TimeoutChecker;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.lang.reflect.Method;
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
@@ -72,7 +73,9 @@ public interface ExceptionDebugger {
   }
 
   public interface CodeOriginRecorder {
-    String captureCodeOrigin(String signature);
+    String captureCodeOrigin(boolean entry);
+
+    String captureCodeOrigin(Method method, boolean entry);
   }
 
   private static volatile ProbeResolver probeResolver;
@@ -351,11 +354,23 @@ public static void commit(
     }
   }
 
-  public static String captureCodeOrigin(String signature) {
+  public static String captureCodeOrigin(boolean entry) {
+    try {
+      CodeOriginRecorder recorder = codeOriginRecorder;
+      if (recorder != null) {
+        return recorder.captureCodeOrigin(entry);
+      }
+    } catch (Exception ex) {
+      LOGGER.debug("Error in captureCodeOrigin: ", ex);
+    }
+    return null;
+  }
+
+  public static String captureCodeOrigin(Method method, boolean entry) {
     try {
       CodeOriginRecorder recorder = codeOriginRecorder;
       if (recorder != null) {
-        return recorder.captureCodeOrigin(signature);
+        return recorder.captureCodeOrigin(method, entry);
       }
     } catch (Exception ex) {
       LOGGER.debug("Error in captureCodeOrigin: ", ex);

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

@@ -1,29 +1,23 @@
 package datadog.trace.bootstrap.debugger.spanorigin;
 
 import static datadog.trace.bootstrap.debugger.DebuggerContext.captureCodeOrigin;
-import static java.util.Arrays.stream;
 
 import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.lang.reflect.Method;
-import java.util.stream.Collectors;
 
 public class CodeOriginInfo {
   public static void entry(Method method) {
     if (InstrumenterConfig.get().isCodeOriginEnabled()) {
-      String signature =
-          stream(method.getParameterTypes())
-              .map(Class::getTypeName)
-              .collect(Collectors.joining(", ", "(", ")"));
-      captureCodeOrigin(signature);
+      captureCodeOrigin(method, true);
     }
   }
 
   public static void exit(AgentSpan span) {
     if (InstrumenterConfig.get().isCodeOriginEnabled()) {
-      String probeId = captureCodeOrigin(null);
+      String probeId = captureCodeOrigin(false);
       if (span != null) {
-        span.getLocalRootSpan().setTag(probeId, span);
+        span.getLocalRootSpan().setTag(probeId, span.getSpanId());
       }
     }
   }

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

@@ -15,6 +15,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.util.AgentTaskScheduler;
 import datadog.trace.util.stacktrace.StackWalkerFactory;
+import java.lang.reflect.Method;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -27,63 +28,75 @@
 public class DefaultCodeOriginRecorder implements CodeOriginRecorder {
   private static final Logger LOG = LoggerFactory.getLogger(DefaultCodeOriginRecorder.class);
 
-  private final Config config;
-
   private final ConfigurationUpdater configurationUpdater;
 
   private final Map<String, CodeOriginProbe> fingerprints = new HashMap<>();
 
   private final Map<String, CodeOriginProbe> probes = new ConcurrentHashMap<>();
 
-  private final AgentTaskScheduler taskScheduler = AgentTaskScheduler.INSTANCE;
+  private final int maxUserFrames;
 
   public DefaultCodeOriginRecorder(Config config, ConfigurationUpdater configurationUpdater) {
-    this.config = config;
     this.configurationUpdater = configurationUpdater;
+    maxUserFrames = config.getDebuggerCodeOriginMaxUserFrames();
   }
 
   @Override
-  public String captureCodeOrigin(String signature) {
+  public String captureCodeOrigin(boolean entry) {
     StackTraceElement element = findPlaceInStack();
     String fingerprint = Fingerprinter.fingerprint(element);
-    if (fingerprint == null) {
-      LOG.debug("Unable to fingerprint stack trace");
-      return null;
-    }
     CodeOriginProbe probe;
 
-    AgentSpan span = AgentTracer.activeSpan();
-    if (!isAlreadyInstrumented(fingerprint)) {
-      Where where =
-          Where.of(
-              element.getClassName(),
-              element.getMethodName(),
-              signature,
-              String.valueOf(element.getLineNumber()));
-
+    if (isAlreadyInstrumented(fingerprint)) {
+      probe = fingerprints.get(fingerprint);
+    } else {
       probe =
-          new CodeOriginProbe(
-              new ProbeId(UUID.randomUUID().toString(), 0),
-              where.getSignature(),
-              where,
-              config.getDebuggerCodeOriginMaxUserFrames());
-      addFingerprint(fingerprint, probe);
-
-      installProbe(probe);
-      if (span != null) {
-        //  committing here manually so that first run probe encounters decorate the span until the
-        // instrumentation gets installed
-        probe.commit(
-            CapturedContext.EMPTY_CONTEXT, CapturedContext.EMPTY_CONTEXT, Collections.emptyList());
-      }
+          createProbe(
+              fingerprint,
+              entry,
+              Where.of(
+                  element.getClassName(),
+                  element.getMethodName(),
+                  null,
+                  String.valueOf(element.getLineNumber())));
+    }
+
+    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 = fingerprints.get(fingerprint);
+      probe = createProbe(fingerPrint, entry, Where.of(method));
     }
 
     return probe.getId();
   }
 
+  private CodeOriginProbe createProbe(String fingerPrint, boolean entry, Where where) {
+    CodeOriginProbe probe;
+    AgentSpan span = AgentTracer.activeSpan();
+
+    probe =
+        new CodeOriginProbe(
+            new ProbeId(UUID.randomUUID().toString(), 0), entry, where, maxUserFrames);
+    addFingerprint(fingerPrint, probe);
+
+    installProbe(probe);
+    // committing here manually so that first run probe encounters decorate the span until the
+    // instrumentation gets installed
+    if (span != null) {
+      probe.commit(
+          CapturedContext.EMPTY_CONTEXT, CapturedContext.EMPTY_CONTEXT, Collections.emptyList());
+    }
+    return probe;
+  }
+
   private StackTraceElement findPlaceInStack() {
     return StackWalkerFactory.INSTANCE.walk(
         stream ->
@@ -104,7 +117,8 @@ void addFingerprint(String fingerprint, CodeOriginProbe probe) {
   public String installProbe(CodeOriginProbe probe) {
     CodeOriginProbe installed = probes.putIfAbsent(probe.getId(), probe);
     if (installed == null) {
-      taskScheduler.execute(() -> configurationUpdater.accept(CODE_ORIGIN, getProbes()));
+      AgentTaskScheduler.INSTANCE.execute(
+          () -> configurationUpdater.accept(CODE_ORIGIN, getProbes()));
       return probe.getId();
     }
     return installed.getId();

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

@@ -9,6 +9,7 @@
 
 import com.datadog.debugger.agent.DebuggerAgent;
 import com.datadog.debugger.instrumentation.InstrumentationResult;
+import com.datadog.debugger.sink.DebuggerSink;
 import com.datadog.debugger.sink.Snapshot;
 import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.DebuggerContext;
@@ -28,14 +29,11 @@ public class CodeOriginProbe extends LogProbe implements ForceMethodInstrumentat
 
   public final int maxFrames;
 
-  private final String signature;
-
   private final boolean entrySpanProbe;
 
-  public CodeOriginProbe(ProbeId probeId, String signature, Where where, int maxFrames) {
+  public CodeOriginProbe(ProbeId probeId, boolean entry, Where where, int maxFrames) {
     super(LANGUAGE, probeId, null, where, MethodLocation.EXIT, null, null, true, null, null, null);
-    this.signature = signature;
-    this.entrySpanProbe = signature != null;
+    this.entrySpanProbe = entry;
     this.maxFrames = maxFrames;
   }
 
@@ -68,16 +66,19 @@ public void commit(
       return;
     }
     String snapshotId = null;
-    if (isDebuggerEnabled(span)) {
+    DebuggerSink sink = DebuggerAgent.getSink();
+    if (isDebuggerEnabled(span) && sink != null) {
       Snapshot snapshot = createSnapshot();
       if (fillSnapshot(entryContext, exitContext, caughtExceptions, snapshot)) {
         snapshotId = snapshot.getId();
         LOGGER.debug("committing code origin probe id={}, snapshot id={}", id, snapshotId);
-        commitSnapshot(snapshot, DebuggerAgent.getSink());
+        commitSnapshot(snapshot, sink);
       }
     }
     applySpanOriginTags(span, snapshotId);
-    DebuggerAgent.getSink().getProbeStatusSink().addEmitting(probeId);
+    if (sink != null) {
+      sink.getProbeStatusSink().addEmitting(probeId);
+    }
     span.getLocalRootSpan().setTag(getId(), (String) null); // clear possible span reference
   }
 
@@ -95,8 +96,8 @@ private void applySpanOriginTags(AgentSpan span, String snapshotId) {
         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 && signature != null) {
-          s.setTag(format(DD_CODE_ORIGIN_FRAME, i, "signature"), signature);
+        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);
@@ -111,6 +112,9 @@ public boolean entrySpanProbe() {
 
   /** look "back" to find exit spans that may have already come and gone */
   private AgentSpan findSpan(AgentSpan candidate) {
+    if (candidate == null) {
+      return null;
+    }
     AgentSpan span = candidate;
     AgentSpan localRootSpan = candidate.getLocalRootSpan();
     if (localRootSpan.getTag(getId()) != null) {

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

@@ -1,18 +1,22 @@
 package com.datadog.debugger.probe;
 
+import static java.util.Arrays.stream;
+
 import com.datadog.debugger.agent.Generated;
 import com.datadog.debugger.instrumentation.Types;
 import com.datadog.debugger.util.ClassFileLines;
 import com.squareup.moshi.JsonAdapter;
 import com.squareup.moshi.JsonReader;
 import com.squareup.moshi.JsonWriter;
 import java.io.IOException;
+import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.tree.MethodNode;
 
@@ -46,8 +50,17 @@ public static Where of(String typeName, String methodName, String signature, Str
     return new Where(typeName, methodName, signature, lines, null);
   }
 
+  public static Where of(Method method) {
+    return of(
+        method.getDeclaringClass().getName(),
+        method.getName(),
+        stream(method.getParameterTypes())
+            .map(Class::getTypeName)
+            .collect(Collectors.joining(", ", "(", ")")));
+  }
+
   protected static SourceLine[] sourceLines(String[] defs) {
-    if (defs == null) {
+    if (defs == null || defs.length == 0) {
       return null;
     }
     SourceLine[] lines = new SourceLine[defs.length];
@@ -72,7 +85,7 @@ public static Where convertLineToMethod(Where lineWhere, ClassFileLines classFil
             null);
       }
     }
-    throw new IllegalArgumentException("Invalid where to convert from line to method " + lineWhere);
+    return lineWhere;
   }
 
   public String getTypeName() {

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

@@ -119,7 +119,7 @@ public void stackDepth() throws IOException, URISyntaxException {
     final String CLASS_NAME = "com.datadog.debugger.CodeOrigin04";
     installProbes(
         new CodeOriginProbe(
-            CODE_ORIGIN_ID1, null, Where.of(CLASS_NAME, "exit", "()", "39"), MAX_FRAMES));
+            CODE_ORIGIN_ID1, true, Where.of(CLASS_NAME, "exit", "()", "39"), MAX_FRAMES));
 
     Class<?> testClass = compileAndLoadClass("com.datadog.debugger.CodeOrigin04");
     countFrames(testClass, 10);
@@ -139,27 +139,54 @@ private void countFrames(Class<?> testClass, int loops) {
   }
 
   @Test
-  public void testCaptureCodeOriginWithSignature() {
+  public void testCaptureCodeOriginEntry() {
     installProbes();
-    CodeOriginProbe probe = codeOriginRecorder.getProbe(codeOriginRecorder.captureCodeOrigin("()"));
+    CodeOriginProbe probe = codeOriginRecorder.getProbe(codeOriginRecorder.captureCodeOrigin(true));
     assertNotNull(probe);
     assertTrue(probe.entrySpanProbe());
   }
 
   @Test
-  public void testCaptureCodeOriginWithNullSignature() {
+  public void testCaptureCodeOriginExit() {
     installProbes();
-    CodeOriginProbe probe = codeOriginRecorder.getProbe(codeOriginRecorder.captureCodeOrigin(null));
+    CodeOriginProbe probe =
+        codeOriginRecorder.getProbe(codeOriginRecorder.captureCodeOrigin(false));
     assertNotNull(probe);
     assertFalse(probe.entrySpanProbe());
   }
 
+  @Test
+  public void testCaptureCodeOriginWithExplicitInfo()
+      throws IOException, URISyntaxException, NoSuchMethodException {
+    final String CLASS_NAME = "com.datadog.debugger.CodeOrigin04";
+    final Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    installProbes();
+    CodeOriginProbe probe =
+        codeOriginRecorder.getProbe(
+            codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true));
+    assertNotNull(probe, "The probe should have been created.");
+    assertTrue(probe.entrySpanProbe(), "Should be an entry probe.");
+  }
+
+  @Test
+  public void testDuplicateInstrumentations()
+      throws IOException, URISyntaxException, NoSuchMethodException {
+    final String CLASS_NAME = "com.datadog.debugger.CodeOrigin04";
+    final Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    installProbes();
+    String probe1 =
+        codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true);
+    String probe2 =
+        codeOriginRecorder.captureCodeOrigin(testClass.getMethod("main", int.class), true);
+    assertEquals(probe1, probe2);
+  }
+
   @NotNull
   private List<LogProbe> codeOriginProbes(String type) {
     CodeOriginProbe entry =
-        new CodeOriginProbe(CODE_ORIGIN_ID1, "()", Where.of(type, "entry", "()", "53"), MAX_FRAMES);
+        new CodeOriginProbe(CODE_ORIGIN_ID1, true, Where.of(type, "entry", "()", "53"), MAX_FRAMES);
     CodeOriginProbe exit =
-        new CodeOriginProbe(CODE_ORIGIN_ID2, null, Where.of(type, "exit", "()", "60"), MAX_FRAMES);
+        new CodeOriginProbe(CODE_ORIGIN_ID2, false, Where.of(type, "exit", "()", "60"), MAX_FRAMES);
     return new ArrayList<>(asList(entry, exit));
   }