Merge pull request #281 from DataDog/ark/jdbc_constructors

Use CallDepthThreadLocalMap in sql Connection init

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/CallDepthThreadLocalMap.java

@@ -0,0 +1,46 @@
+package datadog.trace.bootstrap;
+
+import java.util.Map;
+import java.util.WeakHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Utility to track nested instrumentation.
+ *
+ * <p>For example, this can be used to track nested calls to super() in constructors by calling
+ * #incrementCallDepth at the beginning of each constructor.
+ */
+public class CallDepthThreadLocalMap {
+  private static final ThreadLocal<Map<Object, CallDepthThreadLocalMap>> INSTANCES =
+      new ThreadLocal<>();
+
+  public static CallDepthThreadLocalMap get(Object o) {
+    if (INSTANCES.get() == null) {
+      INSTANCES.set(new WeakHashMap<Object, CallDepthThreadLocalMap>());
+    }
+    if (!INSTANCES.get().containsKey(o)) {
+      INSTANCES.get().put(o, new CallDepthThreadLocalMap());
+    }
+    return INSTANCES.get().get(o);
+  }
+
+  private final ThreadLocal<AtomicInteger> tls = new ThreadLocal<>();
+
+  private CallDepthThreadLocalMap() {}
+
+  public int incrementCallDepth() {
+    AtomicInteger depth = tls.get();
+    if (depth == null) {
+      depth = new AtomicInteger(0);
+      tls.set(depth);
+      return 0;
+    } else {
+      return depth.incrementAndGet();
+    }
+  }
+
+  public void reset() {
+    tls.remove();
+    INSTANCES.get().remove(this);
+  }
+}

dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/CallDepthThreadLocalMap.java

@@ -1,22 +0,0 @@
-package datadog.trace.instrumentation.classloaders;
-
-import java.util.concurrent.atomic.AtomicInteger;
-
-public class CallDepthThreadLocalMap {
-  private static final ThreadLocal<AtomicInteger> tls = new ThreadLocal<>();
-
-  public static int incrementCallDepth() {
-    AtomicInteger depth = tls.get();
-    if (depth == null) {
-      depth = new AtomicInteger(0);
-      tls.set(depth);
-      return 0;
-    } else {
-      return depth.incrementAndGet();
-    }
-  }
-
-  public static void reset() {
-    tls.remove();
-  }
-}

dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/ClassLoaderInstrumentation.java

@@ -9,8 +9,8 @@ import com.google.auto.service.AutoService;
 import datadog.opentracing.DDTracer;
 import datadog.trace.agent.tooling.DDAdvice;
 import datadog.trace.agent.tooling.DDTransformers;
-import datadog.trace.agent.tooling.HelperInjector;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import io.opentracing.util.GlobalTracer;
 import java.lang.reflect.Field;
 import net.bytebuddy.agent.builder.AgentBuilder;
@@ -34,9 +34,6 @@ public final class ClassLoaderInstrumentation extends Instrumenter.Configurable
         .type(
             failSafe(isSubTypeOf(ClassLoader.class)),
             classLoaderHasClasses("io.opentracing.util.GlobalTracer"))
-        .transform(
-            new HelperInjector(
-                "datadog.trace.instrumentation.classloaders.CallDepthThreadLocalMap"))
         .transform(DDTransformers.defaultTransformers())
         .transform(DDAdvice.create().advice(isConstructor(), ClassloaderAdvice.class.getName()))
         .asDecorator();
@@ -48,7 +45,7 @@ public final class ClassLoaderInstrumentation extends Instrumenter.Configurable
     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();
+      return CallDepthThreadLocalMap.get(ClassLoader.class).incrementCallDepth();
     }
 
     // Not sure why, but adding suppress causes a verify error.
@@ -56,7 +53,7 @@ public final class ClassLoaderInstrumentation extends Instrumenter.Configurable
     public static void constructorExit(
         @Advice.This final ClassLoader cl, @Advice.Enter final int depth) {
       if (depth == 0) {
-        CallDepthThreadLocalMap.reset();
+        CallDepthThreadLocalMap.get(ClassLoader.class).reset();
 
         try {
           final Field field = GlobalTracer.class.getDeclaredField("tracer");

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

@@ -13,9 +13,11 @@ 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;
 
@@ -52,9 +54,19 @@ public final class ConnectionInstrumentation extends Instrumenter.Configurable {
   }
 
   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.get(Connection.class).incrementCallDepth();
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void addDBInfo(@Advice.This final Connection connection) {
-      try {
+    public static void constructorExit(
+        @Advice.Enter final int depth, @Advice.This final Connection connection)
+        throws SQLException {
+      if (depth == 0) {
+        CallDepthThreadLocalMap.get(Connection.class).reset();
         final String url = connection.getMetaData().getURL();
         if (url != null) {
           // Remove end of url to prevent passwords from leaking:
@@ -66,9 +78,6 @@ public final class ConnectionInstrumentation extends Instrumenter.Configurable {
           }
           JDBCMaps.connectionInfo.put(connection, new JDBCMaps.DBInfo(sanitizedURL, type, user));
         }
-      } catch (final Throwable t) {
-        // object may not be fully initialized.
-        // calling constructor will populate map
       }
     }
   }