Add Hikari Pool Name tag (#7672)

* Update deprecated MSSQLServerContainer initialization to fix setupSpec error

* add HikariPoolName instrumentation and test for it

* Add seperate HikariDataSource instrumentation, broken muzzle dependency add

* Remove hikariDataSource instrumentation from generic DataSource instrumentation file

* add muzzle dependency

* revert line breaks and changes to unrelated files

* Add not working rework of hikaripoolname instrumentation with DBInfo and JDBCDecorator tag setting

* Add working version of intrumentation that passes statement span test

* Push draft not working code - prepared statement tests dont pass

* unwrap the connection before looking in the instrumentation context

* Refactor tag name to remove hikari specific mention, make generic

* Refactor tag name to remove hikari specific mention, make generic

* clean up print statements and commented out code

* fix tests in jdbc suite to check for pool name tag when testing with hikari

* Add null check before 'get'

* bump dependency lock

* remove non-functional new module that was previously added by mistake

---------

Co-authored-by: Andrea Marziali <[email protected]>

Author: jordan-wong

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

@@ -15,6 +15,7 @@ public class DBInfo {
   private final Integer port;
   private final String warehouse;
   private final String schema;
+  private String poolName;
 
   DBInfo(
       String type,
@@ -27,7 +28,8 @@ public class DBInfo {
       String host,
       Integer port,
       String warehouse,
-      String schema) {
+      String schema,
+      String poolName) {
     this.type = type;
     this.subtype = subtype;
     this.fullPropagationSupport = fullPropagationSupport;
@@ -39,6 +41,7 @@ public class DBInfo {
     this.port = port;
     this.warehouse = warehouse;
     this.schema = schema;
+    this.poolName = poolName;
   }
 
   public static class Builder {
@@ -55,6 +58,7 @@ public static class Builder {
     private String schema;
     private String host;
     private Integer port;
+    private String poolName;
 
     Builder() {}
 
@@ -69,7 +73,8 @@ public static class Builder {
         String host,
         Integer port,
         String warehouse,
-        String schema) {
+        String schema,
+        String poolName) {
       this.type = type;
       this.subtype = subtype;
       this.fullPropagationSupport = fullPropagationSupport;
@@ -81,6 +86,7 @@ public static class Builder {
       this.port = port;
       this.warehouse = warehouse;
       this.schema = schema;
+      this.poolName = poolName;
     }
 
     public Builder type(String type) {
@@ -141,6 +147,11 @@ public Builder port(Integer port) {
       return this;
     }
 
+    public Builder poolName(String poolName) {
+      this.poolName = poolName;
+      return this;
+    }
+
     public DBInfo build() {
       return new DBInfo(
           type,
@@ -153,7 +164,8 @@ public DBInfo build() {
           host,
           port,
           warehouse,
-          schema);
+          schema,
+          poolName);
     }
   }
 
@@ -201,6 +213,14 @@ public String getSchema() {
     return schema;
   }
 
+  public String getPoolName() {
+    return poolName;
+  }
+
+  public void setPoolName(String poolname) {
+    this.poolName = poolname;
+  }
+
   public Builder toBuilder() {
     return new Builder(
         type,
@@ -213,7 +233,8 @@ public Builder toBuilder() {
         host,
         port,
         warehouse,
-        schema);
+        schema,
+        poolName);
   }
 
   @Override
@@ -231,7 +252,8 @@ public boolean equals(Object o) {
         && Objects.equals(host, dbInfo.host)
         && Objects.equals(port, dbInfo.port)
         && Objects.equals(warehouse, dbInfo.warehouse)
-        && Objects.equals(schema, dbInfo.schema);
+        && Objects.equals(schema, dbInfo.schema)
+        && Objects.equals(poolName, dbInfo.poolName);
   }
 
   @Override

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

@@ -7,6 +7,7 @@ muzzle {
   pass {
     coreJdk()
     extraDependency 'com.ibm.db2:jcc:11.1.4.4'
+    extraDependency 'com.zaxxer:HikariCP:2.4.0'
   }
 }
 
@@ -19,6 +20,7 @@ addTestSuiteForDir('latestDepTest', 'test')
 addTestSuiteExtendingForDir('latestDepJava11Test', 'latestDepTest', 'test')
 
 dependencies {
+  compileOnly group: 'com.zaxxer', name: 'HikariCP', version: '2.4.0'
   testImplementation(testFixtures(project(':dd-java-agent:agent-iast')))
 
   // jdbc unit testing

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

@@ -56,7 +56,7 @@ com.squareup.okhttp3:logging-interceptor:3.12.12=jmhRuntimeClasspath,latestDepJa
 com.squareup.okhttp3:okhttp:3.12.12=jmhRuntimeClasspath,latestDepJava11TestCompileClasspath,latestDepJava11TestRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,oldH2TestCompileClasspath,oldH2TestRuntimeClasspath,oldPostgresTestCompileClasspath,oldPostgresTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath
 com.squareup.okio:okio:1.17.5=compileClasspath,csiCompileClasspath,instrumentPluginClasspath,jmhCompileClasspath,jmhRuntimeClasspath,latestDepJava11TestCompileClasspath,latestDepJava11TestRuntimeClasspath,latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,muzzleTooling,oldH2TestCompileClasspath,oldH2TestRuntimeClasspath,oldPostgresTestCompileClasspath,oldPostgresTestRuntimeClasspath,runtimeClasspath,testCompileClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath
 com.thoughtworks.qdox:qdox:1.12.1=jmhRuntimeClasspath,latestDepJava11TestRuntimeClasspath,latestDepTestRuntimeClasspath,oldH2TestRuntimeClasspath,oldPostgresTestRuntimeClasspath,testRuntimeClasspath
-com.zaxxer:HikariCP:2.4.0=jmhRuntimeClasspath,oldH2TestCompileClasspath,oldH2TestRuntimeClasspath,oldPostgresTestCompileClasspath,oldPostgresTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath
+com.zaxxer:HikariCP:2.4.0=compileClasspath,csiCompileClasspath,jmhCompileClasspath,jmhRuntimeClasspath,oldH2TestCompileClasspath,oldH2TestRuntimeClasspath,oldPostgresTestCompileClasspath,oldPostgresTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath
 com.zaxxer:HikariCP:4.0.3=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath
 com.zaxxer:HikariCP:6.0.0=latestDepJava11TestCompileClasspath,latestDepJava11TestRuntimeClasspath
 commons-codec:commons-codec:1.15=spotbugs

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

@@ -0,0 +1,76 @@
+package datadog.trace.instrumentation.jdbc;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static java.util.Collections.singletonMap;
+
+import com.google.auto.service.AutoService;
+import com.zaxxer.hikari.HikariDataSource;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import java.sql.Connection;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@AutoService(InstrumenterModule.class)
+public final class HikariDataSourceInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType {
+
+  private static final Logger log = LoggerFactory.getLogger(HikariDataSourceInstrumentation.class);
+
+  public HikariDataSourceInstrumentation() {
+    super("jdbc-datasource");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "com.zaxxer.hikari.HikariDataSource";
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("java.sql.Connection", DBInfo.class.getName());
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("getConnection"),
+        HikariDataSourceInstrumentation.class.getName() + "$HikariGetConnectionAdvice");
+  }
+
+  public static class HikariGetConnectionAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void start(
+        @Advice.This final HikariDataSource ds, @Advice.Return Connection con) {
+      if (con == null) {
+        // Exception was probably thrown.
+        return;
+      }
+      // connection pools wraps connection with their own type (ProxyConnection in this case)
+      // jdbc drivers have a standard way to ask for the unwrapped instance (calling unwrap).
+      // we need the unwrapped version in order to be able to lookup the instrumentation context
+      // since we stored dbInfo for that instance and not for the wrapped one
+      Connection unwrapped = con;
+      try {
+        if (con.isWrapperFor(Connection.class)) {
+          unwrapped = con.unwrap(Connection.class);
+        }
+      } catch (Throwable t) {
+        return;
+      }
+      String hikariPoolname = ds.getPoolName();
+      if (unwrapped == null) {
+        return;
+      }
+      DBInfo dbInfo = InstrumentationContext.get(Connection.class, DBInfo.class).get(unwrapped);
+      if (dbInfo == null) {
+        return;
+      }
+      dbInfo.setPoolName(hikariPoolname);
+    }
+  }
+}

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

@@ -1,9 +1,7 @@
 package datadog.trace.instrumentation.jdbc;
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
-import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_OPERATION;
-import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_SCHEMA;
-import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_WAREHOUSE;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.*;
 
 import datadog.trace.api.Config;
 import datadog.trace.api.DDSpanId;
@@ -140,6 +138,7 @@ public AgentSpan onConnection(final AgentSpan span, DBInfo dbInfo) {
 
       setTagIfPresent(span, DB_WAREHOUSE, dbInfo.getWarehouse());
       setTagIfPresent(span, DB_SCHEMA, dbInfo.getSchema());
+      setTagIfPresent(span, DB_POOL_NAME, dbInfo.getPoolName());
     }
     return super.onConnection(span, dbInfo);
   }
@@ -263,6 +262,7 @@ public boolean isSqlServer(final DBInfo dbInfo) {
   public long setContextInfo(Connection connection, DBInfo dbInfo) {
     final byte VERSION = 0;
     final long spanID = Config.get().getIdGenerationStrategy().generateSpanId();
+    // potentially get build span like here
     AgentSpan instrumentationSpan =
         AgentTracer.get().buildSpan("set context_info").withTag("dd.instrumentation", true).start();
     DECORATE.afterStart(instrumentationSpan);