Merge pull request #322 from DataDog/gary/hikari-handling

cache JDBC dbInfo after connection constructor throws exception

dd-java-agent-ittests/dd-java-agent-ittests.gradle

@@ -23,10 +23,6 @@ dependencies {
   testCompile group: 'org.apache.httpcomponents', name: 'httpclient', version: '4.3'
   testCompile group: 'com.squareup.okhttp3', name: 'okhttp', version: '3.6.0'
 
-  // JDBC tests:
-  testCompile group: 'com.h2database', name: 'h2', version: '1.4.196'
-  testCompile group: 'org.hsqldb', name: 'hsqldb', version: '2.3.+'
-  testCompile group: 'org.apache.derby', name: 'derby', version: '10.12.1.1'
 }
 
 tasks.withType(Test) {

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

@@ -7,4 +7,12 @@ dependencies {
   compile deps.opentracing
   annotationProcessor deps.autoservice
   implementation deps.autoservice
+
+  testCompile project(':dd-java-agent:testing')
+  // jdbc unit testing
+  testCompile group: 'com.h2database', name: 'h2', version: '1.4.197'
+  testCompile group: 'org.apache.derby', name: 'derby', version: '10.12.1.1'
+  // not using hsqldb 2.4 because org/hsqldb/jdbc/JDBCDriver : Unsupported major.minor version 52.0
+  testCompile group: 'org.hsqldb', name: 'hsqldb', version: '2.3.+'
+
 }

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

@@ -1,7 +1,6 @@
 package datadog.trace.instrumentation.jdbc;
 
 import static net.bytebuddy.matcher.ElementMatchers.failSafe;
-import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
 import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
@@ -13,11 +12,9 @@ import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.DDAdvice;
 import datadog.trace.agent.tooling.DDTransformers;
 import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.JDBCMaps;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
-import java.sql.SQLException;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 
@@ -40,8 +37,6 @@ public final class ConnectionInstrumentation extends Instrumenter.Configurable {
                         .and(takesArgument(0, String.class))
                         .and(returns(PreparedStatement.class)),
                     ConnectionPrepareAdvice.class.getName()))
-        .transform(
-            DDAdvice.create().advice(isConstructor(), ConnectionConstructorAdvice.class.getName()))
         .asDecorator();
   }
 
@@ -52,34 +47,4 @@ public final class ConnectionInstrumentation extends Instrumenter.Configurable {
       JDBCMaps.preparedStatements.put(statement, sql);
     }
   }
-
-  public static class ConnectionConstructorAdvice {
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static int constructorEnter() {
-      // We use this to make sure we only apply the exit instrumentation
-      // after the constructors are done calling their super constructors.
-      return CallDepthThreadLocalMap.incrementCallDepth(Connection.class);
-    }
-
-    // Since we're instrumenting the constructor, we can't add onThrowable.
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void constructorExit(
-        @Advice.Enter final int depth, @Advice.This final Connection connection)
-        throws SQLException {
-      if (depth == 0) {
-        CallDepthThreadLocalMap.reset(Connection.class);
-        final String url = connection.getMetaData().getURL();
-        if (url != null) {
-          // Remove end of url to prevent passwords from leaking:
-          final String sanitizedURL = url.replaceAll("[?;].*", "");
-          final String type = url.split(":", -1)[1];
-          String user = connection.getMetaData().getUserName();
-          if (user != null && user.trim().equals("")) {
-            user = null;
-          }
-          JDBCMaps.connectionInfo.put(connection, new JDBCMaps.DBInfo(sanitizedURL, type, user));
-        }
-      }
-    }
-  }
 }

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

@@ -23,13 +23,13 @@ import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracer;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
+import java.sql.SQLException;
 import java.util.Collections;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 
 @AutoService(Instrumenter.class)
 public final class PreparedStatementInstrumentation extends Instrumenter.Configurable {
-
   public PreparedStatementInstrumentation() {
     super("jdbc");
   }
@@ -65,9 +65,39 @@ public final class PreparedStatementInstrumentation extends Instrumenter.Configu
       }
 
       JDBCMaps.DBInfo dbInfo = JDBCMaps.connectionInfo.get(connection);
