🪞 9251 - Add a span when waiting for an available database connection from a pool (#9636)

* HikariCP: Add a span when waiting on the pool
* Apache DBCP2: Add a span when waiting on the pool
* Disable pool.waiting by default, other tweaks
* Use dd.trace.experimental.jdbc.pool.waiting.enabled=true to enable.
* Change instrumentation name from jdbc-datasource to jdbc.
* Record exceptions.
* Use default instrumentation name (jdbc).
* Set resource name to {dbcp2,hikari}.waiting
* Docs: Update formatting settings for later versions of IntelliJ
* Merge SaturatedPoolBlockingTest into JDBCInstrumentationTestBase
* Add PoolWaitingDecorator.
* Cleanup Hikari instrumentation
* Comment on the overall structure and highlight related instrumentations.
* Avoid using reflection and breakup HikariConcurrentBagInstrumentation.ConstructorAdvice
  into separate, simpler instrumentations: HikariPoolInstrumentation and
  HikariConcurrentBagHandoffQueueInstrumentation.
* Move HikariBlockedTrackingSynchronousQueue into HikariConcurrentBagHandoffQueueInstrumentation.
* Make sure the instrumentationName is consistently "jdbc".
* Fix: injected bytecode shouldn't link back to instrumentation class
* Add dbcp2/hikari aliases for pool waiting instrumentation
* Avoid leaking thread-locals in HikariBlockedTracker

---------

Co-authored-by: DJ Gregor <[email protected]>
Co-authored-by: Stuart McCulloch <[email protected]>

Author: mcculls

CONTRIBUTING.md

@@ -52,6 +52,7 @@ The default JVM to build and run tests from the command line should be Java 8.
   * `Project Structure` -> `Project` -> `SDK` -> `Download JDK...` -> `Version: 1.8` -> `Download`
 * Configure Java and Groovy import formatting:  
   * `Settings...` ->`Editor` > `Code Style` > `Java` > `Imports`
+    * `Use single class import`: checked
     * `Class count to use import with '*'`: `9999` (some number sufficiently large that is unlikely to matter)
     * `Names count to use static import with '*'`: `9999`
     * Use the following import layout to ensure consistency with google-java-format:

dd-java-agent/instrumentation/jdbc/build.gradle

@@ -31,6 +31,7 @@ dependencies {
   testImplementation group: 'com.h2database', name: 'h2', version: '[1.3.168,1.3.169]'// first jdk 1.6 compatible
   testImplementation group: 'org.apache.derby', name: 'derby', version: '10.6.1.0'
   testImplementation group: 'org.hsqldb', name: 'hsqldb', version: '2.0.0'
+  testImplementation group: 'org.apache.commons', name: 'commons-dbcp2', version: '2.10.0'
 
   testImplementation group: 'org.apache.tomcat', name: 'tomcat-jdbc', version: '7.0.19'
   // tomcat needs this to run

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/Dbcp2LinkedBlockingDequeInstrumentation.java

@@ -0,0 +1,72 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.jdbc.PoolWaitingDecorator.DECORATE;
+import static datadog.trace.instrumentation.jdbc.PoolWaitingDecorator.POOL_WAITING;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public final class Dbcp2LinkedBlockingDequeInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public Dbcp2LinkedBlockingDequeInstrumentation() {
+    super("jdbc", "dbcp2");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isJdbcPoolWaitingEnabled();
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.apache.commons.pool2.impl.LinkedBlockingDeque", // standalone
+      "org.apache.tomcat.dbcp.pool2.impl.LinkedBlockingDeque" // bundled with Tomcat
+    };
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".PoolWaitingDecorator"};
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("pollFirst").and(takesArguments(1)),
+        Dbcp2LinkedBlockingDequeInstrumentation.class.getName() + "$PollFirstAdvice");
+  }
+
+  public static class PollFirstAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentSpan onEnter() {
+      if (CallDepthThreadLocalMap.getCallDepth(PoolWaitingDecorator.class) > 0) {
+        AgentSpan span = startSpan(POOL_WAITING);
+        DECORATE.afterStart(span);
+        span.setResourceName("dbcp2.waiting");
+        return span;
+      } else {
+        return null;
+      }
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Enter final AgentSpan span, @Advice.Thrown final Throwable throwable) {
+      if (span != null) {
+        DECORATE.onError(span, throwable);
+        span.finish();
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/Dbcp2ManagedConnectionInstrumentation.java

@@ -0,0 +1,56 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public final class Dbcp2ManagedConnectionInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public Dbcp2ManagedConnectionInstrumentation() {
+    super("jdbc", "dbcp2");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isJdbcPoolWaitingEnabled();
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.apache.commons.dbcp2.managed.ManagedConnection", // standalone
+      "org.apache.tomcat.dbcp.dbcp2.managed.ManagedConnection" // bundled with Tomcat
+    };
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".PoolWaitingDecorator"};
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("updateTransactionStatus"),
+        Dbcp2ManagedConnectionInstrumentation.class.getName() + "$UpdateTransactionStatusAdvice");
+  }
+
+  public static class UpdateTransactionStatusAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter() {
+      CallDepthThreadLocalMap.incrementCallDepth(PoolWaitingDecorator.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(PoolWaitingDecorator.class);
+    }
+  }
+}

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/Dbcp2PerUserPoolDataSourceInstrumentation.java

