Change from enum to object for map lookup key

This commit is contained in:
Tyler Benson 2018-05-07 10:09:46 +10:00
parent 122c482dd5
commit c1ac920e97
6 changed files with 17 additions and 28 deletions

View File

@ -11,16 +11,16 @@ import java.util.concurrent.atomic.AtomicInteger;
* #incrementCallDepth at the beginning of each constructor. * #incrementCallDepth at the beginning of each constructor.
*/ */
public class CallDepthThreadLocalMap { public class CallDepthThreadLocalMap {
private static final ThreadLocal<Map<Key, AtomicInteger>> TLS = private static final ThreadLocal<Map<Object, AtomicInteger>> TLS =
new ThreadLocal<Map<Key, AtomicInteger>>() { new ThreadLocal<Map<Object, AtomicInteger>>() {
@Override @Override
public Map<Key, AtomicInteger> initialValue() { public Map<Object, AtomicInteger> initialValue() {
return new HashMap<>(); return new HashMap<>();
} }
}; };
public static int incrementCallDepth(final Key k) { public static int incrementCallDepth(final Object k) {
final Map<Key, AtomicInteger> map = TLS.get(); final Map<Object, AtomicInteger> map = TLS.get();
AtomicInteger depth = map.get(k); AtomicInteger depth = map.get(k);
if (depth == null) { if (depth == null) {
depth = new AtomicInteger(0); depth = new AtomicInteger(0);
@ -31,17 +31,10 @@ public class CallDepthThreadLocalMap {
} }
} }
public static void reset(final Key k) { public static void reset(final Object k) {
final Map<Key, AtomicInteger> map = TLS.get(); final Map<Object, AtomicInteger> map = TLS.get();
if (map != null) { if (map != null) {
map.remove(k); map.remove(k);
} }
} }
public enum Key {
CLASSLOADER,
CONNECTION,
PREPARED_STATEMENT,
STATEMENT
}
} }

View File

@ -6,8 +6,8 @@ class CallDepthThreadLocalMapTest extends Specification {
def "test CallDepthThreadLocalMap"() { def "test CallDepthThreadLocalMap"() {
setup: setup:
def k1 = CallDepthThreadLocalMap.Key.CLASSLOADER def k1 = new Object()
def k2 = CallDepthThreadLocalMap.Key.CONNECTION def k2 = new Object()
expect: expect:
CallDepthThreadLocalMap.incrementCallDepth(k1) == 0 CallDepthThreadLocalMap.incrementCallDepth(k1) == 0

View File

@ -1,7 +1,6 @@
package datadog.trace.instrumentation.classloaders; package datadog.trace.instrumentation.classloaders;
import static datadog.trace.agent.tooling.ClassLoaderMatcher.classLoaderHasClasses; 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.failSafe;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor; import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf; import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
@ -46,7 +45,7 @@ public final class ClassLoaderInstrumentation extends Instrumenter.Configurable
public static int constructorEnter() { public static int constructorEnter() {
// We use this to make sure we only apply the exit instrumentation // We use this to make sure we only apply the exit instrumentation
// after the constructors are done calling their super constructors. // 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. // 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( public static void constructorExit(
@Advice.This final ClassLoader cl, @Advice.Enter final int depth) { @Advice.This final ClassLoader cl, @Advice.Enter final int depth) {
if (depth == 0) { if (depth == 0) {
CallDepthThreadLocalMap.reset(CLASSLOADER); CallDepthThreadLocalMap.reset(ClassLoader.class);
try { try {
final Field field = GlobalTracer.class.getDeclaredField("tracer"); final Field field = GlobalTracer.class.getDeclaredField("tracer");

View File

@ -1,6 +1,5 @@
package datadog.trace.instrumentation.jdbc; 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.failSafe;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor; import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isInterface; import static net.bytebuddy.matcher.ElementMatchers.isInterface;
@ -59,7 +58,7 @@ public final class ConnectionInstrumentation extends Instrumenter.Configurable {
public static int constructorEnter() { public static int constructorEnter() {
// We use this to make sure we only apply the exit instrumentation // We use this to make sure we only apply the exit instrumentation
// after the constructors are done calling their super constructors. // 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. // 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) @Advice.Enter final int depth, @Advice.This final Connection connection)
throws SQLException { throws SQLException {
if (depth == 0) { if (depth == 0) {
CallDepthThreadLocalMap.reset(CONNECTION); CallDepthThreadLocalMap.reset(Connection.class);
final String url = connection.getMetaData().getURL(); final String url = connection.getMetaData().getURL();
if (url != null) { if (url != null) {
// Remove end of url to prevent passwords from leaking: // Remove end of url to prevent passwords from leaking:

View File

@ -1,6 +1,5 @@
package datadog.trace.instrumentation.jdbc; package datadog.trace.instrumentation.jdbc;
import static datadog.trace.bootstrap.CallDepthThreadLocalMap.Key.PREPARED_STATEMENT;
import static io.opentracing.log.Fields.ERROR_OBJECT; import static io.opentracing.log.Fields.ERROR_OBJECT;
import static net.bytebuddy.matcher.ElementMatchers.failSafe; import static net.bytebuddy.matcher.ElementMatchers.failSafe;
import static net.bytebuddy.matcher.ElementMatchers.isInterface; import static net.bytebuddy.matcher.ElementMatchers.isInterface;
@ -52,7 +51,7 @@ public final class PreparedStatementInstrumentation extends Instrumenter.Configu
@Advice.OnMethodEnter(suppress = Throwable.class) @Advice.OnMethodEnter(suppress = Throwable.class)
public static Scope startSpan(@Advice.This final PreparedStatement statement) { 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) { if (callDepth > 0) {
return null; return null;
} }
@ -99,7 +98,7 @@ public final class PreparedStatementInstrumentation extends Instrumenter.Configu
span.log(Collections.singletonMap(ERROR_OBJECT, throwable)); span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
} }
scope.close(); scope.close();
CallDepthThreadLocalMap.reset(PREPARED_STATEMENT); CallDepthThreadLocalMap.reset(PreparedStatement.class);
} }
} }
} }

View File

@ -1,6 +1,5 @@
package datadog.trace.instrumentation.jdbc; package datadog.trace.instrumentation.jdbc;
import static datadog.trace.bootstrap.CallDepthThreadLocalMap.Key.STATEMENT;
import static io.opentracing.log.Fields.ERROR_OBJECT; import static io.opentracing.log.Fields.ERROR_OBJECT;
import static net.bytebuddy.matcher.ElementMatchers.failSafe; import static net.bytebuddy.matcher.ElementMatchers.failSafe;
import static net.bytebuddy.matcher.ElementMatchers.isInterface; import static net.bytebuddy.matcher.ElementMatchers.isInterface;
@ -53,7 +52,7 @@ public final class StatementInstrumentation extends Instrumenter.Configurable {
@Advice.OnMethodEnter(suppress = Throwable.class) @Advice.OnMethodEnter(suppress = Throwable.class)
public static Scope startSpan( public static Scope startSpan(
@Advice.Argument(0) final String sql, @Advice.This final Statement statement) { @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) { if (callDepth > 0) {
return null; return null;
} }
@ -101,7 +100,7 @@ public final class StatementInstrumentation extends Instrumenter.Configurable {
span.log(Collections.singletonMap(ERROR_OBJECT, throwable)); span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
} }
scope.close(); scope.close();
CallDepthThreadLocalMap.reset(STATEMENT); CallDepthThreadLocalMap.reset(Statement.class);
} }
} }
} }