Add a span when waiting for an available database connection from a pool

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

@@ -27,6 +27,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,56 @@
+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 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.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-datasource");
+  }
+
+  @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 void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("pollFirst").and(takesArguments(1)),
+        Dbcp2LinkedBlockingDequeInstrumentation.class.getName() + "$PollFirstAdvice");
+  }
+
+  public static class PollFirstAdvice {
+    private static final String POOL_WAITING = "pool.waiting";
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentSpan onEnter() {
+      if (CallDepthThreadLocalMap.getCallDepth(Dbcp2LinkedBlockingDequeInstrumentation.class) > 0) {
+        return startSpan("dbcp2", POOL_WAITING);
+      } else {
+        return null;
+      }
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Enter final AgentSpan span) {
+      if (span != null) {
+        span.finish();
+      }
+    }
+  }
+}

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

@@ -0,0 +1,45 @@
+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.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-datasource");
+  }
+
+  @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 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(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+  }
+}

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

@@ -0,0 +1,46 @@
+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.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-datasource");
+  }
+
+  @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 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(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+  }
+}

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

@@ -0,0 +1,45 @@
+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.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-datasource");
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.apache.commons.dbcp2.PoolingDataSource", // standalone
+      "org.apache.tomcat.dbcp.dbcp2.PoolingDataSource" // bundled with Tomcat
+    };
+  }
+
+  @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(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+  }
+}

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

@@ -0,0 +1,44 @@
+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.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-datasource");
+  }
+
+  @Override
+  public String[] knownMatchingTypes() {
+    return new String[] {
+      "org.apache.commons.dbcp2.PoolingDriver", // standalone
+      "org.apache.tomcat.dbcp.dbcp2.PoolingDriver" // bundled with Tomcat
+    };
+  }
+
+  @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(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+  }
+}

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

@@ -0,0 +1,46 @@
+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.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-datasource");
+  }
+
+  @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 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(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit() {
+      CallDepthThreadLocalMap.decrementCallDepth(Dbcp2LinkedBlockingDequeInstrumentation.class);
+    }
+  }
+}