From c1ac920e97512bbf82d2b16b7ebde2c357db1beb Mon Sep 17 00:00:00 2001 From: Tyler Benson Date: Mon, 7 May 2018 10:09:46 +1000 Subject: [PATCH] Change from enum to object for map lookup key --- .../bootstrap/CallDepthThreadLocalMap.java | 21 +++++++------------ .../CallDepthThreadLocalMapTest.groovy | 4 ++-- .../ClassLoaderInstrumentation.java | 5 ++--- .../jdbc/ConnectionInstrumentation.java | 5 ++--- .../PreparedStatementInstrumentation.java | 5 ++--- .../jdbc/StatementInstrumentation.java | 5 ++--- 6 files changed, 17 insertions(+), 28 deletions(-) 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 8f94401e30..4b03d5d4f2 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 @@ -11,16 +11,16 @@ import java.util.concurrent.atomic.AtomicInteger; * #incrementCallDepth at the beginning of each constructor. */ public class CallDepthThreadLocalMap { - private static final ThreadLocal> TLS = - new ThreadLocal>() { + private static final ThreadLocal> TLS = + new ThreadLocal>() { @Override - public Map initialValue() { + public Map initialValue() { return new HashMap<>(); } }; - public static int incrementCallDepth(final Key k) { - final Map map = TLS.get(); + public static int incrementCallDepth(final Object k) { + final Map map = TLS.get(); AtomicInteger depth = map.get(k); if (depth == null) { depth = new AtomicInteger(0); @@ -31,17 +31,10 @@ public class CallDepthThreadLocalMap { } } - public static void reset(final Key k) { - final Map map = TLS.get(); + public static void reset(final Object 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 index 0df2d3219b..d85ff5fd13 100644 --- 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 @@ -6,8 +6,8 @@ class CallDepthThreadLocalMapTest extends Specification { def "test CallDepthThreadLocalMap"() { setup: - def k1 = CallDepthThreadLocalMap.Key.CLASSLOADER - def k2 = CallDepthThreadLocalMap.Key.CONNECTION + def k1 = new Object() + def k2 = new Object() expect: CallDepthThreadLocalMap.incrementCallDepth(k1) == 0 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 b411705902..0e526c2ef4 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,7 +1,6 @@ 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; @@ -46,7 +45,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.incrementCallDepth(CLASSLOADER); + return CallDepthThreadLocalMap.incrementCallDepth(ClassLoader.class); } // Not sure why, but adding suppress causes a verify error. @@ -54,7 +53,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.reset(CLASSLOADER); + CallDepthThreadLocalMap.reset(ClassLoader.class); 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 2d87a5b268..28549d06d3 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,6 +1,5 @@ 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; @@ -59,7 +58,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.incrementCallDepth(CONNECTION); + return CallDepthThreadLocalMap.incrementCallDepth(Connection.class); } // Since we're instrumenting the constructor, we can't add onThrowable. @@ -68,7 +67,7 @@ public final class ConnectionInstrumentation extends Instrumenter.Configurable { @Advice.Enter final int depth, @Advice.This final Connection connection) throws SQLException { if (depth == 0) { - CallDepthThreadLocalMap.reset(CONNECTION); + CallDepthThreadLocalMap.reset(Connection.class); 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 5ff3d2e22c..da1fbde984 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,6 +1,5 @@ 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; @@ -52,7 +51,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.incrementCallDepth(PREPARED_STATEMENT); + final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(PreparedStatement.class); if (callDepth > 0) { return null; } @@ -99,7 +98,7 @@ public final class PreparedStatementInstrumentation extends Instrumenter.Configu span.log(Collections.singletonMap(ERROR_OBJECT, throwable)); } scope.close(); - CallDepthThreadLocalMap.reset(PREPARED_STATEMENT); + CallDepthThreadLocalMap.reset(PreparedStatement.class); } } } 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 d67be7b6ba..4457c76f76 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,6 +1,5 @@ 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; @@ -53,7 +52,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.incrementCallDepth(STATEMENT); + final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Statement.class); if (callDepth > 0) { return null; } @@ -101,7 +100,7 @@ public final class StatementInstrumentation extends Instrumenter.Configurable { span.log(Collections.singletonMap(ERROR_OBJECT, throwable)); } scope.close(); - CallDepthThreadLocalMap.reset(STATEMENT); + CallDepthThreadLocalMap.reset(Statement.class); } } }