Add code origin support to the grpc server integration

dd-java-agent/agent-bootstrap/build.gradle

@@ -20,7 +20,7 @@ dependencies {
   api project(':dd-trace-api')
   api project(':internal-api')
   api project(':internal-api:internal-api-9')
-  api project(':dd-java-agent:agent-logging')
+  api project(':dd-java-agent:agent-debugger:debugger-bootstrap')
   api libs.slf4j
   // ^ Generally a bad idea for libraries, but we're shadowing.
 

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

@@ -7,6 +7,7 @@
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Function;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -73,6 +74,9 @@ public interface ExceptionDebugger {
 
   public interface CodeOriginRecorder {
     String captureCodeOrigin(String signature);
+
+    String captureCodeOrigin(
+        String name, Class<?> target, String method, Class<?>[] types, boolean entry);
   }
 
   private static volatile ProbeResolver probeResolver;
@@ -363,6 +367,31 @@ public static String captureCodeOrigin(String signature) {
     return null;
   }
 
+  public static String captureCodeOrigin(Function<CodeOriginRecorder, String> function) {
+    try {
+      CodeOriginRecorder recorder = codeOriginRecorder;
+      if (recorder != null) {
+        return function.apply(recorder);
+      }
+    } catch (Exception ex) {
+      LOGGER.debug("Error in captureCodeOrigin: ", ex);
+    }
+    return null;
+  }
+
+  public static String captureCodeOrigin(
+      String name, Class<?> target, String method, Class<?>[] types, boolean entry) {
+    try {
+      CodeOriginRecorder recorder = codeOriginRecorder;
+      if (recorder != null) {
+        return recorder.captureCodeOrigin(name, target, method, types, entry);
+      }
+    } catch (Exception ex) {
+      LOGGER.debug("Error in captureCodeOrigin: ", ex);
+    }
+    return null;
+  }
+
   public static void handleException(Throwable t, AgentSpan span) {
     try {
       ExceptionDebugger exDebugger = exceptionDebugger;

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

@@ -19,11 +19,17 @@ public static void entry(Method method) {
     }
   }
 
+  public static void entry(String name, Class<?> target, String method, Class<?>[] types) {
+    if (InstrumenterConfig.get().isCodeOriginEnabled()) {
+      captureCodeOrigin(name, target, method, types, true);
+    }
+  }
+
   public static void exit(AgentSpan span) {
     if (InstrumenterConfig.get().isCodeOriginEnabled()) {
-      String probeId = captureCodeOrigin(null);
+      String probeId = captureCodeOrigin((String) null);
       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

@@ -1,6 +1,7 @@
 package com.datadog.debugger.codeorigin;
 
 import static com.datadog.debugger.agent.ConfigurationAcceptor.Source.CODE_ORIGIN;
+import static java.util.Arrays.stream;
 
 import com.datadog.debugger.agent.ConfigurationUpdater;
 import com.datadog.debugger.exception.Fingerprinter;
@@ -21,69 +22,92 @@
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 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) {
     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(),
+    if (isAlreadyInstrumented(fingerprint)) {
+      probe = fingerprints.get(fingerprint);
+    } else {
+      probe =
+          createProbe(
               signature,
-              String.valueOf(element.getLineNumber()));
+              Where.of(
+                  element.getClassName(),
+                  element.getMethodName(),
+                  signature,
+                  String.valueOf(element.getLineNumber())));
+    }
 
-      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());
-      }
+    return probe.getId();
+  }
 
+  @Override
+  public String captureCodeOrigin(
+      String name, Class<?> target, String method, Class<?>[] types, boolean entry) {
+    CodeOriginProbe probe;
+
+    if (isAlreadyInstrumented(name)) {
+      probe = fingerprints.get(name);
     } else {
-      probe = fingerprints.get(fingerprint);
+      probe =
+          createProbe(
+              name,
+              Where.of(
+                  target.getName(),
+                  method,
+                  stream(types)
+                      .map(Class::getTypeName)
+                      .collect(Collectors.joining(", ", "(", ")"))));
     }
 
     return probe.getId();
   }
 
+  private CodeOriginProbe createProbe(String fingerPrint, Where where) {
+    CodeOriginProbe probe;
+    AgentSpan span = AgentTracer.activeSpan();
+
+    probe =
+        new CodeOriginProbe(
+            new ProbeId(UUID.randomUUID().toString(), 0),
+            where.getSignature(),
+            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 +128,10 @@ 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()));
+      if (configurationUpdater != null) {
+        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/exception/Fingerprinter.java

@@ -42,6 +42,17 @@ public static String fingerprint(Throwable t, ClassNameFilter classNameFiltering
     return bytesToHex(digest.digest());
   }
 
+  public static String fingerprint(CharSequence resourceName) {
+    try {
+      MessageDigest digest = MessageDigest.getInstance("SHA-256");
+      digest.update(resourceName.toString().getBytes());
+      return bytesToHex(digest.digest());
+    } catch (NoSuchAlgorithmException e) {
+      LOGGER.debug("Unable to find digest algorithm SHA-256", e);
+      return null;
+    }
+  }
+
   public static String fingerprint(StackTraceElement element) {
     try {
       MessageDigest digest = MessageDigest.getInstance("SHA-256");

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;
@@ -68,16 +69,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
   }
 
@@ -111,6 +115,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

@@ -47,7 +47,7 @@ public static Where of(String typeName, String methodName, String signature, Str
   }
 
   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 +72,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/agent/CapturingTestBase.java

@@ -3,6 +3,7 @@
 import static com.datadog.debugger.util.MoshiSnapshotHelper.NOT_CAPTURED_REASON;
 import static com.datadog.debugger.util.MoshiSnapshotTestHelper.VALUE_ADAPTER;
 import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
+import static datadog.trace.util.AgentThreadFactory.AgentThread.TASK_SCHEDULER;
 import static java.util.Arrays.asList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -30,6 +31,7 @@
 import datadog.trace.bootstrap.debugger.ProbeImplementation;
 import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
 import datadog.trace.bootstrap.debugger.util.Redaction;
+import datadog.trace.util.AgentTaskScheduler;
 import java.io.File;
 import java.io.IOException;
 import java.lang.instrument.ClassFileTransformer;
@@ -351,6 +353,15 @@ public static LogProbe.Builder createProbeBuilder(
 
   protected TestSnapshotListener installProbes(
       Configuration configuration, ProbeDefinition... probes) {
+
+    AgentTaskScheduler.INSTANCE =
+        new AgentTaskScheduler(TASK_SCHEDULER) {
+          @Override
+          public void execute(final Runnable target) {
+            target.run();
+          }
+        };
+
     config = mock(Config.class);
     when(config.isDebuggerEnabled()).thenReturn(true);
     when(config.isDebuggerClassFileDumpEnabled()).thenReturn(true);

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

@@ -154,6 +154,33 @@ public void testCaptureCodeOriginWithNullSignature() {
     assertFalse(probe.entrySpanProbe());
   }
 
+  @Test
+  public void testCaptureCodeOriginWithExplicitInfo() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "com.datadog.debugger.CodeOrigin04";
+    final Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    installProbes();
+    CodeOriginProbe probe =
+        codeOriginRecorder.getProbe(
+            codeOriginRecorder.captureCodeOrigin(
+                "explicit", testClass, "main", new Class[] {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 {
+    final String CLASS_NAME = "com.datadog.debugger.CodeOrigin04";
+    final Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    installProbes();
+    String probe1 =
+        codeOriginRecorder.captureCodeOrigin(
+            "explicit", testClass, "main", new Class[] {int.class}, true);
+    String probe2 =
+        codeOriginRecorder.captureCodeOrigin(
+            "explicit", testClass, "main", new Class[] {int.class}, true);
+    assertEquals(probe1, probe2);
+  }
+
   @NotNull
   private List<LogProbe> codeOriginProbes(String type) {
     CodeOriginProbe entry =

dd-java-agent/instrumentation/grpc-1.5/build.gradle

@@ -39,6 +39,8 @@ dependencies {
   testImplementation group: 'io.grpc', name: 'grpc-protobuf', version: grpcVersion
   testImplementation group: 'io.grpc', name: 'grpc-stub', version: grpcVersion
   testImplementation group: 'javax.annotation', name: 'javax.annotation-api', version: '1.3.2'
+  testImplementation project(':dd-java-agent:agent-debugger')
+  testImplementation libs.bundles.mockito
 
   latestDepTestImplementation sourceSets.test.output // include the protobuf generated classes
   latestDepTestCompileOnly group: 'io.grpc', name: 'grpc-core', version: '1.+'