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 00000000000..1c0660e9bfc
--- /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.
+ *
+ *
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> INSTANCES =
+ new ThreadLocal<>();
+
+ public static CallDepthThreadLocalMap get(Object o) {
+ if (INSTANCES.get() == null) {
+ INSTANCES.set(new WeakHashMap());
+ }
+ if (!INSTANCES.get().containsKey(o)) {
+ INSTANCES.get().put(o, new CallDepthThreadLocalMap());
+ }
+ return INSTANCES.get().get(o);
+ }
+
+ private final ThreadLocal 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/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/CallDepthThreadLocalMap.java b/dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/CallDepthThreadLocalMap.java
deleted file mode 100644
index 6eadaaac785..00000000000
--- a/dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/CallDepthThreadLocalMap.java
+++ /dev/null
@@ -1,22 +0,0 @@
-package datadog.trace.instrumentation.classloaders;
-
-import java.util.concurrent.atomic.AtomicInteger;
-
-public class CallDepthThreadLocalMap {
- private static final ThreadLocal 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();
- }
-}
diff --git a/dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/ClassLoaderInstrumentation.java b/dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/ClassLoaderInstrumentation.java
index 5a3237f4bf5..116b69ff43e 100644
--- a/dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/ClassLoaderInstrumentation.java
+++ b/dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/ClassLoaderInstrumentation.java
@@ -9,8 +9,8 @@
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 AgentBuilder apply(final AgentBuilder agentBuilder) {
.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 static class ClassloaderAdvice {
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 static int constructorEnter() {
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");
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 05995b6ccfa..557cafc5b65 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 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 static void addDBInfo(
}
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 static void addDBInfo(@Advice.This final Connection connection) {
}
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
}
}
}