Fix memory leak in Exception Replay (#7885)

The weak map stateByThrowable keep the throwable as the key but this
exception is also strong referenced by snapshots stored inside the
ThrowableState in CapturedThrowable and in locals and extensions for
@exception.
Fixing by storing weak reference inside the CapturedThrowable and
clearing the other ref for @exception at commit time

Author: jpbempel

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

@@ -9,6 +9,7 @@
 import datadog.trace.bootstrap.debugger.util.Redaction;
 import datadog.trace.bootstrap.debugger.util.TimeoutChecker;
 import datadog.trace.bootstrap.debugger.util.WellKnownClasses;
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -183,6 +184,10 @@ public ValueReferenceResolver withExtensions(Map<String, Object> extensions) {
     return new CapturedContext(this, extensions);
   }
 
+  public void removeExtension(String name) {
+    extensions.remove(name);
+  }
+
   private void addExtension(String name, Object value) {
     extensions.put(name, value);
   }
@@ -642,7 +647,7 @@ public TimeoutChecker getTimeoutChecker() {
   public static class CapturedThrowable {
     private final String type;
     private final String message;
-    private final transient Throwable throwable;
+    private final transient WeakReference<Throwable> throwable;
 
     /*
      * Need to exclude stacktrace from equals/hashCode computation.
@@ -663,7 +668,7 @@ public CapturedThrowable(
       this.type = type;
       this.message = message;
       this.stacktrace = new ArrayList<>(stacktrace);
-      this.throwable = t;
+      this.throwable = new WeakReference<>(t);
     }
 
     public String getType() {
@@ -679,7 +684,7 @@ public List<CapturedStackFrame> getStacktrace() {
     }
 
     public Throwable getThrowable() {
-      return throwable;
+      return throwable.get();
     }
 
     private static List<CapturedStackFrame> captureFrames(StackTraceElement[] stackTrace) {

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

@@ -146,11 +146,14 @@ boolean shouldCaptureException(String fingerprint, Clock clock) {
 
   public void addSnapshot(Snapshot snapshot) {
     Throwable throwable = snapshot.getCaptures().getReturn().getCapturedThrowable().getThrowable();
+    if (throwable == null) {
+      LOGGER.debug("Snapshot has no throwable: {}", snapshot.getId());
+      return;
+    }
     throwable = ExceptionHelper.getInnerMostThrowable(throwable);
     if (throwable == null) {
-      LOGGER.debug(
-          "Unable to find root cause of exception: {}",
-          snapshot.getCaptures().getReturn().getCapturedThrowable().getThrowable().toString());
+      throwable = snapshot.getCaptures().getReturn().getCapturedThrowable().getThrowable();
+      LOGGER.debug("Unable to find root cause of exception: {}", String.valueOf(throwable));
       return;
     }
     ThrowableState state =
@@ -172,6 +175,10 @@ void updateLastCapture(String fingerprint, Clock clock) {
     fingerprints.put(fingerprint, Instant.now(clock));
   }
 
+  boolean hasExceptionStateTracked() {
+    return !snapshotsByThrowable.isEmpty();
+  }
+
   public static class ThrowableState {
     private final String exceptionId;
     private List<Snapshot> snapshots;

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/ExceptionProbe.java

@@ -13,6 +13,7 @@
 import datadog.trace.bootstrap.debugger.ProbeId;
 import datadog.trace.bootstrap.debugger.ProbeImplementation;
 import datadog.trace.bootstrap.debugger.ProbeLocation;
+import datadog.trace.bootstrap.debugger.el.ValueReferences;
 import java.util.List;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -67,6 +68,10 @@ public void evaluate(
       return;
     }
     Throwable throwable = context.getCapturedThrowable().getThrowable();
+    if (throwable == null) {
+      LOGGER.debug("Throwable cleared by GC");
+      return;
+    }
     Throwable innerMostThrowable = getInnerMostThrowable(throwable);
     String fingerprint =
         Fingerprinter.fingerprint(innerMostThrowable, exceptionProbeManager.getClassNameFilter());
@@ -107,6 +112,9 @@ public void commit(
        * - DebuggerContext.commit()
        */
       snapshot.recordStackTrace(4);
+      // clear any strong ref to the exception before adding the snapshot to avoid leaking snapshots
+      // inside the stateByThrowable map
+      clearExceptionRefs(snapshot);
       // add snapshot for later to wait for triggering point (ExceptionDebugger::handleException)
       exceptionProbeManager.addSnapshot(snapshot);
       LOGGER.debug(
@@ -117,6 +125,11 @@ public void commit(
     }
   }
 
+  private void clearExceptionRefs(Snapshot snapshot) {
+    snapshot.getCaptures().getReturn().getLocals().remove(ValueReferences.EXCEPTION_REF);
+    snapshot.getCaptures().getReturn().removeExtension(ValueReferences.EXCEPTION_EXTENSION_NAME);
+  }
+
   @Override
   public void buildLocation(InstrumentationResult result) {
     String type = where.getTypeName();

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

@@ -143,6 +143,14 @@ public void nestedException() {
         expectedFrameIndex,
         "com.datadog.debugger.exception.DefaultExceptionDebuggerTest",
         "createTest1Exception");
+    // make sure we are not leaking references
+    exception = null; // release strong reference
+    System.gc();
+    // calling ExceptionProbeManager#hasExceptionStateTracked() will call WeakIdentityHashMap#size()
+    // through isEmpty() an will purge stale entries
+    assertWithTimeout(
+        () -> !exceptionDebugger.getExceptionProbeManager().hasExceptionStateTracked(),
+        Duration.ofSeconds(30));
   }
 
   @Test

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

@@ -131,7 +131,6 @@ public void instrumentAndCaptureSnapshots() throws Exception {
     Snapshot snapshot0 = listener.snapshots.get(0);
     assertProbeId(probeIdsByMethodName, "processWithException", snapshot0.getProbe().getId());
     assertEquals("oops", snapshot0.getCaptures().getReturn().getCapturedThrowable().getMessage());
-    assertTrue(snapshot0.getCaptures().getReturn().getLocals().containsKey("@exception"));
     ProbeLocation location = snapshot0.getProbe().getLocation();
     assertEquals(
         location.getType() + "." + location.getMethod(), snapshot0.getStack().get(0).getFunction());
@@ -212,7 +211,6 @@ public void recursive() throws Exception {
     assertProbeId(probeIdsByMethodName, "fiboException", snapshot0.getProbe().getId());
     assertEquals(
         "oops fibo", snapshot0.getCaptures().getReturn().getCapturedThrowable().getMessage());
-    assertTrue(snapshot0.getCaptures().getReturn().getLocals().containsKey("@exception"));
     assertEquals("1", getValue(snapshot0.getCaptures().getReturn().getArguments().get("n")));
     Snapshot snapshot1 = listener.snapshots.get(1);
     assertEquals("2", getValue(snapshot1.getCaptures().getReturn().getArguments().get("n")));