-      if (dbInfo == null) {
-        dbInfo = JDBCMaps.DBInfo.DEFAULT;
+      /**
+       * Logic to get the DBInfo from a JDBC Connection, if the connection was never seen before,
+       * the connectionInfo map will return null and will attempt to extract DBInfo from the
+       * connection. If the DBInfo can't be extracted, then the connection will be stored with the
+       * DEFAULT DBInfo as the value in the connectionInfo map to avoid retry overhead.
+       *
+       * <p>This should be a util method to be shared between PreparedStatementInstrumentation and
+       * StatementInstrumentation class, but java.sql.* are on the platform classloaders in Java 9,
+       * which prevents us from referencing them in the bootstrap utils.
+       */
+      {
+        if (dbInfo == null) {
+          try {
+            final String url = connection.getMetaData().getURL();
+            if (url != null) {
+              // Remove end of url to prevent passwords from leaking:
+              final String sanitizedURL = url.replaceAll("[?;].*", "");
+              final String type = url.split(":", -1)[1];
+              String user = connection.getMetaData().getUserName();
+              if (user != null && user.trim().equals("")) {
+                user = null;
+              }
+              dbInfo = new JDBCMaps.DBInfo(sanitizedURL, type, user);
+            } else {
+              dbInfo = JDBCMaps.DBInfo.DEFAULT;
+            }
+          } catch (SQLException se) {
+            dbInfo = JDBCMaps.DBInfo.DEFAULT;
+          }
+          JDBCMaps.connectionInfo.put(connection, dbInfo);
+        }
       }
+
       final Scope scope =
           GlobalTracer.get().buildSpan(dbInfo.getType() + ".query").startActive(true);
 

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

@@ -22,6 +22,7 @@ import io.opentracing.noop.NoopScopeManager.NoopScope;
 import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracer;
 import java.sql.Connection;
+import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.Collections;
 import net.bytebuddy.agent.builder.AgentBuilder;
@@ -65,15 +66,43 @@ public final class StatementInstrumentation extends Instrumenter.Configurable {
       }
 
       JDBCMaps.DBInfo dbInfo = JDBCMaps.connectionInfo.get(connection);
-      if (dbInfo == null) {
-        dbInfo = JDBCMaps.DBInfo.DEFAULT;
+      /**
+       * Logic to get the DBInfo from a JDBC Connection, if the connection was never seen before,
+       * the connectionInfo map will return null and will attempt to extract DBInfo from the
+       * connection. If the DBInfo can't be extracted, then the connection will be stored with the
+       * DEFAULT DBInfo as the value in the connectionInfo map to avoid retry overhead.
+       *
+       * <p>This should be a util method to be shared between PreparedStatementInstrumentation and
+       * StatementInstrumentation class, but java.sql.* are on the platform classloaders in Java 9,
+       * which prevents us from referencing them in the bootstrap utils.
+       */
+      {
+        if (dbInfo == null) {
+          try {
+            final String url = connection.getMetaData().getURL();
+            if (url != null) {
+              // Remove end of url to prevent passwords from leaking:
+              final String sanitizedURL = url.replaceAll("[?;].*", "");
+              final String type = url.split(":", -1)[1];
+              String user = connection.getMetaData().getUserName();
+              if (user != null && user.trim().equals("")) {
+                user = null;
+              }
+              dbInfo = new JDBCMaps.DBInfo(sanitizedURL, type, user);
+            } else {
+              dbInfo = JDBCMaps.DBInfo.DEFAULT;
+            }
+          } catch (SQLException se) {
+            dbInfo = JDBCMaps.DBInfo.DEFAULT;
+          }
+          JDBCMaps.connectionInfo.put(connection, dbInfo);
+        }
       }
 
       final Scope scope =
           GlobalTracer.get().buildSpan(dbInfo.getType() + ".query").startActive(true);
 
       final Span span = scope.span();
-
       Tags.DB_TYPE.set(span, dbInfo.getType());
       Tags.SPAN_KIND.set(span, Tags.SPAN_KIND_CLIENT);
       Tags.COMPONENT.set(span, "java-jdbc-statement");

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

@@ -0,0 +1,296 @@
+import java.sql.Array
+import java.sql.Blob
+import java.sql.CallableStatement
+import java.sql.Clob
+import java.sql.Connection
+import java.sql.DatabaseMetaData
+import java.sql.NClob
+import java.sql.PreparedStatement
+import java.sql.SQLClientInfoException
+import java.sql.SQLException
+import java.sql.SQLWarning
+import java.sql.SQLXML
+import java.sql.Savepoint
+import java.sql.Statement
+import java.sql.Struct
+import java.util.concurrent.Executor
+
+
+/**
+ * A JDBC connection class that throws an exception in the constructor, used to test
+ */
+class DummyThrowingConnection implements Connection {
+  DummyThrowingConnection() {
+    throw new RuntimeException("Dummy exception")
+  }
+
+  @Override
+  Statement createStatement() throws SQLException {
+    return null
+  }
+
+  @Override
+  PreparedStatement prepareStatement(String sql) throws SQLException {
+    return null
+  }
+
+  @Override
+  CallableStatement prepareCall(String sql) throws SQLException {
+    return null
+  }
+
+  @Override
+  String nativeSQL(String sql) throws SQLException {
+    return null
+  }
+
+  @Override
+  void setAutoCommit(boolean autoCommit) throws SQLException {
+
+  }
+
+  @Override
+  boolean getAutoCommit() throws SQLException {
+    return false
+  }
+
+  @Override
+  void commit() throws SQLException {
+
+  }
+
+  @Override
+  void rollback() throws SQLException {
+
+  }
+
+  @Override
+  void close() throws SQLException {
+
+  }
+
+  @Override
+  boolean isClosed() throws SQLException {
+    return false
+  }
+
+  @Override
+  DatabaseMetaData getMetaData() throws SQLException {
+    return null
+  }
+
+  @Override
+  void setReadOnly(boolean readOnly) throws SQLException {
+
+  }
+
+  @Override
+  boolean isReadOnly() throws SQLException {
+    return false
+  }
+
+  @Override
+  void setCatalog(String catalog) throws SQLException {
+
+  }
+
+  @Override
+  String getCatalog() throws SQLException {
+    return null
+  }
+
+  @Override
+  void setTransactionIsolation(int level) throws SQLException {
+
+  }
+
+  @Override
+  int getTransactionIsolation() throws SQLException {
+    return 0
+  }
+
+  @Override
+  SQLWarning getWarnings() throws SQLException {
+    return null
+  }
+
+  @Override
+  void clearWarnings() throws SQLException {
+
+  }
+
+  @Override
+  Statement createStatement(int resultSetType, int resultSetConcurrency) throws SQLException {
+    return null
+  }
+
+  @Override
+  PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
+    return null
+  }
+
+  @Override
+  CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
+    return null
+  }
+
+  @Override
+  Map<String, Class<?>> getTypeMap() throws SQLException {
+    return null
+  }
+
+  @Override
+  void setTypeMap(Map<String, Class<?>> map) throws SQLException {
+
+  }
+
+  @Override
+  void setHoldability(int holdability) throws SQLException {
+
+  }
+
+  @Override
+  int getHoldability() throws SQLException {
+    return 0
+  }
+
+  @Override
+  Savepoint setSavepoint() throws SQLException {
+    return null
+  }
+
+  @Override
+  Savepoint setSavepoint(String name) throws SQLException {
+    return null
+  }
+
+  @Override
+  void rollback(Savepoint savepoint) throws SQLException {
+
+  }
+
+  @Override
+  void releaseSavepoint(Savepoint savepoint) throws SQLException {
+
+  }
+
+  @Override
+  Statement createStatement(int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
+    return null
+  }
+
+  @Override
+  PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
+    return null
+  }
+
+  @Override
+  CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
+    return null
+  }
+
+  @Override
+  PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
+    return null
+  }
+
+  @Override
+  PreparedStatement prepareStatement(String sql, int[] columnIndexes) throws SQLException {
+    return null
+  }
+
+  @Override
+  PreparedStatement prepareStatement(String sql, String[] columnNames) throws SQLException {
+    return null
+  }
+
+  @Override
+  Clob createClob() throws SQLException {
+    return null
+  }
+
+  @Override
+  Blob createBlob() throws SQLException {
+    return null
+  }
+
+  @Override
+  NClob createNClob() throws SQLException {
+    return null
+  }
+
+  @Override
+  SQLXML createSQLXML() throws SQLException {
+    return null
+  }
+
+  @Override
+  boolean isValid(int timeout) throws SQLException {
+    return false
+  }
+
+  @Override
+  void setClientInfo(String name, String value) throws SQLClientInfoException {
+
+  }
+
+  @Override
+  void setClientInfo(Properties properties) throws SQLClientInfoException {
+
+  }
+
+  @Override
+  String getClientInfo(String name) throws SQLException {
+    return null
+  }
+
+  @Override
+  Properties getClientInfo() throws SQLException {
+    return null
+  }
+
+  @Override
+  Array createArrayOf(String typeName, Object[] elements) throws SQLException {
+    return null
+  }
+
+  @Override
+  Struct createStruct(String typeName, Object[] attributes) throws SQLException {
+    return null
+  }
+
+  @Override
+  void setSchema(String schema) throws SQLException {
+
+  }
+
+  @Override
+  String getSchema() throws SQLException {
+    return null
+  }
+
+  @Override
+  void abort(Executor executor) throws SQLException {
+
+  }
+
+  @Override
+  void setNetworkTimeout(Executor executor, int milliseconds) throws SQLException {
+
+  }
+
+  @Override
+  int getNetworkTimeout() throws SQLException {
+    return 0
+  }
+
+  @Override
+  def <T> T unwrap(Class<T> iface) throws SQLException {
+    return null
+  }
+
+  @Override
+  boolean isWrapperFor(Class<?> iface) throws SQLException {
+    return false
+  }
+}

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

