Use CallDepthThreadLocalMap in sql Connection init
ADD
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/CallDepthThreadLocalMap.java
+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);
+ }
+}

MODIFY
dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/ConnectionInstrumentation.java
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
}
}
}

Use CallDepthThreadLocalMap in sql Connection init