JAIPilot campaign evaluation: DataDog/dd-trace-java #12254

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/exception/AbstractExceptionDebugger.java

@@ -59,7 +59,6 @@ public void handleException(Throwable t, AgentSpan span) {
       }
       return;
     }
-
     String fingerprint = Fingerprinter.fingerprint(t, classNameFiltering);
     if (fingerprint == null) {
       LOGGER.debug("Unable to fingerprint exception", t);
@@ -80,7 +79,13 @@ public void handleException(Throwable t, AgentSpan span) {
         return;
       }
       processSnapshotsAndSetTags(
-          t, span, state, chainedExceptionsList, fingerprint, maxCapturedFrames);
+          t,
+          innerMostException,
+          span,
+          state,
+          chainedExceptionsList,
+          fingerprint,
+          maxCapturedFrames);
       exceptionProbeManager.updateLastCapture(fingerprint);
     } else {
       // climb up the exception chain to find the first exception that has instrumented frames
@@ -126,6 +131,7 @@ protected void addStackFrameTags(
 
   private void processSnapshotsAndSetTags(
       Throwable t,
+      Throwable innerMostException,
       AgentSpan span,
       ExceptionProbeManager.ThrowableState state,
       List<Throwable> chainedExceptions,
@@ -190,6 +196,9 @@ private void processSnapshotsAndSetTags(
           state.getExceptionId());
       span.setTag(Tags.ERROR_DEBUG_INFO_CAPTURED, true);
       span.setTag(DD_DEBUG_ERROR_EXCEPTION_HASH, fingerprint);
+      // Remove ThrowableState to avoid growing indefinitely for singleton exception instances
+      // like ones generated for FastThrow optimization (OmitStackTraceInFastThrow)
+      exceptionProbeManager.removeThrowableState(innerMostException);
     }
   }
 

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/exception/ExceptionProbeManager.java

@@ -67,6 +67,10 @@ public ClassNameFilter getClassNameFilter() {
     return classNameFiltering;
   }
 
+  public void removeThrowableState(Throwable t) {
+    snapshotsByThrowable.remove(t);
+  }
+
   static class CreationResult {
     final int probesCreated;
     final int thirdPartyFrames;
@@ -198,6 +202,10 @@ public List<Snapshot> getSnapshots() {
     }
 
     public void addSnapshot(Snapshot snapshot) {
+      if (snapshots.size() > 256) {
+        LOGGER.debug("Too many (256) snapshots for exceptionId={}, dropping snapshot", exceptionId);
+        return;
+      }
       snapshots.add(snapshot);
     }
 

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/exception/DefaultExceptionDebuggerTest.java

@@ -6,6 +6,7 @@
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toMap;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -110,13 +111,14 @@ public void nestedException() {
         Duration.ofSeconds(30));
     generateSnapshots(exception);
     exception.printStackTrace();
-    exceptionDebugger.handleException(exception, span);
+    Throwable innerMostException = ExceptionHelper.getInnerMostThrowable(exception);
+    // capture the state before the second call: a successful assignment removes it from the
+    // manager, so it can no longer be looked up by throwable afterwards
     ExceptionProbeManager.ThrowableState state =
-        exceptionDebugger
-            .getExceptionProbeManager()
-            .getStateByThrowable(ExceptionHelper.getInnerMostThrowable(exception));
-    assertEquals(
-        state.getExceptionId(), spanTags.get(DefaultExceptionDebugger.DD_DEBUG_ERROR_EXCEPTION_ID));
+        exceptionDebugger.getExceptionProbeManager().getStateByThrowable(innerMostException);
+    String exceptionId = state.getExceptionId();
+    exceptionDebugger.handleException(exception, span);
+    assertEquals(exceptionId, spanTags.get(DefaultExceptionDebugger.DD_DEBUG_ERROR_EXCEPTION_ID));
     Map<String, Snapshot> snapshotMap =
         listener.snapshots.stream().collect(toMap(Snapshot::getId, Function.identity()));
     List<String> lines = parseStackTrace(exception);
@@ -149,8 +151,13 @@ public void nestedException() {
         expectedFrameIndex,
         "com.datadog.debugger.exception.DefaultExceptionDebuggerTest",
         "createTest1Exception");
+    // ThrowableState is keyed by the innermost throwable, not the top-level one passed to
+    // handleException: removal must use that same key or it silently leaks for chained exceptions
+    assertNull(
+        exceptionDebugger.getExceptionProbeManager().getStateByThrowable(innerMostException));
     // make sure we are not leaking references
     exception = null; // release strong reference
+    innerMostException = null;
     System.gc();
     // calling ExceptionProbeManager#hasExceptionStateTracked() will call WeakIdentityHashMap#size()
     // through isEmpty() an will purge stale entries
@@ -184,13 +191,14 @@ public void doubleNestedException() {
     generateSnapshots(simpleException);
     exceptionDebugger.handleException(simpleException, span);
     nestedException.printStackTrace();
-    exceptionDebugger.handleException(nestedException, span);
+    Throwable innerMostException = ExceptionHelper.getInnerMostThrowable(nestedException);
+    // capture the state before the second call: a successful assignment removes it from the
+    // manager, so it can no longer be looked up by throwable afterwards
     ExceptionProbeManager.ThrowableState state =
-        exceptionDebugger
-            .getExceptionProbeManager()
-            .getStateByThrowable(ExceptionHelper.getInnerMostThrowable(nestedException));
-    assertEquals(
-        state.getExceptionId(), spanTags.get(DefaultExceptionDebugger.DD_DEBUG_ERROR_EXCEPTION_ID));
+        exceptionDebugger.getExceptionProbeManager().getStateByThrowable(innerMostException);
+    String exceptionId = state.getExceptionId();
+    exceptionDebugger.handleException(nestedException, span);
+    assertEquals(exceptionId, spanTags.get(DefaultExceptionDebugger.DD_DEBUG_ERROR_EXCEPTION_ID));
     Map<String, Snapshot> snapshotMap =
         listener.snapshots.stream().collect(toMap(Snapshot::getId, Function.identity()));
     List<String> lines = parseStackTrace(nestedException);
@@ -224,6 +232,10 @@ public void doubleNestedException() {
         expectedFrameIndex,
         "com.datadog.debugger.exception.DefaultExceptionDebuggerTest",
         "createTest1Exception");
+    // ThrowableState is keyed by the innermost throwable, not the top-level one passed to
+    // handleException: removal must use that same key or it silently leaks for chained exceptions
+    assertNull(
+        exceptionDebugger.getExceptionProbeManager().getStateByThrowable(innerMostException));
   }
 
   @Test
@@ -249,13 +261,14 @@ public void nestedExceptionFullThirdParty() {
         Duration.ofSeconds(30));
     generateSnapshots(exception);
     exception.printStackTrace();
-    exceptionDebugger.handleException(exception, span);
+    Throwable innerMostException = ExceptionHelper.getInnerMostThrowable(exception);
+    // capture the state before the second call: a successful assignment removes it from the
+    // manager, so it can no longer be looked up by throwable afterwards
     ExceptionProbeManager.ThrowableState state =
-        exceptionDebugger
-            .getExceptionProbeManager()
-            .getStateByThrowable(ExceptionHelper.getInnerMostThrowable(exception));
-    assertEquals(
-        state.getExceptionId(), spanTags.get(DefaultExceptionDebugger.DD_DEBUG_ERROR_EXCEPTION_ID));
+        exceptionDebugger.getExceptionProbeManager().getStateByThrowable(innerMostException);
+    String exceptionId = state.getExceptionId();
+    exceptionDebugger.handleException(exception, span);
+    assertEquals(exceptionId, spanTags.get(DefaultExceptionDebugger.DD_DEBUG_ERROR_EXCEPTION_ID));
     Map<String, Snapshot> snapshotMap =
         listener.snapshots.stream().collect(toMap(Snapshot::getId, Function.identity()));
     List<String> lines = parseStackTrace(exception);
@@ -269,6 +282,8 @@ public void nestedExceptionFullThirdParty() {
         expectedFrameIndex,
         "com.datadog.debugger.exception.DefaultExceptionDebuggerTest",
         "createTest2Exception");
+    assertNull(
+        exceptionDebugger.getExceptionProbeManager().getStateByThrowable(innerMostException));
   }
 
   @Test
@@ -308,12 +323,36 @@ public StackTraceElement[] getStackTrace() {
     generateSnapshots(exception);
     ExceptionProbeManager.ThrowableState state =
         exceptionDebugger.getExceptionProbeManager().getStateByThrowable(exception);
-    List<Snapshot> snapshots = state.getSnapshots();
+    String firstSnapshotId = state.getSnapshots().get(0).getId();
     // This should hit the `currentIdx < 0` branch and fallback to i=0
     exceptionDebugger.handleException(exception, span);
     String tagName = String.format(SNAPSHOT_ID_TAG_FMT, 0);
     assertTrue(spanTags.containsKey(tagName));
-    assertEquals(snapshots.get(0).getId(), spanTags.get(tagName));
+    assertEquals(firstSnapshotId, spanTags.get(tagName));
+    // ThrowableState is removed once assigned & sent, so the state does not grow indefinitely
+    assertNull(exceptionDebugger.getExceptionProbeManager().getStateByThrowable(exception));
+  }
+
+  @Test
+  public void throwableStateRemovedAfterSendingToAvoidUnboundedGrowth() {
+    // simulates a singleton exception instance re-thrown repeatedly, as generated by the JIT's
+    // OmitStackTraceInFastThrow optimization
+    RuntimeException exception = new RuntimeException("test");
+    String fingerprint = Fingerprinter.fingerprint(exception, classNameFiltering);
+    AgentSpan span = mock(AgentSpan.class);
+    doAnswer(this::recordTags).when(span).setTag(anyString(), anyString());
+    when(span.getTag(anyString())).thenAnswer(inv -> spanTags.get(inv.getArgument(0)));
+    when(span.getTags()).thenReturn(spanTags);
+    exceptionDebugger.handleException(exception, span);
+    assertWithTimeout(
+        () -> exceptionDebugger.getExceptionProbeManager().isAlreadyInstrumented(fingerprint),
+        Duration.ofSeconds(30));
+    generateSnapshots(exception);
+    ExceptionProbeManager.ThrowableState state =
+        exceptionDebugger.getExceptionProbeManager().getStateByThrowable(exception);
+    assertTrue(state.getSnapshots().size() > 0);
+    exceptionDebugger.handleException(exception, span);
+    assertNull(exceptionDebugger.getExceptionProbeManager().getStateByThrowable(exception));
   }
 
   private Object recordTags(InvocationOnMock invocationOnMock) {

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/exception/ExceptionProbeManagerTest.java

@@ -7,8 +7,10 @@
 import static org.mockito.Mockito.when;
 
 import com.datadog.debugger.probe.ExceptionProbe;
+import com.datadog.debugger.sink.Snapshot;
 import com.datadog.debugger.util.ClassNameFiltering;
 import datadog.trace.api.Config;
+import datadog.trace.bootstrap.debugger.CapturedContext;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
@@ -46,7 +48,7 @@ void instrumentSingleFrame() {
     ExceptionProbeManager exceptionProbeManager = new ExceptionProbeManager(classNameFiltering);
 
     String fingerprint = Fingerprinter.fingerprint(exception, classNameFiltering);
-    assertEquals("4974b2b4853e6152d8f218fb79a42a761a45335447e22e53d75f5325e742655", fingerprint);
+    assertEquals("b3cf344f48cb7d44f2f75ac267667fbdda167e494f24d55149843896c16584ca", fingerprint);
     exceptionProbeManager.createProbesForException(exception.getStackTrace(), 0);
     assertEquals(1, exceptionProbeManager.getProbes().size());
     ExceptionProbe exceptionProbe = exceptionProbeManager.getProbes().iterator().next();
@@ -121,4 +123,29 @@ RuntimeException level2() {
   RuntimeException level3() {
     return new RuntimeException("3 level deep exception");
   }
+
+  @Test
+  void addSnapshotCapsGrowth() {
+    ClassNameFiltering classNameFiltering = ClassNameFiltering.allowAll();
+    ExceptionProbeManager exceptionProbeManager = new ExceptionProbeManager(classNameFiltering);
+    RuntimeException throwable = new RuntimeException("test");
+    for (int i = 0; i < 300; i++) {
+      exceptionProbeManager.addSnapshot(createSnapshot(throwable));
+    }
+    ExceptionProbeManager.ThrowableState state =
+        exceptionProbeManager.getStateByThrowable(throwable);
+    int cappedSize = state.getSnapshots().size();
+    for (int i = 0; i < 50; i++) {
+      exceptionProbeManager.addSnapshot(createSnapshot(throwable));
+    }
+    assertEquals(cappedSize, state.getSnapshots().size());
+  }
+
+  private static Snapshot createSnapshot(Throwable throwable) {
+    Snapshot snapshot = new Snapshot(Thread.currentThread(), null, 1);
+    CapturedContext context = new CapturedContext();
+    context.addThrowable(throwable);
+    snapshot.setExit(context);
+    return snapshot;
+  }
 }