@@ -0,0 +1,57 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public final class Dbcp2PerUserPoolDataSourceInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public Dbcp2PerUserPoolDataSourceInstrumentation() {
+    super("jdbc", "dbcp2");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isJdbcPoolWaitingEnabled();
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.apache.commons.dbcp2.datasources.PerUserPoolDataSource", // standalone
+      "org.apache.tomcat.dbcp.dbcp2.datasources.PerUserPoolDataSource" // bundled with Tomcat
+    };
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".PoolWaitingDecorator"};
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("getPooledConnectionAndInfo"),
+        Dbcp2PerUserPoolDataSourceInstrumentation.class.getName()
+            + "$GetPooledConnectionAndInfoAdvice");
+  }
+
+  public static class GetPooledConnectionAndInfoAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter() {
+      CallDepthThreadLocalMap.incrementCallDepth(PoolWaitingDecorator.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(PoolWaitingDecorator.class);
+    }
+  }
+}

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/Dbcp2PoolingDataSourceInstrumentation.java

@@ -0,0 +1,56 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public final class Dbcp2PoolingDataSourceInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public Dbcp2PoolingDataSourceInstrumentation() {
+    super("jdbc", "dbcp2");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isJdbcPoolWaitingEnabled();
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.apache.commons.dbcp2.PoolingDataSource", // standalone
+      "org.apache.tomcat.dbcp.dbcp2.PoolingDataSource" // bundled with Tomcat
+    };
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".PoolWaitingDecorator"};
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("getConnection"),
+        Dbcp2PoolingDataSourceInstrumentation.class.getName() + "$GetConnectionAdvice");
+  }
+
+  public static class GetConnectionAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter() {
+      CallDepthThreadLocalMap.incrementCallDepth(PoolWaitingDecorator.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(PoolWaitingDecorator.class);
+    }
+  }
+}

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/Dbcp2PoolingDriverInstrumentation.java

@@ -0,0 +1,55 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public final class Dbcp2PoolingDriverInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public Dbcp2PoolingDriverInstrumentation() {
+    super("jdbc", "dbcp2");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isJdbcPoolWaitingEnabled();
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.apache.commons.dbcp2.PoolingDriver", // standalone
+      "org.apache.tomcat.dbcp.dbcp2.PoolingDriver" // bundled with Tomcat
+    };
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".PoolWaitingDecorator"};
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("connect"), Dbcp2PoolingDriverInstrumentation.class.getName() + "$ConnectAdvice");
+  }
+
+  public static class ConnectAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter() {
+      CallDepthThreadLocalMap.incrementCallDepth(PoolWaitingDecorator.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(PoolWaitingDecorator.class);
+    }
+  }
+}

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/Dbcp2SharedPoolDataSourceInstrumentation.java

@@ -0,0 +1,57 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public final class Dbcp2SharedPoolDataSourceInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  public Dbcp2SharedPoolDataSourceInstrumentation() {
+    super("jdbc", "dbcp2");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return InstrumenterConfig.get().isJdbcPoolWaitingEnabled();
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.apache.commons.dbcp2.datasources.SharePoolDataSource", // standalone
+      "org.apache.tomcat.dbcp.dbcp2.datasources.SharedPoolPoolDataSource" // bundled with Tomcat
+    };
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".PoolWaitingDecorator"};
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("getPooledConnectionAndInfo"),
+        Dbcp2SharedPoolDataSourceInstrumentation.class.getName()
+            + "$GetPooledConnectionAndInfoAdvice");
+  }
+
+  public static class GetPooledConnectionAndInfoAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter() {
+      CallDepthThreadLocalMap.incrementCallDepth(PoolWaitingDecorator.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(PoolWaitingDecorator.class);
+    }
+  }
+}

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/HikariBlockedTracker.java

@@ -0,0 +1,20 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static java.lang.Boolean.TRUE;
+
+/** Shared blocked getConnection() tracking {@link ThreadLocal} for Hikari. */
+public class HikariBlockedTracker {
+  private static final ThreadLocal<Boolean> tracker = new ThreadLocal<>();
+
+  public static void clearBlocked() {
+    tracker.remove();
+  }
+
+  public static void setBlocked() {
+    tracker.set(TRUE);
+  }
+
+  public static boolean wasBlocked() {
+    return TRUE.equals(tracker.get());
+  }
+}