Make source file tracking asynchronous

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerAgent.java

@@ -74,10 +74,7 @@ public static synchronized void run(Instrumentation inst, SharedCommunicationObj
     Config config = Config.get();
     DebuggerContext.initProductConfigUpdater(new DefaultProductConfigUpdater());
     classesToRetransformFinder = new ClassesToRetransformFinder();
-    if (config.isDynamicInstrumentationEnabled() || config.isDebuggerExceptionEnabled()) {
-      // only activate Source File Tracking if DI or ER is enabled from the start
-      setupSourceFileTracking(instrumentation, classesToRetransformFinder);
-    }
+    setupSourceFileTracking(instrumentation, classesToRetransformFinder);
     if (config.isDebuggerCodeOriginEnabled()) {
       startCodeOriginForSpans();
     }
@@ -314,7 +311,10 @@ private static String getDiagnosticEndpoint(
 
   private static void setupSourceFileTracking(
       Instrumentation instrumentation, ClassesToRetransformFinder finder) {
-    instrumentation.addTransformer(new SourceFileTrackingTransformer(finder));
+    SourceFileTrackingTransformer sourceFileTrackingTransformer =
+        new SourceFileTrackingTransformer(finder);
+    sourceFileTrackingTransformer.start();
+    instrumentation.addTransformer(sourceFileTrackingTransformer);
   }
 
   private static void loadFromFile(

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/SourceFileTrackingTransformer.java

@@ -4,22 +4,57 @@
 import static com.datadog.debugger.util.ClassFileHelper.stripPackagePath;
 
 import com.datadog.debugger.util.ClassFileHelper;
+import datadog.trace.util.AgentTaskScheduler;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.IllegalClassFormatException;
 import java.security.ProtectionDomain;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.TimeUnit;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Permanent Transformer to track all Inner or Top-Level classes associated with the same SourceFile
  * (String.java) Allows to get all classes that are dependent from a source file and be able to
  * trigger {@link java.lang.instrument.Instrumentation#retransformClasses(Class[])} on them
  */
 public class SourceFileTrackingTransformer implements ClassFileTransformer {
+  private static final Logger LOGGER = LoggerFactory.getLogger(SourceFileTrackingTransformer.class);
+
   private final ClassesToRetransformFinder finder;
+  private final Queue<SourceFileItem> queue = new ConcurrentLinkedQueue<>();
+  private final AgentTaskScheduler scheduler = AgentTaskScheduler.INSTANCE;
+  private AgentTaskScheduler.Scheduled<Runnable> scheduled;
 
   public SourceFileTrackingTransformer(ClassesToRetransformFinder finder) {
     this.finder = finder;
   }
 
+  public void start() {
+    scheduled = scheduler.scheduleAtFixedRate(this::flush, 0, 1, TimeUnit.SECONDS);
+  }
+
+  public void stop() {
+    if (scheduled != null) {
+      scheduled.cancel();
+    }
+  }
+
+  void flush() {
+    if (queue.isEmpty()) {
+      return;
+    }
+    int size = queue.size();
+    long start = System.nanoTime();
+    SourceFileItem item;
+    while ((item = queue.poll()) != null) {
+      registerSourceFile(item.className, item.classfileBuffer);
+    }
+    LOGGER.debug(
+        "flushing {} source file items in {}ms", size, (System.nanoTime() - start) / 1_000_000);
+  }
+
   @Override
   public byte[] transform(
       ClassLoader loader,
@@ -31,16 +66,30 @@ public byte[] transform(
     if (className == null) {
       return null;
     }
+    queue.add(new SourceFileItem(className, classfileBuffer));
+    return null;
+  }
+
+  private void registerSourceFile(String className, byte[] classfileBuffer) {
     String sourceFile = ClassFileHelper.extractSourceFile(classfileBuffer);
     if (sourceFile == null) {
-      return null;
+      return;
     }
     String simpleClassName = stripPackagePath(className);
     String simpleSourceFile = removeExtension(sourceFile);
     if (simpleClassName.equals(simpleSourceFile)) {
-      return null;
+      return;
     }
     finder.register(sourceFile, className);
-    return null;
+  }
+
+  private static class SourceFileItem {
+    final String className;
+    final byte[] classfileBuffer;
+
+    public SourceFileItem(String className, byte[] classfileBuffer) {
+      this.className = className;
+      this.classfileBuffer = classfileBuffer;
+    }
   }
 }

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/SourceFileTrackingTransformerTest.java

@@ -32,6 +32,7 @@ void transformTopLevel() throws IllegalClassFormatException {
         null,
         null,
         getClassFileBytes(MyTopLevelClass.class));
+    sourceFileTrackingTransformer.flush();
     changedClasses =
         finder.getAllLoadedChangedClasses(
             new Class[] {TopLevelHelper.class, MyTopLevelClass.class}, comparer);
@@ -48,6 +49,7 @@ void transformInner() throws IllegalClassFormatException {
     ConfigurationComparer comparer = createComparer("InnerHelper.java");
     sourceFileTrackingTransformer.transform(
         null, getInternalName(InnerHelper.class), null, null, getClassFileBytes(InnerHelper.class));
+    sourceFileTrackingTransformer.flush();
     List<Class<?>> changedClasses =
         finder.getAllLoadedChangedClasses(new Class[] {InnerHelper.class}, comparer);
     assertEquals(1, changedClasses.size());
@@ -64,6 +66,7 @@ void transformInner() throws IllegalClassFormatException {
         null,
         null,
         getClassFileBytes(InnerHelper.MySecondInner.class));
+    sourceFileTrackingTransformer.flush();
     changedClasses =
         finder.getAllLoadedChangedClasses(
             new Class[] {

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

@@ -43,6 +43,7 @@ public static void registerMethods() {
     methodsByName.put("tracedMethod", ServerDebuggerTestApplication::runTracedMethod);
     methodsByName.put("loopingFullMethod", ServerDebuggerTestApplication::runLoopingFullMethod);
     methodsByName.put("loopingTracedMethod", ServerDebuggerTestApplication::runLoopingTracedMethod);
+    methodsByName.put("topLevelMethod", ServerDebuggerTestApplication::runTopLevelMethod);
   }
 
   public ServerDebuggerTestApplication(String controlServerUrl) {
@@ -163,6 +164,10 @@ private static void runLoopingTracedMethod(String arg) {
     }
   }
 
+  private static String runTopLevelMethod(String arg) {
+    return TopLevel.process(arg);
+  }
+
   private static String fullMethod(
       int argInt, String argStr, double argDouble, Map<String, String> argMap, String... argVar) {
     try {
@@ -290,3 +295,9 @@ public MockResponse dispatch(RecordedRequest request) throws InterruptedExceptio
     }
   }
 }
+
+class TopLevel {
+  public static String process(String arg) {
+    return "TopLevel.process: " + arg;
+  }
+}

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

@@ -45,7 +45,7 @@ void testCodeOriginTraceAnnotation() throws Exception {
                 assertEquals("runTracedMethod", span.getMeta().get(DD_CODE_ORIGIN_FRAME_METHOD));
                 assertEquals(
                     "(java.lang.String)", span.getMeta().get(DD_CODE_ORIGIN_FRAME_SIGNATURE));
-                assertEquals("145", span.getMeta().get(DD_CODE_ORIGIN_FRAME_LINE));
+                assertEquals("146", span.getMeta().get(DD_CODE_ORIGIN_FRAME_LINE));
                 codeOrigin.set(true);
               }
             }

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

@@ -39,6 +39,28 @@ void testDynamicInstrumentationEnablement() throws Exception {
     waitForReTransformation(appUrl); // wait for retransformation of removed probe
   }
 
+  @Test
+  @DisplayName("testDynamicInstrumentationEnablementWithLineProbe")
+  void testDynamicInstrumentationEnablementWithLineProbe() throws Exception {
+    appUrl = startAppAndAndGetUrl();
+    setConfigOverrides(createConfigOverrides(true, false));
+    LogProbe probe =
+        LogProbe.builder()
+            .probeId(LINE_PROBE_ID1)
+            .where("ServerDebuggerTestApplication.java", 301)
+            .build();
+    setCurrentConfiguration(createConfig(probe));
+    waitForFeatureStarted(appUrl, "Dynamic Instrumentation");
+    execute(appUrl, "topLevelMethod", "");
+    waitForInstrumentation(appUrl, "datadog.smoketest.debugger.TopLevel");
+    // disable DI
+    setConfigOverrides(createConfigOverrides(false, false));
+    waitForFeatureStopped(appUrl, "Dynamic Instrumentation");
+    waitForReTransformation(
+        appUrl,
+        "datadog.smoketest.debugger.TopLevel"); // wait for retransformation of removed probe
+  }
+
   @Test
   @DisplayName("testDynamicInstrumentationEnablementStaticallyDisabled")
   void testDynamicInstrumentationEnablementStaticallyDisabled() throws Exception {

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

@@ -250,6 +250,33 @@ void testMultiProbes() throws Exception {
     assertFalse(logHasErrors(logFilePath, it -> false));
   }
 
+  @Test
+  @DisplayName("testLineProbe")
+  void testLineProbe() throws Exception {
+    final String METHOD_NAME = "fullMethod";
+    final String EXPECTED_UPLOADS = "4"; // 3 statuses + 1 snapshot
+    LogProbe probe =
+        LogProbe.builder()
+            .probeId(LINE_PROBE_ID1)
+            .where("DebuggerTestApplication.java", 88)
+            .captureSnapshot(true)
+            .build();
+    setCurrentConfiguration(createConfig(probe));
+    targetProcess = createProcessBuilder(logFilePath, METHOD_NAME, EXPECTED_UPLOADS).start();
+    AtomicBoolean snapshotReceived = new AtomicBoolean();
+    registerSnapshotListener(
+        snapshot -> {
+          assertEquals(LINE_PROBE_ID1.getId(), snapshot.getProbe().getId());
+          CapturedContext capturedContext = snapshot.getCaptures().getLines().get(88);
+          assertFullMethodCaptureArgs(capturedContext);
+          assertNull(capturedContext.getLocals());
+          assertNull(capturedContext.getCapturedThrowable());
+          snapshotReceived.set(true);
+        });
+    AtomicBoolean statusResult = registerCheckReceivedInstalledEmitting();
+    processRequests(() -> snapshotReceived.get() && statusResult.get());
+  }
+
   @Test
   @DisplayName("testSamplingSnapshotDefault")
   @DisabledIf(value = "datadog.trace.api.Platform#isJ9", disabledReason = "Flaky on J9 JVMs")

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

@@ -31,7 +31,8 @@ public class ServerAppDebuggerIntegrationTest extends BaseIntegrationTest {
       "datadog.smoketest.debugger.ServerDebuggerTestApplication";
   protected static final String CONTROL_URL = "/control";
   protected static final ProbeId PROBE_ID = new ProbeId("123356536", 0);
-  protected static final ProbeId PROBE_ID2 = new ProbeId("1233565367", 12);
+  protected static final ProbeId LINE_PROBE_ID1 =
+      new ProbeId("beae1817-f3b0-4ea8-a74f-000000000001", 0);
   protected static final String TEST_APP_CLASS_NAME = "ServerDebuggerTestApplication";
   protected static final String FULL_METHOD_NAME = "fullMethod";
   protected static final String TRACED_METHOD_NAME = "tracedMethod";
@@ -106,9 +107,11 @@ protected void execute(String appUrl, String methodName, String arg) throws IOEx
   }
 
   protected void waitForInstrumentation(String appUrl) throws Exception {
-    String url =
-        String.format(
-            appUrl + "/waitForInstrumentation?classname=%s", SERVER_DEBUGGER_TEST_APP_CLASS);
+    waitForInstrumentation(appUrl, SERVER_DEBUGGER_TEST_APP_CLASS);
+  }
+
+  protected void waitForInstrumentation(String appUrl, String className) throws Exception {
+    String url = String.format(appUrl + "/waitForInstrumentation?classname=%s", className);
     LOG.info("waitForInstrumentation with url={}", url);
     sendRequest(url);
     AtomicBoolean received = new AtomicBoolean();
@@ -136,9 +139,11 @@ protected void waitForAProbeStatus(ProbeStatus.Status status) throws Exception {
   }
 
   protected void waitForReTransformation(String appUrl) throws IOException {
-    String url =
-        String.format(
-            appUrl + "/waitForReTransformation?classname=%s", SERVER_DEBUGGER_TEST_APP_CLASS);
+    waitForReTransformation(appUrl, SERVER_DEBUGGER_TEST_APP_CLASS);
+  }
+
+  protected void waitForReTransformation(String appUrl, String className) throws IOException {
+    String url = String.format(appUrl + "/waitForReTransformation?classname=%s", className);
     sendRequest(url);
     LOG.info("re-transformation done");
   }

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

@@ -11,7 +11,8 @@ public class SimpleAppDebuggerIntegrationTest extends BaseIntegrationTest {
   protected static final String DEBUGGER_TEST_APP_CLASS =
       "datadog.smoketest.debugger.DebuggerTestApplication";
   protected static final ProbeId PROBE_ID = new ProbeId("123356536", 0);
-  protected static final ProbeId PROBE_ID2 = new ProbeId("1233565368", 12);
+  protected static final ProbeId LINE_PROBE_ID1 =
+      new ProbeId("beae1817-f3b0-4ea8-a74f-000000000001", 0);
   protected static final String MAIN_CLASS_NAME = "Main";
 
   @Override