Fix TLS usage
MODIFY
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/CallDepthThreadLocalMap.java
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/CallDepthThreadLocalMap.java
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<Map<Object, CallDepthThreadLocalMap>> INSTANCES =
- new ThreadLocal<>();
-
- private static final ThreadLocal<AtomicInteger> tls = new ThreadLocal<>();
-
- public static CallDepthThreadLocalMap get(final 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 static final ThreadLocal<Map<Key, AtomicInteger>> TLS =
+ new ThreadLocal<Map<Key, AtomicInteger>>() {
+ @Override
+ public Map<Key, AtomicInteger> initialValue() {
+ return new HashMap<>();
}
+ };
- private CallDepthThreadLocalMap() {}
-
- public int incrementCallDepth() {
- AtomicInteger depth = tls.get();
+ public static int incrementCallDepth(final Key k) {
+ final Map<Key, AtomicInteger> 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<Key, AtomicInteger> map = TLS.get();
+ if (map != null) {
+ map.remove(k);
+ }
+ }
+
+ public enum Key {
+ CLASSLOADER,
+ CONNECTION,
+ PREPARED_STATEMENT,
+ STATEMENT
}
}

ADD
dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/CallDepthThreadLocalMapTest.groovy
+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
+ }
+}

MODIFY
dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/ClassLoaderInstrumentation.java
dd-java-agent/instrumentation/classloaders/src/main/java/datadog/trace/instrumentation/classloaders/ClassLoaderInstrumentation.java
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");

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
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:

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

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

Fix TLS usage