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/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
       }
     }
   }