Fix TLS usage

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<Map<Object, CallDepthThreadLocalMap>> INSTANCES =
+  private static final ThreadLocal<Map<Key, AtomicInteger>> TLS =
-      new ThreadLocal<>();
+      new ThreadLocal<Map<Key, AtomicInteger>>() {
-
+        @Override
-  private static final ThreadLocal<AtomicInteger> tls = new ThreadLocal<>();
+        public Map<Key, AtomicInteger> initialValue() {
-
+          return new HashMap<>();
-  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 CallDepthThreadLocalMap() {}
+  public static int incrementCallDepth(final Key k) {
-
+    final Map<Key, AtomicInteger> map = TLS.get();
-  public int incrementCallDepth() {
+    AtomicInteger depth = map.get(k);
-    AtomicInteger depth = tls.get();
     if (depth == null) {
       depth = new AtomicInteger(0);
-      tls.set(depth);
+      map.put(k, depth);
       return 0;
     } else {
       return depth.incrementAndGet();
     }
   }
 
-  public void reset() {
+  public static void reset(final Key k) {
-    tls.remove();
+    final Map<Key, AtomicInteger> map = TLS.get();
-    INSTANCES.get().remove(this);
+    if (map != null) {
+      map.remove(k);
+    }
+  }
+
+  public enum Key {
+    CLASSLOADER,
+    CONNECTION,
+    PREPARED_STATEMENT,
+    STATEMENT
   }
 }

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
+  }
+}

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");

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:

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 =
+      final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(PREPARED_STATEMENT);
-          CallDepthThreadLocalMap.get(PreparedStatement.class).incrementCallDepth();
       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);
       }
     }
   }

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);
       }
     }
   }