From 31c7d2d328fe1e2c9757285392c1d07ff8050aa1 Mon Sep 17 00:00:00 2001
From: Andrew Kent <andrew.kent@datadoghq.com>
Date: Wed, 4 Apr 2018 15:34:43 -0400
Subject: [PATCH] Use CallDepthThreadLocalMap in sql Connection init

---
 .../bootstrap/CallDepthThreadLocalMap.java    | 46 +++++++++++++++++++
 .../jdbc/ConnectionInstrumentation.java       | 19 ++++++--
 2 files changed, 60 insertions(+), 5 deletions(-)
 create mode 100644 dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/CallDepthThreadLocalMap.java

diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/CallDepthThreadLocalMap.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/CallDepthThreadLocalMap.java
new file mode 100644
index 0000000000..1c0660e9bf
--- /dev/null
+++ b/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);
+  }
+}
diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/ConnectionInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/ConnectionInstrumentation.java
index 05995b6ccf..557cafc5b6 100644
--- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/ConnectionInstrumentation.java
+++ b/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
       }
     }
   }