Fix CodeOrigin for interface endpoints

dd-java-agent/instrumentation/datadog/dynamic-instrumentation/span-origin/src/main/java/datadog/trace/instrumentation/codeorigin/CodeOriginInstrumentation.java

@@ -1,5 +1,10 @@
 package datadog.trace.instrumentation.codeorigin;
 
+import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule.Tracing;
 import datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers;
@@ -11,6 +16,7 @@
 import net.bytebuddy.description.NamedElement;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.ElementMatchers;
 
 public abstract class CodeOriginInstrumentation extends Tracing
     implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
@@ -37,13 +43,28 @@ public String hierarchyMarkerType() {
 
   @Override
   public ElementMatcher<TypeDescription> hierarchyMatcher() {
-    return HierarchyMatchers.declaresMethod(HierarchyMatchers.isAnnotatedWith(matcher));
+    ElementMatcher.Junction<TypeDescription> matcher =
+        HierarchyMatchers.declaresMethod(HierarchyMatchers.isAnnotatedWith(this.matcher));
+    if (InstrumenterConfig.get().isCodeOriginInterfaceSupport()) {
+      matcher =
+          matcher.or(
+              HierarchyMatchers.implementsInterface(
+                  HierarchyMatchers.declaresMethod(
+                      HierarchyMatchers.isAnnotatedWith(this.matcher))));
+    }
+    return matcher;
   }
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         HierarchyMatchers.isAnnotatedWith(matcher),
         "datadog.trace.instrumentation.codeorigin.EntrySpanOriginAdvice");
+    if (InstrumenterConfig.get().isCodeOriginInterfaceSupport()) {
+      transformer.applyAdvice(
+          ElementMatchers.isDeclaredBy(
+              hasSuperType(isInterface().and(declaresMethod(isAnnotatedWith(matcher))))),
+          "datadog.trace.instrumentation.codeorigin.EntrySpanOriginAdvice");
+    }
   }
 }

dd-smoke-tests/debugger-integration-tests/src/main/java/datadog/smoketest/debugger/controller/InterfacedController.java

@@ -0,0 +1,18 @@
+package datadog.smoketest.debugger.controller;
+
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+
+@RestController
+public class InterfacedController implements InterfaceApi {
+
+  @Override
+  public String process() {
+    return "OK";
+  }
+}
+
+interface InterfaceApi {
+  @RequestMapping("/process")
+  String process();
+}

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

@@ -0,0 +1,79 @@
+package datadog.smoketest;
+
+import com.datadog.debugger.probe.LogProbe;
+import datadog.trace.agent.test.utils.PortUtils;
+import datadog.trace.util.TagsHelper;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.List;
+import java.util.concurrent.locks.LockSupport;
+import okhttp3.OkHttpClient;
+import okhttp3.Request;
+
+public class SpringBasedIntegrationTest extends BaseIntegrationTest {
+  protected static final String DEBUGGER_TEST_APP_CLASS =
+      "datadog.smoketest.debugger.SpringBootTestApplication";
+
+  @Override
+  protected String getAppClass() {
+    return DEBUGGER_TEST_APP_CLASS;
+  }
+
+  @Override
+  protected String getAppId() {
+    return TagsHelper.sanitize("SpringBootTestApplication");
+  }
+
+  protected String startSpringApp(List<LogProbe> logProbes) throws Exception {
+    return startSpringApp(logProbes, false);
+  }
+
+  protected String startSpringApp(List<LogProbe> logProbes, boolean enableProcessTags)
+      throws Exception {
+    setCurrentConfiguration(createConfig(logProbes));
+    String httpPort = String.valueOf(PortUtils.randomOpenPort());
+    ProcessBuilder processBuilder = createProcessBuilder(logFilePath, "--server.port=" + httpPort);
+    if (!enableProcessTags) {
+      processBuilder.environment().put("DD_EXPERIMENTAL_PROPAGATE_PROCESS_TAGS_ENABLED", "false");
+    }
+    targetProcess = processBuilder.start();
+    // assert in logs app started
+    waitForSpecificLogLine(
+        logFilePath,
+        "datadog.smoketest.debugger.SpringBootTestApplication - Started SpringBootTestApplication",
+        Duration.ofMillis(100),
+        Duration.ofSeconds(30));
+    return httpPort;
+  }
+
+  protected void sendRequest(String httpPort, String urlPath) {
+    OkHttpClient client = new OkHttpClient.Builder().build();
+    Request request =
+        new Request.Builder().url("http://localhost:" + httpPort + urlPath).get().build();
+    try {
+      client.newCall(request).execute();
+    } catch (Exception ex) {
+      ex.printStackTrace();
+    }
+  }
+
+  protected static void waitForSpecificLogLine(
+      Path logFilePath, String line, Duration sleep, Duration timeout) throws IOException {
+    boolean[] result = new boolean[] {false};
+    long total = sleep.toNanos() == 0 ? 0 : timeout.toNanos() / sleep.toNanos();
+    int i = 0;
+    while (i < total && !result[0]) {
+      Files.lines(logFilePath)
+          .forEach(
+              it -> {
+                if (it.contains(line)) {
+                  result[0] = true;
+                }
+              });
+      LockSupport.parkNanos(sleep.toNanos());
+      i++;
+    }
+  }
+}

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