@@ -1,31 +1,26 @@
-package datadog.trace.agent.integration.jdbc
-
-import datadog.opentracing.DDTracer
-import datadog.trace.agent.test.IntegrationTestUtils
-import datadog.trace.common.writer.ListWriter
+import datadog.trace.agent.test.AgentTestRunner
 import org.apache.derby.jdbc.EmbeddedDriver
 import org.h2.Driver
 import org.hsqldb.jdbc.JDBCDriver
 import spock.lang.Shared
-import spock.lang.Specification
+import spock.lang.Unroll
 
 import java.sql.Connection
 import java.sql.PreparedStatement
 import java.sql.ResultSet
 import java.sql.Statement
 
-class JDBCInstrumentationTest extends Specification {
-
-  final ListWriter writer = new ListWriter()
-  final DDTracer tracer = new DDTracer(writer)
+class JDBCInstrumentationTest extends AgentTestRunner {
+  @Shared
+  def dbName = "jdbcUnitTest"
 
   @Shared
   private Map<String, Connection> connections
 
   def setupSpec() {
-    Connection h2Connection = new Driver().connect("jdbc:h2:mem:integ-test", null)
-    Connection hsqlConnection = new JDBCDriver().connect("jdbc:hsqldb:mem:integTest", null)
-    Connection derbyConnection = new EmbeddedDriver().connect("jdbc:derby:memory:integTest;create=true", null)
+    Connection h2Connection = new Driver().connect("jdbc:h2:mem:" + dbName, null)
+    Connection hsqlConnection = new JDBCDriver().connect("jdbc:hsqldb:mem:" + dbName, null)
+    Connection derbyConnection = new EmbeddedDriver().connect("jdbc:derby:memory:" + dbName + ";create=true", null)
 
     connections = [
       h2    : h2Connection,
@@ -40,11 +35,7 @@ class JDBCInstrumentationTest extends Specification {
     }
   }
 
-  def setup() {
-    IntegrationTestUtils.registerOrReplaceGlobalTracer(tracer)
-    writer.start()
-  }
-
+  @Unroll
   def "basic statement on #driver generates spans"() {
     setup:
     Statement statement = connection.createStatement()
@@ -53,9 +44,9 @@ class JDBCInstrumentationTest extends Specification {
     expect:
     resultSet.next()
     resultSet.getInt(1) == 3
-    writer.size() == 1
+    TEST_WRITER.size() == 1
 
-    def trace = writer.firstTrace()
+    def trace = TEST_WRITER.firstTrace()
     trace.size() == 1
     def span = trace[0]
 
@@ -90,6 +81,7 @@ class JDBCInstrumentationTest extends Specification {
     "hsqldb" | connections.get("hsqldb") | "SA"     | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
   }
 
+  @Unroll
   def "prepared statement execute on #driver generates a span"() {
     setup:
     PreparedStatement statement = connection.prepareStatement(query)
@@ -99,9 +91,9 @@ class JDBCInstrumentationTest extends Specification {
     expect:
     resultSet.next()
     resultSet.getInt(1) == 3
-    writer.size() == 1
+    TEST_WRITER.size() == 1
 
-    def trace = writer.firstTrace()
+    def trace = TEST_WRITER.firstTrace()
     trace.size() == 1
     def span = trace[0]
 
@@ -136,6 +128,7 @@ class JDBCInstrumentationTest extends Specification {
     "hsqldb" | connections.get("hsqldb") | "SA"     | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
   }
 
+  @Unroll
   def "prepared statement query on #driver generates a span"() {
     setup:
     PreparedStatement statement = connection.prepareStatement(query)
@@ -144,9 +137,9 @@ class JDBCInstrumentationTest extends Specification {
     expect:
     resultSet.next()
     resultSet.getInt(1) == 3
-    writer.size() == 1
+    TEST_WRITER.size() == 1
 
-    def trace = writer.firstTrace()
+    def trace = TEST_WRITER.firstTrace()
     trace.size() == 1
     def span = trace[0]
 
@@ -181,6 +174,7 @@ class JDBCInstrumentationTest extends Specification {
     "hsqldb" | connections.get("hsqldb") | "SA"     | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
   }
 
+  @Unroll
   def "statement update on #driver generates a span"() {
     setup:
     Statement statement = connection.createStatement()
@@ -190,9 +184,9 @@ class JDBCInstrumentationTest extends Specification {
     !statement.execute(sql)
     statement.updateCount == 0
 
-    writer.size() == 1
+    TEST_WRITER.size() == 1
 
-    def trace = writer.firstTrace()
+    def trace = TEST_WRITER.firstTrace()
     trace.size() == 1
     def span = trace[0]
 
@@ -227,6 +221,7 @@ class JDBCInstrumentationTest extends Specification {
     "hsqldb" | connections.get("hsqldb") | "SA"     | "CREATE TABLE PUBLIC.S_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
   }
 
+  @Unroll
   def "prepared statement update on #driver generates a span"() {
     setup:
     def sql = connection.nativeSQL(query)
@@ -234,9 +229,9 @@ class JDBCInstrumentationTest extends Specification {
 
     expect:
     statement.executeUpdate() == 0
-    writer.size() == 1
+    TEST_WRITER.size() == 1
 
-    def trace = writer.firstTrace()
+    def trace = TEST_WRITER.firstTrace()
     trace.size() == 1
     def span = trace[0]
 
@@ -270,4 +265,77 @@ class JDBCInstrumentationTest extends Specification {
     "derby"  | connections.get("derby")  | "APP"    | "CREATE TABLE PS_DERBY (id INTEGER not NULL, PRIMARY KEY ( id ))"
     "hsqldb" | connections.get("hsqldb") | "SA"     | "CREATE TABLE PUBLIC.PS_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
   }
+  
+  @Unroll
+  def "connection constructor throwing then generating correct spans after recovery using #driver connection (prepare statement = #prepareStatement)"() {
+    setup:
+    Connection connection = null
+
+    when:
+    try {
+      connection = new DummyThrowingConnection()
+    } catch (Exception e) {
+      connection = driverClass.connect(url, null)
+    }
+
+    Statement statement = null
+    ResultSet rs = null
+    if (prepareStatement) {
+      statement = connection.prepareStatement(query)
+      rs = statement.executeQuery()
+    } else {
+      statement = connection.createStatement()
+      rs = statement.executeQuery(query)
+    }
+
+    then:
+    rs.next()
+    rs.getInt(1) == 3
+    TEST_WRITER.size() == 1
+
+    def trace = TEST_WRITER.firstTrace()
+    trace.size() == 1
+
+    and:
+    def span = trace[0]
+    span.context().operationName == "${driver}.query"
+    span.serviceName == driver
+    span.resourceName == query
+    span.type == "sql"
+    !span.context().getErrorFlag()
+    span.context().parentId == 0
+
+    def tags = span.context().tags
+    tags["db.type"] == driver
+    tags["db.user"] == user
+    tags["span.kind"] == "client"
+    if (prepareStatement) {
+      tags["component"] == "java-jdbc-prepared_statement"
+    } else {
+      tags["component"] == "java-jdbc-statement"
+    }
+
+    tags["db.jdbc.url"].contains(driver)
+    tags["span.origin.type"] != null
+
+    tags["thread.name"] != null
+    tags["thread.id"] != null
+    tags.size() == user == null ? 7 : 8
+
+    cleanup:
+    if (statement != null) {
+      statement.close()
+    }
+    if (connection != null) {
+      connection.close()
+    }
+
+    where:
+    prepareStatement | driver   | driverClass          | url                                            | user  | query
+    true             | "h2"     | new Driver()         | "jdbc:h2:mem:" + dbName                        | null  | "SELECT 3;"
+    true             | "derby"  | new EmbeddedDriver() | "jdbc:derby:memory:" + dbName + ";create=true" | "APP" | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
+    false            | "h2"     | new Driver()         | "jdbc:h2:mem:" + dbName                        | null  | "SELECT 3;"
+    false            | "derby"  | new EmbeddedDriver() | "jdbc:derby:memory:" + dbName + ";create=true" | "APP" | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
+  }
+
 }