Exclude ProxyLeakTask exception from exception profiling (#8666)

Author: jbachorik

dd-java-agent/agent-bootstrap/src/main/java11/datadog/trace/bootstrap/instrumentation/jfr/exceptions/ExceptionProfiling.java

@@ -1,6 +1,7 @@
 package datadog.trace.bootstrap.instrumentation.jfr.exceptions;
 
 import datadog.trace.api.Config;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 
 /**
  * JVM-wide singleton exception profiling service. Uses {@linkplain Config} class to configure
@@ -13,6 +14,24 @@ private static final class Holder {
     static final ExceptionProfiling INSTANCE = new ExceptionProfiling(Config.get());
   }
 
+  /**
+   * Support for excluding certain exception types because they are used for control flow or leak
+   * detection.
+   */
+  public static final class Exclusion {
+    public static void enter() {
+      CallDepthThreadLocalMap.incrementCallDepth(Exclusion.class);
+    }
+
+    public static void exit() {
+      CallDepthThreadLocalMap.decrementCallDepth(Exclusion.class);
+    }
+
+    public static boolean isEffective() {
+      return CallDepthThreadLocalMap.getCallDepth(Exclusion.class) > 0;
+    }
+  }
+
   /**
    * Get a pre-configured shared instance.
    *

dd-java-agent/instrumentation/exception-profiling/build.gradle

@@ -11,6 +11,7 @@ apply from: "$rootDir/gradle/java.gradle"
 apply plugin: "idea"
 
 dependencies {
+  testImplementation 'de.thetaphi:forbiddenapis:3.8'
   testImplementation libs.bundles.junit5
   testImplementation libs.bundles.jmc
   testImplementation libs.commons.math
@@ -29,6 +30,11 @@ forbiddenApisMain_java11 {
   failOnMissingClasses = false
 }
 
+test {
+  useJUnit()
+  useJUnitPlatform()
+}
+
 idea {
   module {
     jdkName = '11'

dd-java-agent/instrumentation/exception-profiling/src/main/java/datadog/exceptions/instrumentation/KnownExcludesInstrumentation.java

@@ -0,0 +1,37 @@
+package datadog.exceptions.instrumentation;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Platform;
+
+/**
+ * Provides instrumentation to exclude exception types known to be used for control flow or
+ * 'connection leak' detection.
+ */
+@AutoService(InstrumenterModule.class)
+public final class KnownExcludesInstrumentation extends InstrumenterModule.Profiling
+    implements Instrumenter.ForBootstrap, Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public KnownExcludesInstrumentation() {
+    // this instrumentation is controlled together with 'throwables' instrumentation
+    super("throwables");
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return Platform.hasJfr() && super.isEnabled();
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(isConstructor(), packageName + ".ExclusionAdvice");
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {"com.zaxxer.hikari.pool.ProxyLeakTask"};
+  }
+}

dd-java-agent/instrumentation/exception-profiling/src/main/java11/datadog/exceptions/instrumentation/ExclusionAdvice.java

@@ -0,0 +1,16 @@
+package datadog.exceptions.instrumentation;
+
+import datadog.trace.bootstrap.instrumentation.jfr.exceptions.ExceptionProfiling;
+import net.bytebuddy.asm.Advice;
+
+public class ExclusionAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static void onEnter() {
+    ExceptionProfiling.Exclusion.enter();
+  }
+
+  @Advice.OnMethodExit(suppress = Throwable.class)
+  public static void onExit(@Advice.This final Object t) {
+    ExceptionProfiling.Exclusion.exit();
+  }
+}

dd-java-agent/instrumentation/exception-profiling/src/main/java11/datadog/exceptions/instrumentation/ThrowableInstanceAdvice.java

@@ -12,6 +12,9 @@
 public class ThrowableInstanceAdvice {
   @Advice.OnMethodExit(suppress = Throwable.class)
   public static void onExit(@Advice.This final Object t) {
+    if (ExceptionProfiling.Exclusion.isEffective()) {
+      return;
+    }
     /*
      * This instrumentation handler is sensitive to any throwables thrown from its body -
      * it will go into infinite loop of trying to handle the new throwable instance and generating
@@ -40,7 +43,7 @@ public static void onExit(@Advice.This final Object t) {
       }
       /*
        * JFR will assign the stacktrace depending on the place where the event is committed.
-       * Therefore we need to commit the event here, right in the 'Exception' constructor
+       * Therefore, we need to commit the event here, right in the 'Exception' constructor
        */
       final ExceptionSampleEvent event = ExceptionProfiling.getInstance().process((Throwable) t);
       if (event != null && event.shouldCommit()) {

dd-java-agent/instrumentation/exception-profiling/src/test/groovy/datadog/trace/bootstrap/instrumentation/jfr/exceptions/KnownExcludesForkedTest.groovy

@@ -0,0 +1,56 @@
+import com.zaxxer.hikari.pool.ProxyLeakTask
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.bootstrap.instrumentation.jfr.InstrumentationBasedProfiling
+import jdk.jfr.Recording
+import org.openjdk.jmc.common.item.Attribute
+import org.openjdk.jmc.common.item.IAttribute
+import org.openjdk.jmc.common.item.ItemFilters
+import org.openjdk.jmc.common.unit.UnitLookup
+import org.openjdk.jmc.flightrecorder.JfrLoaderToolkit
+import spock.lang.Shared
+
+import java.nio.file.Files
+
+class KnownExcludesForkedTest extends AgentTestRunner {
+  private static final IAttribute<String> TYPE =
+  Attribute.attr("type", "type", "Exception type", UnitLookup.PLAIN_TEXT)
+
+  @Shared
+  Recording recording
+
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig("profiling.enabled", "true")
+    injectSysConfig("profiling.start-force-first", "true")
+  }
+
+  def setupSpec() {
+    recording = new Recording()
+    recording.enable("datadog.ExceptionCount")
+    recording.start()
+    InstrumentationBasedProfiling.enableInstrumentationBasedProfiling()
+  }
+
+  def "obey excluded"() {
+    when:
+    println("Generating exceptions ...")
+    for (int i = 0; i < 50; i++) {
+      new ProxyLeakTask()
+      new NullPointerException()
+    }
+    println("Exceptions generated")
+
+    def tempPath = Files.createTempDirectory("test-recording")
+    def recFile = tempPath.resolve(KnownExcludesForkedTest.name.replace('.', '_') + ".jfr")
+    recFile.toFile().deleteOnExit()
+    recording.dump(recFile)
+    recording.stop()
+
+    def events = JfrLoaderToolkit.loadEvents(recFile.toFile()).apply(ItemFilters.type("datadog.ExceptionCount"))
+
+    then:
+    events.apply(ItemFilters.equals(TYPE, NullPointerException.canonicalName)).hasItems()
+    !events.apply(ItemFilters.equals(TYPE, ProxyLeakTask.canonicalName)).hasItems()
+  }
+}

dd-java-agent/instrumentation/exception-profiling/src/test/java/com/zaxxer/hikari/pool/ProxyLeakTask.java

@@ -0,0 +1,21 @@
+package com.zaxxer.hikari.pool;
+
+public class ProxyLeakTask extends Exception {
+  private static final long serialVersionUID = 1L;
+
+  public ProxyLeakTask() {
+    super("Proxy leak detected");
+  }
+
+  public ProxyLeakTask(String message) {
+    super(message);
+  }
+
+  public ProxyLeakTask(String message, Throwable cause) {
+    super(message, cause);
+  }
+
+  public ProxyLeakTask(Throwable cause) {
+    super(cause);
+  }
+}