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 index 58a363f180..8f94401e30 100644 --- 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 @@ -1,7 +1,7 @@ package datadog.trace.bootstrap; +import java.util.HashMap; import java.util.Map; -import java.util.WeakHashMap; import java.util.concurrent.atomic.AtomicInteger; /** @@ -11,36 +11,37 @@ import java.util.concurrent.atomic.AtomicInteger; * #incrementCallDepth at the beginning of each constructor. */ public class CallDepthThreadLocalMap { - private static final ThreadLocal> INSTANCES = - new ThreadLocal<>(); + private static final ThreadLocal> TLS = + new ThreadLocal>() { + @Override + public Map initialValue() { + return new HashMap<>(); + } + }; - private static final ThreadLocal tls = new ThreadLocal<>(); - - public static CallDepthThreadLocalMap get(final 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 CallDepthThreadLocalMap() {} - - public int incrementCallDepth() { - AtomicInteger depth = tls.get(); + public static int incrementCallDepth(final Key k) { + final Map map = TLS.get(); + AtomicInteger depth = map.get(k); if (depth == null) { depth = new AtomicInteger(0); - tls.set(depth); + map.put(k, depth); return 0; } else { return depth.incrementAndGet(); } } - public void reset() { - tls.remove(); - INSTANCES.get().remove(this); + public static void reset(final Key k) { + final Map map = TLS.get(); + if (map != null) { + map.remove(k); + } + } + + public enum Key { + CLASSLOADER, + CONNECTION, + PREPARED_STATEMENT, + STATEMENT } } diff --git a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/CallDepthThreadLocalMapTest.groovy b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/CallDepthThreadLocalMapTest.groovy new file mode 100644 index 0000000000..0df2d3219b --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/CallDepthThreadLocalMapTest.groovy @@ -0,0 +1,35 @@ +package datadog.trace.bootstrap + +import spock.lang.Specification + +class CallDepthThreadLocalMapTest extends Specification { + + def "test CallDepthThreadLocalMap"() { + setup: + def k1 = CallDepthThreadLocalMap.Key.CLASSLOADER + def k2 = CallDepthThreadLocalMap.Key.CONNECTION + + expect: + CallDepthThreadLocalMap.incrementCallDepth(k1) == 0 + CallDepthThreadLocalMap.incrementCallDepth(k2) == 0 + + CallDepthThreadLocalMap.incrementCallDepth(k1) == 1 + CallDepthThreadLocalMap.incrementCallDepth(k2) == 1 + + when: + CallDepthThreadLocalMap.reset(k1) + + then: + CallDepthThreadLocalMap.incrementCallDepth(k2) == 2 + + when: + CallDepthThreadLocalMap.reset(k2) + + then: + CallDepthThreadLocalMap.incrementCallDepth(k1) == 0 + CallDepthThreadLocalMap.incrementCallDepth(k2) == 0 + + CallDepthThreadLocalMap.incrementCallDepth(k1) == 1 + CallDepthThreadLocalMap.incrementCallDepth(k2) == 1 + } +} 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 116b69ff43..b411705902 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 @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.classloaders; import static datadog.trace.agent.tooling.ClassLoaderMatcher.classLoaderHasClasses; +import static datadog.trace.bootstrap.CallDepthThreadLocalMap.Key.CLASSLOADER; import static net.bytebuddy.matcher.ElementMatchers.failSafe; import static net.bytebuddy.matcher.ElementMatchers.isConstructor; import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf; @@ -45,7 +46,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.get(ClassLoader.class).incrementCallDepth(); + return CallDepthThreadLocalMap.incrementCallDepth(CLASSLOADER); } // Not sure why, but adding suppress causes a verify error. @@ -53,7 +54,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.get(ClassLoader.class).reset(); + CallDepthThreadLocalMap.reset(CLASSLOADER); 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 4f9e25e32e..2d87a5b268 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 @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.jdbc; +import static datadog.trace.bootstrap.CallDepthThreadLocalMap.Key.CONNECTION; import static net.bytebuddy.matcher.ElementMatchers.failSafe; import static net.bytebuddy.matcher.ElementMatchers.isConstructor; import static net.bytebuddy.matcher.ElementMatchers.isInterface; @@ -58,7 +59,7 @@ public final class ConnectionInstrumentation 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.get(Connection.class).incrementCallDepth(); + return CallDepthThreadLocalMap.incrementCallDepth(CONNECTION); } // Since we're instrumenting the constructor, we can't add onThrowable. @@ -67,7 +68,7 @@ public final class ConnectionInstrumentation extends Instrumenter.Configurable { @Advice.Enter final int depth, @Advice.This final Connection connection) throws SQLException { if (depth == 0) { - CallDepthThreadLocalMap.get(Connection.class).reset(); + CallDepthThreadLocalMap.reset(CONNECTION); final String url = connection.getMetaData().getURL(); if (url != null) { // Remove end of url to prevent passwords from leaking: diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java index bc1f3b49b6..5ff3d2e22c 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.jdbc; +import static datadog.trace.bootstrap.CallDepthThreadLocalMap.Key.PREPARED_STATEMENT; import static io.opentracing.log.Fields.ERROR_OBJECT; import static net.bytebuddy.matcher.ElementMatchers.failSafe; import static net.bytebuddy.matcher.ElementMatchers.isInterface; @@ -51,8 +52,7 @@ public final class PreparedStatementInstrumentation extends Instrumenter.Configu @Advice.OnMethodEnter(suppress = Throwable.class) public static Scope startSpan(@Advice.This final PreparedStatement statement) { - final int callDepth = - CallDepthThreadLocalMap.get(PreparedStatement.class).incrementCallDepth(); + final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(PREPARED_STATEMENT); if (callDepth > 0) { return null; } @@ -99,7 +99,7 @@ public final class PreparedStatementInstrumentation extends Instrumenter.Configu span.log(Collections.singletonMap(ERROR_OBJECT, throwable)); } scope.close(); - CallDepthThreadLocalMap.get(PreparedStatement.class).reset(); + CallDepthThreadLocalMap.reset(PREPARED_STATEMENT); } } } diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java index 9013042463..d67be7b6ba 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.jdbc; +import static datadog.trace.bootstrap.CallDepthThreadLocalMap.Key.STATEMENT; import static io.opentracing.log.Fields.ERROR_OBJECT; import static net.bytebuddy.matcher.ElementMatchers.failSafe; import static net.bytebuddy.matcher.ElementMatchers.isInterface; @@ -52,7 +53,7 @@ public final class StatementInstrumentation extends Instrumenter.Configurable { @Advice.OnMethodEnter(suppress = Throwable.class) public static Scope startSpan( @Advice.Argument(0) final String sql, @Advice.This final Statement statement) { - final int callDepth = CallDepthThreadLocalMap.get(Statement.class).incrementCallDepth(); + final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(STATEMENT); if (callDepth > 0) { return null; } @@ -100,7 +101,7 @@ public final class StatementInstrumentation extends Instrumenter.Configurable { span.log(Collections.singletonMap(ERROR_OBJECT, throwable)); } scope.close(); - CallDepthThreadLocalMap.get(Statement.class).reset(); + CallDepthThreadLocalMap.reset(STATEMENT); } } }