Apache DBCP2: Add a span when waiting on the pool

Author: deejgregor

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,54 @@
+package datadog.trace.instrumentation.jdbc;
+
+import com.google.auto.service.*;
+import datadog.trace.agent.tooling.*;
+import datadog.trace.bootstrap.instrumentation.api.*;
+import net.bytebuddy.asm.*;
+import org.slf4j.*;
+
+import java.util.concurrent.*;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.*;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+@AutoService(InstrumenterModule.class)
+public final class Dbcp2LinkedBlockingDequeInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice {
+
+  private static final Logger log = LoggerFactory.getLogger(Dbcp2LinkedBlockingDequeInstrumentation.class);
+
+  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() {
+      return startSpan("dbcp2", POOL_WAITING);
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Enter final AgentSpan span) {
+      span.finish();
+    }
+  }
+}

dd-java-agent/instrumentation/jdbc/src/test/groovy/SaturatedPoolBlockingTest.groovy

@@ -1,25 +1,26 @@
 import datadog.trace.agent.test.AgentTestRunner
 import com.zaxxer.hikari.HikariConfig
 import com.zaxxer.hikari.HikariDataSource
+import org.apache.commons.dbcp2.BasicDataSource
 import test.TestDataSource
+import test.TestDriver
 
+import javax.sql.DataSource
+import java.sql.SQLException
 import java.sql.SQLTimeoutException
 import java.sql.SQLTransientConnectionException
+import java.time.Duration
 
 /**
  * Ideas taken from Hikari's com.zaxxer.hikari.pool.TestSaturatedPool830.
  */
 class SaturatedPoolBlockingTest extends AgentTestRunner {
-  def "saturated pool test"(int connectionTimeout, Long exhaustPoolForMillis, int expectedWaitingSpans, boolean expectedTimeout) {
+  public static final int CONNECTION_TIMEOUT = 1000
+
+  def "saturated pool test"(Closure<DataSource> createDataSource, Long exhaustPoolForMillis, int expectedWaitingSpans, boolean expectedTimeout) {
     setup:
     TEST_WRITER.setFilter((trace) -> trace.get(0).getOperationName() == "test.when")
-
-    final HikariConfig config = new HikariConfig()
-    config.setPoolName("testPool")
-    config.setMaximumPoolSize(1)
-    config.setConnectionTimeout(connectionTimeout)
-    config.setDataSourceClassName(TestDataSource.class.getName())
-    final HikariDataSource ds = new HikariDataSource(config)
+    final DataSource ds = createDataSource()
 
     when:
     if (exhaustPoolForMillis != null) {
@@ -43,6 +44,12 @@ class SaturatedPoolBlockingTest extends AgentTestRunner {
       }
     } catch (SQLTimeoutException ignored) { // Hikari, older
       timedOut = true
+    } catch (SQLException e) {
+      if (e.getMessage().contains("pool error Timeout waiting for idle object")) { // dbcp2
+        timedOut = true
+      } else {
+        throw e
+      }
     }
     span.finish()
 
@@ -59,10 +66,32 @@ class SaturatedPoolBlockingTest extends AgentTestRunner {
     }
 
     where:
-    connectionTimeout | exhaustPoolForMillis | expectedWaitingSpans | expectedTimeout
-    1000              | null                 | 0                    | false
-    1000              | null                 | 0                    | false
-    1000              | 500                  | 1                    | false
-    1000              | 1500                 | 1                    | true
+    createDataSource       | exhaustPoolForMillis | expectedWaitingSpans | expectedTimeout
+    this.&hikariDataSource | null                 | 0                    | false
+    this.&hikariDataSource | null                 | 0                    | false
+    this.&hikariDataSource | 500                  | 1                    | false
+    this.&hikariDataSource | 1500                 | 1                    | true
+    this.&dbcp2DataSource  | null                 | 0                    | false
+    this.&dbcp2DataSource  | null                 | 0                    | false
+    this.&dbcp2DataSource  | 500                  | 1                    | false
+    this.&dbcp2DataSource  | 1500                 | 1                    | true
+  }
+
+  private static DataSource hikariDataSource() {
+    final HikariConfig config = new HikariConfig()
+    config.setPoolName("testPool")
+    config.setMaximumPoolSize(1)
+    config.setConnectionTimeout(CONNECTION_TIMEOUT)
+    config.setDataSourceClassName(TestDataSource.class.getName())
+    return new HikariDataSource(config)
+  }
+
+  private static DataSource dbcp2DataSource() {
+    final BasicDataSource ds = new BasicDataSource()
+    ds.setMaxTotal(1)
+    ds.setMaxWait(Duration.ofMillis(CONNECTION_TIMEOUT))
+    ds.setDriverClassName(TestDriver.class.getName())
+    ds.start()
+    return ds
   }
 }

dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestConnection.groovy

@@ -227,7 +227,7 @@ class TestConnection implements Connection {
 
   @Override
   boolean isValid(int timeout) throws SQLException {
-    return false
+    return true
   }
 
   @Override

dd-java-agent/instrumentation/jdbc/src/test/groovy/test/TestDriver.groovy

@@ -15,7 +15,7 @@ class TestDriver implements Driver {
 
   @Override
   boolean acceptsURL(String url) throws SQLException {
-    return false
+    return true
   }
 
   @Override