@@ -0,0 +1,109 @@
+package datadog.smoketest;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+
+import datadog.trace.api.DDTags;
+import datadog.trace.test.agent.decoder.DecodedSpan;
+import datadog.trace.test.agent.decoder.DecodedTrace;
+import datadog.trace.test.util.NonRetryable;
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.List;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.DisabledIf;
+
+@NonRetryable
+public class SpringCodeOriginIntegrationTest extends SpringBasedIntegrationTest {
+
+  private boolean traceReceived;
+
+  @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");
+    commandParams.add("-Ddd.code.origin.for.spans.interface.support=true");
+    return ProcessBuilderHelper.createProcessBuilder(
+        commandParams, logFilePath, getAppClass(), params);
+  }
+
+  @Test
+  @DisplayName("testRegularController")
+  @DisabledIf(
+      value = "datadog.environment.JavaVirtualMachine#isJ9",
+      disabledReason = "Flaky on J9 JVMs")
+  void testRegularController() throws Exception {
+    registerTraceListener(this::receiveGreetingTrace);
+    String httpPort = startSpringApp(Collections.emptyList());
+    sendRequest(httpPort, "/greeting"); // trigger CodeOriginProbe instrumentation
+    waitForSpecificLogLine(
+        logFilePath,
+        "DEBUG com.datadog.debugger.agent.ConfigurationUpdater - Re-transformation done",
+        Duration.ofMillis(100),
+        Duration.ofSeconds(30)); // wait for instrumentation to be done
+    sendRequest(httpPort, "/greeting"); // generate first span with tags
+    processRequests(
+        () -> traceReceived, () -> String.format("Timeout! traceReceived=%s", traceReceived));
+  }
+
+  @Test
+  @DisplayName("testInterfacedController")
+  @DisabledIf(
+      value = "datadog.environment.JavaVirtualMachine#isJ9",
+      disabledReason = "Flaky on J9 JVMs")
+  void testInterfacedController() throws Exception {
+    registerTraceListener(this::receiveProcessTrace);
+    String httpPort = startSpringApp(Collections.emptyList());
+    sendRequest(httpPort, "/process"); // trigger CodeOriginProbe instrumentation
+    waitForSpecificLogLine(
+        logFilePath,
+        "DEBUG com.datadog.debugger.agent.ConfigurationUpdater - Re-transformation done",
+        Duration.ofMillis(100),
+        Duration.ofSeconds(30)); // wait for instrumentation to be done
+    sendRequest(httpPort, "/process"); // generate first span with tags
+    processRequests(
+        () -> traceReceived, () -> String.format("Timeout! traceReceived=%s", traceReceived));
+  }
+
+  private void receiveGreetingTrace(DecodedTrace decodedTrace) {
+    for (DecodedSpan span : decodedTrace.getSpans()) {
+      if (span.getName().equals("spring.handler")
+          && span.getResource().equals("WebController.greeting")
+          && span.getMeta().containsKey(DDTags.DD_CODE_ORIGIN_FRAME_TYPE)) {
+        assertEquals("entry", span.getMeta().get(DDTags.DD_CODE_ORIGIN_TYPE));
+        assertEquals(
+            "datadog.smoketest.debugger.controller.WebController",
+            span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_TYPE));
+        assertEquals("WebController.java", span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_FILE));
+        assertEquals("10", span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_LINE));
+        assertEquals("greeting", span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_METHOD));
+        assertEquals("()", span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_SIGNATURE));
+        assertFalse(traceReceived);
+        traceReceived = true;
+      }
+    }
+  }
+
+  private void receiveProcessTrace(DecodedTrace decodedTrace) {
+    for (DecodedSpan span : decodedTrace.getSpans()) {
+      if (span.getName().equals("spring.handler")
+          && span.getResource().equals("InterfacedController.process")
+          && span.getMeta().containsKey(DDTags.DD_CODE_ORIGIN_FRAME_TYPE)) {
+        assertEquals("entry", span.getMeta().get(DDTags.DD_CODE_ORIGIN_TYPE));
+        assertEquals(
+            "datadog.smoketest.debugger.controller.InterfacedController",
+            span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_TYPE));
+        assertEquals(
+            "InterfacedController.java", span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_FILE));
+        assertEquals("11", span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_LINE));
+        assertEquals("process", span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_METHOD));
+        assertEquals("()", span.getMeta().get(DDTags.DD_CODE_ORIGIN_FRAME_SIGNATURE));
+        assertFalse(traceReceived);
+        traceReceived = true;
+      }
+    }
+  }
+}