Add capping on SourceFile tracking queue (#9245)

tracks number of item put into the queue and dequeued.
limits to 4096 items.
Avoid using queue.size that is O(n) and keep the size in AtomicInteger
dump the number of entries in the SourceFile map into tracer flare

Author: jpbempel

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

@@ -478,7 +478,9 @@ private static void addReportToFlare(ZipOutputStream zip) throws IOException {
             "SymbolDB stats:",
             symbolDBStats,
             "Exception Fingerprints:",
-            exceptionFingerprints);
+            exceptionFingerprints,
+            "SourceFile tracking entries:",
+            String.valueOf(classesToRetransformFinder.getClassNamesBySourceFile().size()));
     TracerFlare.addText(zip, "dynamic_instrumentation.txt", content);
   }
 }

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

@@ -6,6 +6,7 @@
 import com.datadog.debugger.util.ClassFileHelper;
 import com.datadog.debugger.util.ClassNameFiltering;
 import datadog.trace.api.Config;
+import datadog.trace.relocate.api.RatelimitedLogger;
 import datadog.trace.util.AgentTaskScheduler;
 import datadog.trace.util.Strings;
 import java.lang.instrument.ClassFileTransformer;
@@ -14,6 +15,7 @@
 import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -24,10 +26,15 @@
  */
 public class SourceFileTrackingTransformer implements ClassFileTransformer {
   private static final Logger LOGGER = LoggerFactory.getLogger(SourceFileTrackingTransformer.class);
+  private static final int MINUTES_BETWEEN_ERROR_LOG = 5;
+  static final int MAX_QUEUE_SIZE = 4096;
 
+  private final RatelimitedLogger ratelimitedLogger =
+      new RatelimitedLogger(LOGGER, MINUTES_BETWEEN_ERROR_LOG, TimeUnit.MINUTES);
   private final ClassesToRetransformFinder finder;
   private final Queue<SourceFileItem> queue = new ConcurrentLinkedQueue<>();
   private final AgentTaskScheduler scheduler = AgentTaskScheduler.INSTANCE;
+  private final AtomicInteger queueSize = new AtomicInteger(0);
   private AgentTaskScheduler.Scheduled<Runnable> scheduled;
   // this field MUST only be used in flush() calling thread
   private ClassNameFiltering classNameFilter;
@@ -55,14 +62,23 @@ void flush() {
     if (queue.isEmpty()) {
       return;
     }
-    int size = queue.size();
+    int itemCount = 0;
     long start = System.nanoTime();
     SourceFileItem item;
     while ((item = queue.poll()) != null) {
+      queueSize.decrementAndGet();
       registerSourceFile(item.className, item.classfileBuffer);
+      itemCount++;
     }
     LOGGER.debug(
-        "flushing {} source file items in {}ms", size, (System.nanoTime() - start) / 1_000_000);
+        "flushing {} source file items in {}ms, totalentries: {}",
+        itemCount,
+        (System.nanoTime() - start) / 1_000_000,
+        finder.getClassNamesBySourceFile().size());
+  }
+
+  int getQueueSize() {
+    return queueSize.get();
   }
 
   @Override
@@ -76,7 +92,12 @@ public byte[] transform(
     if (className == null) {
       return null;
     }
+    if (queueSize.get() >= MAX_QUEUE_SIZE) {
+      ratelimitedLogger.warn("SourceFile Tracking queue full, dropping class: {}", className);
+      return null;
+    }
     queue.add(new SourceFileItem(className, classfileBuffer));
+    queueSize.incrementAndGet();
     return null;
   }
 

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

@@ -1,5 +1,6 @@
 package com.datadog.debugger.agent;
 
+import static com.datadog.debugger.agent.SourceFileTrackingTransformer.MAX_QUEUE_SIZE;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static utils.TestClassFileHelper.getClassFileBytes;
 
@@ -138,6 +139,22 @@ void thirdPartyExcludes() throws IllegalClassFormatException {
     assertEquals(0, finder.getClassNamesBySourceFile().size());
   }
 
+  @Test
+  void maxQueue() throws IllegalClassFormatException {
+    ClassesToRetransformFinder finder = new ClassesToRetransformFinder();
+    SourceFileTrackingTransformer sourceFileTrackingTransformer =
+        new SourceFileTrackingTransformer(finder);
+    for (int i = 0; i < MAX_QUEUE_SIZE + 10; i++) {
+      sourceFileTrackingTransformer.transform(
+          null,
+          getInternalName(TopLevelHelper.class),
+          null,
+          null,
+          getClassFileBytes(TopLevelHelper.class));
+    }
+    assertEquals(MAX_QUEUE_SIZE, sourceFileTrackingTransformer.getQueueSize());
+  }
+
   private static void replaceInByteArray(byte[] buffer, byte[] oldBytes, byte[] newBytes) {
     int oldIdx = 0;
     for (int i = 0; i < buffer.length; i++) {