Add automatic instrumentation for JDBC

This instrumentation creates spans for Statements and PreparedStatements. It also captures the corresponding SQL and additional connection info.  ResultSet could be considered for future instrumentation to capture even more of the DB interaction time.

This integration uses Bytebuddy instead of Byteman as the many methods to instrument would have been messy in Byteman.

Author: tylerbenson

README.md

@@ -62,7 +62,7 @@ Finally, add the following JVM argument when starting your application—in your
 -javaagent:/path/to/the/dd-java-agent.jar
 ```
 
-The Java Agent—once passed to your application—automatically traces requests to the frameworks, application servers, and databases shown below. It does this by using various libraries from [opentracing-contrib](https://github.com/opentracing-contrib). In most cases you don't need to install or configure anything; traces will automatically show up in your Datadog dashboards. The exception is [any database library that uses JDBC](#jdbc).
+The Java Agent—once passed to your application—automatically traces requests to the frameworks, application servers, and databases shown below. It does this by using various libraries from [opentracing-contrib](https://github.com/opentracing-contrib). In most cases you don't need to install or configure anything; traces will automatically show up in your Datadog dashboards.
 
 #### Application Servers
 
@@ -87,8 +87,7 @@ Also, frameworks like Spring Boot and Dropwizard inherently work because they us
 
 | Database      | Versions           | Comments  |
 | ------------- |:-------------:| ----- |
-| Spring JDBC| 4.x | **NOT traced automatically**—see [JDBC instructions](#jdbc) |
-| Hibernate | 5.x | **NOT traced automatically**—see [JDBC instructions](#jdbc) |
+| JDBC | 4.x | Intercepts calls to JDBC compatible clients |
 | [MongoDB](https://github.com/opentracing-contrib/java-mongo-driver) | 3.x | Intercepts all the calls from the MongoDB client |
 | [Cassandra](https://github.com/opentracing-contrib/java-cassandra-driver) | 3.2.x | Intercepts all the calls from the Cassandra client |
 
@@ -103,22 +102,6 @@ disabledInstrumentations: ["opentracing-apache-httpclient", "opentracing-mongo-d
 
 See [this YAML file](dd-java-agent/src/main/resources/dd-trace-supported-framework.yaml) for the proper names of all supported libraries (i.e. the names as you must list them in `disabledInstrumentations`).
 
-#### JDBC
-
-The Java Agent doesn't automatically trace requests to databases whose drivers are JDBC-based. For such databases, you must:
-
-1. Add the opentracing-jdbc dependency to your project, e.g. for Maven, add this to pom.xml:
-
-```
-<dependency>
-    <groupId>io.opentracing.contrib</groupId>
-    <artifactId>opentracing-jdbc</artifactId>
-    <version>0.0.3</version>
-</dependency>
-```
-
-2. Modify your code's database connection strings, e.g. for a connection string `jdbc:h2:mem:test`, make it `jdbc:tracing:h2:mem:test`.
-
 ### The `@Trace` Annotation
 
 The Java Agent lets you add a `@Trace` annotation to any method to measure its execution time. Setup the [Java Agent](#java-agent-setup) first if you haven't done so.

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

@@ -20,6 +20,7 @@ dependencies {
   compile project(':dd-trace')
   compile project(':dd-trace-annotations')
 
+  compile group: 'net.bytebuddy', name: 'byte-buddy', version: '1.7.6'
   compile group: 'org.jboss.byteman', name: 'byteman', version: '3.0.10'
 
   compile group: 'org.reflections', name: 'reflections', version: '0.9.11'
@@ -36,9 +37,6 @@ dependencies {
   testCompile(project(path: ':dd-java-agent:integrations:helpers')) {
     transitive = false
   }
-
-  // Not bundled in with the agent.  Usage requires being on the app's classpath (eg. Spring Boot's executable jar)
-  compileOnly group: 'io.opentracing.contrib', name: 'opentracing-jdbc', version: '0.0.3'
 }
 
 project(':dd-java-agent:integrations:helpers').afterEvaluate { helperProject ->
@@ -90,6 +88,7 @@ shadowJar {
     relocate 'com.fasterxml', 'dd.deps.com.fasterxml'
 
     relocate 'javassist', 'dd.deps.javassist'
+    relocate 'net.bytebuddy', 'dd.deps.net.bytebuddy'
     relocate 'org.reflections', 'dd.deps.org.reflections'
     relocate('org.jboss.byteman', 'dd.deps.org.jboss.byteman') {
       // Renaming these causes a verify error in the tests.

dd-java-agent/src/main/java/com/datadoghq/agent/TracingAgent.java

@@ -16,8 +16,16 @@
  */
 package com.datadoghq.agent;
 
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+
+import com.datadoghq.agent.instrumentation.Instrumenter;
 import java.lang.instrument.Instrumentation;
+import java.util.ServiceLoader;
 import lombok.extern.slf4j.Slf4j;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.utility.JavaModule;
 
 /**
  * This class provides a wrapper around the ByteMan agent, to establish required system properties
@@ -27,12 +35,14 @@
 public class TracingAgent {
 
   public static void premain(String agentArgs, final Instrumentation inst) throws Exception {
+    addByteBuddy(inst);
     agentArgs = addManager(agentArgs);
     log.debug("Using premain for loading {}", TracingAgent.class.getSimpleName());
     org.jboss.byteman.agent.Main.premain(agentArgs, inst);
   }
 
   public static void agentmain(String agentArgs, final Instrumentation inst) throws Exception {
+    addByteBuddy(inst);
     agentArgs = addManager(agentArgs);
     log.debug("Using agentmain for loading {}", TracingAgent.class.getSimpleName());
     org.jboss.byteman.agent.Main.agentmain(agentArgs, inst);
@@ -48,4 +58,65 @@ protected static String addManager(String agentArgs) {
     log.debug("Agent args=: {}", agentArgs);
     return agentArgs;
   }
+
+  public static void addByteBuddy(final Instrumentation inst) {
+
+    AgentBuilder agentBuilder =
+        new AgentBuilder.Default()
+            .disableClassFormatChanges()
+            .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
+            .with(new Listener())
+            .ignore(nameStartsWith("com.datadoghq.agent.integration"));
+
+    for (final Instrumenter instrumenter : ServiceLoader.load(Instrumenter.class)) {
+      agentBuilder = instrumenter.instrument(agentBuilder);
+    }
+
+    agentBuilder.installOn(inst);
+  }
+
+  @Slf4j
+  static class Listener implements AgentBuilder.Listener {
+
+    @Override
+    public void onError(
+        final String typeName,
+        final ClassLoader classLoader,
+        final JavaModule module,
+        final boolean loaded,
+        final Throwable throwable) {
+      log.warn("Failed to handle " + typeName + " for transformation", throwable);
+    }
+
+    @Override
+    public void onTransformation(
+        final TypeDescription typeDescription,
+        final ClassLoader classLoader,
+        final JavaModule module,
+        final boolean loaded,
+        final DynamicType dynamicType) {
+      log.debug("Transformed {0}", typeDescription);
+    }
+
+    @Override
+    public void onIgnored(
+        final TypeDescription typeDescription,
+        final ClassLoader classLoader,
+        final JavaModule module,
+        final boolean loaded) {}
+
+    @Override
+    public void onComplete(
+        final String typeName,
+        final ClassLoader classLoader,
+        final JavaModule module,
+        final boolean loaded) {}
+
+    @Override
+    public void onDiscovery(
+        final String typeName,
+        final ClassLoader classLoader,
+        final JavaModule module,
+        final boolean loaded) {}
+  }
 }

dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/Instrumenter.java

@@ -0,0 +1,7 @@
+package com.datadoghq.agent.instrumentation;
+
+import net.bytebuddy.agent.builder.AgentBuilder;
+
+public interface Instrumenter {
+  AgentBuilder instrument(AgentBuilder agentBuilder);
+}

dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/ConnectionInstrumentation.java

@@ -0,0 +1,44 @@
+package com.datadoghq.agent.instrumentation.jdbc;
+
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.datadoghq.agent.instrumentation.Instrumenter;
+import com.google.auto.service.AutoService;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.util.Map;
+import java.util.WeakHashMap;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public final class ConnectionInstrumentation implements Instrumenter {
+  public static final Map<PreparedStatement, String> preparedStatements = new WeakHashMap<>();
+
+  @Override
+  public AgentBuilder instrument(final AgentBuilder agentBuilder) {
+    return agentBuilder
+        .type(not(isInterface()).and(hasSuperType(named(Connection.class.getName()))))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(
+                    nameStartsWith("prepare")
+                        .and(takesArgument(0, String.class))
+                        .and(returns(PreparedStatement.class)),
+                    ConnectionAdvice.class.getName()));
+  }
+
+  public static class ConnectionAdvice {
+    @Advice.OnMethodExit
+    public static void addDBInfo(
+        @Advice.Argument(0) final String sql, @Advice.Return final PreparedStatement statement) {
+      preparedStatements.put(statement, sql);
+    }
+  }
+}

dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/DriverInstrumentation.java

@@ -0,0 +1,55 @@
+package com.datadoghq.agent.instrumentation.jdbc;
+
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.datadoghq.agent.instrumentation.Instrumenter;
+import com.google.auto.service.AutoService;
+import java.sql.Connection;
+import java.sql.Driver;
+import java.util.Map;
+import java.util.Properties;
+import java.util.WeakHashMap;
+import lombok.Data;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public final class DriverInstrumentation implements Instrumenter {
+  public static final Map<Connection, DBInfo> connectionInfo = new WeakHashMap<>();
+
+  @Override
+  public AgentBuilder instrument(final AgentBuilder agentBuilder) {
+    return agentBuilder
+        .type(not(isInterface()).and(hasSuperType(named(Driver.class.getName()))))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(
+                    named("connect").and(takesArguments(String.class, Properties.class)),
+                    DriverAdvice.class.getName()));
+  }
+
+  public static class DriverAdvice {
+    @Advice.OnMethodExit
+    public static void addDBInfo(
+        @Advice.Argument(0) final String url,
+        @Advice.Argument(1) final Properties info,
+        @Advice.Return final Connection connection) {
+      // Remove end of url to prevent passwords from leaking:
+      final String sanitizedURL = url.replaceAll("[?;].*", "");
+      final String type = url.split(":")[1];
+      final String dbUser = info.getProperty("user");
+      connectionInfo.put(connection, new DBInfo(sanitizedURL, type, dbUser));
+    }
+  }
+
+  @Data
+  public static class DBInfo {
+    private final String url;
+    private final String type;
+    private final String user;
+  }
+}

dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/PreparedStatementInstrumentation.java

@@ -0,0 +1,82 @@
+package com.datadoghq.agent.instrumentation.jdbc;
+
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.datadoghq.agent.instrumentation.Instrumenter;
+import com.google.auto.service.AutoService;
+import io.opentracing.ActiveSpan;
+import io.opentracing.NoopActiveSpanSource;
+import io.opentracing.tag.Tags;
+import io.opentracing.util.GlobalTracer;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.util.Collections;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public final class PreparedStatementInstrumentation implements Instrumenter {
+
+  @Override
+  public AgentBuilder instrument(final AgentBuilder agentBuilder) {
+    return agentBuilder
+        .type(not(isInterface()).and(hasSuperType(named(PreparedStatement.class.getName()))))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(
+                    nameStartsWith("execute").and(takesArguments(0)),
+                    PreparedStatementAdvice.class.getName()))
+        .asDecorator();
+  }
+
+  public static class PreparedStatementAdvice {
+
+    @Advice.OnMethodEnter
+    public static ActiveSpan startSpan(@Advice.This final PreparedStatement statement) {
+      // TODO: Should this happen always instead of just inside an existing tracer?
+      if (GlobalTracer.get().activeSpan() == null) {
+        return NoopActiveSpanSource.NoopActiveSpan.INSTANCE;
+      }
+
+      final String sql = ConnectionInstrumentation.preparedStatements.get(statement);
+
+      final ActiveSpan span = GlobalTracer.get().buildSpan("sql.prepared_statement").startActive();
+      Tags.SPAN_KIND.set(span, Tags.SPAN_KIND_CLIENT);
+      Tags.COMPONENT.set(span, "java-jdbc");
+      Tags.DB_STATEMENT.set(span, sql);
+      span.setTag("span.origin.type", statement.getClass().getName());
+
+      try {
+        final Connection connection = statement.getConnection();
+        final DriverInstrumentation.DBInfo dbInfo =
+            DriverInstrumentation.connectionInfo.get(connection);
+
+        span.setTag("db.jdbc.url", dbInfo.getUrl());
+        span.setTag("db.schema", connection.getSchema());
+
+        Tags.DB_TYPE.set(span, dbInfo.getType());
+        if (dbInfo.getUser() != null) {
+          Tags.DB_USER.set(span, dbInfo.getUser());
+        }
+      } finally {
+        return span;
+      }
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void stopSpan(
+        @Advice.Enter final ActiveSpan activeSpan,
+        @Advice.Thrown(readOnly = false) final Throwable throwable) {
+      if (throwable != null) {
+        Tags.ERROR.set(activeSpan, true);
+        activeSpan.log(Collections.singletonMap("error.object", throwable));
+      }
+      activeSpan.deactivate();
+    }
+  }
+}