WeakCache instead of WeakMap (#1256)

WeakCache abstraction in addition to WeakMap

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakCache.java

@@ -0,0 +1,19 @@
+package datadog.trace.bootstrap;
+
+import java.util.concurrent.Callable;
+
+public interface WeakCache<K, V> {
+  interface Provider<K, V> {
+    WeakCache<K, V> newWeakCache();
+
+    WeakCache<K, V> newWeakCache(final long maxSize);
+  }
+
+  V getIfPresent(final K key);
+
+  V getIfPresentOrCompute(final K key, final Callable<? extends V> loader);
+
+  V get(final K key, final Callable<? extends V> loader);
+
+  void put(final K key, final V value);
+}

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/WeakMapTest.groovy

@@ -6,11 +6,11 @@ class WeakMapTest extends Specification {
 
   def supplier = new CounterSupplier()
 
-  def sut = new WeakMap.MapAdapter<String, Integer>(new WeakHashMap<>())
+  def weakMap = new WeakMap.MapAdapter<String, Integer>(new WeakHashMap<>())
 
   def "getOrCreate a value"() {
     when:
-    def count = sut.computeIfAbsent('key', supplier)
+    def count = weakMap.computeIfAbsent('key', supplier)
 
     then:
     count == 1
@@ -19,8 +19,8 @@ class WeakMapTest extends Specification {
 
   def "getOrCreate a value multiple times same class loader same key"() {
     when:
-    def count1 = sut.computeIfAbsent('key', supplier)
-    def count2 = sut.computeIfAbsent('key', supplier)
+    def count1 = weakMap.computeIfAbsent('key', supplier)
+    def count2 = weakMap.computeIfAbsent('key', supplier)
 
     then:
     count1 == 1
@@ -30,8 +30,8 @@ class WeakMapTest extends Specification {
 
   def "getOrCreate a value multiple times same class loader different keys"() {
     when:
-    def count1 = sut.computeIfAbsent('key1', supplier)
-    def count2 = sut.computeIfAbsent('key2', supplier)
+    def count1 = weakMap.computeIfAbsent('key1', supplier)
+    def count2 = weakMap.computeIfAbsent('key2', supplier)
 
     then:
     count1 == 1

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentTooling.java

@@ -2,7 +2,11 @@ package datadog.trace.agent.tooling;
 
 import datadog.trace.agent.tooling.bytebuddy.DDCachingPoolStrategy;
 import datadog.trace.agent.tooling.bytebuddy.DDLocationStrategy;
+import datadog.trace.bootstrap.WeakCache;
+import datadog.trace.bootstrap.WeakCache.Provider;
 import datadog.trace.bootstrap.WeakMap;
+import java.util.Iterator;
+import java.util.ServiceLoader;
 
 /**
  * This class contains class references for objects shared by the agent installer as well as muzzle
@@ -16,9 +20,41 @@ public class AgentTooling {
     registerWeakMapProvider();
   }
 
+  private static void registerWeakMapProvider() {
+    if (!WeakMap.Provider.isProviderRegistered()) {
+      WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent(new Cleaner()));
+      //    WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent.Inline());
+      //    WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.Guava());
+    }
+  }
+
+  private static <K, V> Provider loadWeakCacheProvider() {
+    final Iterator<Provider> providers =
+        ServiceLoader.load(Provider.class, AgentInstaller.class.getClassLoader()).iterator();
+    if (providers.hasNext()) {
+      final Provider provider = providers.next();
+      if (providers.hasNext()) {
+        throw new IllegalStateException(
+            "Only one implementation of WeakCache.Provider suppose to be in classpath");
+      }
+      return provider;
+    }
+    throw new IllegalStateException("Can't load implementation of WeakCache.Provider");
+  }
+
+  private static final Provider weakCacheProvider = loadWeakCacheProvider();
+
   private static final DDLocationStrategy LOCATION_STRATEGY = new DDLocationStrategy();
   private static final DDCachingPoolStrategy POOL_STRATEGY = new DDCachingPoolStrategy();
 
+  public static <K, V> WeakCache<K, V> newWeakCache() {
+    return weakCacheProvider.newWeakCache();
+  }
+
+  public static <K, V> WeakCache<K, V> newWeakCache(final long maxSize) {
+    return weakCacheProvider.newWeakCache(maxSize);
+  }
+
   public static DDLocationStrategy locationStrategy() {
     return LOCATION_STRATEGY;
   }
@@ -26,12 +62,4 @@ public class AgentTooling {
   public static DDCachingPoolStrategy poolStrategy() {
     return POOL_STRATEGY;
   }
-
-  private static void registerWeakMapProvider() {
-    if (!WeakMap.Provider.isProviderRegistered()) {
-      WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent(new Cleaner()));
-      //    WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent.Inline());
-      //    WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.Guava());
-    }
-  }
 }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassHierarchyIterable.java

@@ -28,7 +28,7 @@ import java.util.Set;
 public class ClassHierarchyIterable implements Iterable<Class<?>> {
   private final Class<?> baseClass;
 
-  public ClassHierarchyIterable(final Class baseClass) {
+  public ClassHierarchyIterable(final Class<?> baseClass) {
     this.baseClass = baseClass;
   }
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderMatcher.java

@@ -1,8 +1,7 @@
 package datadog.trace.agent.tooling;
 
-import com.google.common.cache.Cache;
-import com.google.common.cache.CacheBuilder;
 import datadog.trace.bootstrap.PatchLogger;
+import datadog.trace.bootstrap.WeakCache;
 import io.opentracing.util.GlobalTracer;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -10,9 +9,6 @@ import net.bytebuddy.matcher.ElementMatcher;
 @Slf4j
 public final class ClassLoaderMatcher {
   public static final ClassLoader BOOTSTRAP_CLASSLOADER = null;
-  public static final int CACHE_MAX_SIZE = 25; // limit number of cached responses for each matcher.
-  public static final int CACHE_CONCURRENCY =
-      Math.max(8, Runtime.getRuntime().availableProcessors());
 
   /** A private constructor that must not be invoked. */
   private ClassLoaderMatcher() {
@@ -39,10 +35,9 @@ public final class ClassLoaderMatcher {
       extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
     public static final SkipClassLoaderMatcher INSTANCE = new SkipClassLoaderMatcher();
     /* Cache of classloader-instance -> (true|false). True = skip instrumentation. False = safe to instrument. */
-    private static final Cache<ClassLoader, Boolean> skipCache =
-        CacheBuilder.newBuilder().weakKeys().concurrencyLevel(CACHE_CONCURRENCY).build();
     private static final String DATADOG_CLASSLOADER_NAME =
         "datadog.trace.bootstrap.DatadogClassLoader";
+    private static final WeakCache<ClassLoader, Boolean> skipCache = AgentTooling.newWeakCache();
 
     private SkipClassLoaderMatcher() {}
 
@@ -115,12 +110,7 @@ public final class ClassLoaderMatcher {
   private static class ClassLoaderHasClassesNamedMatcher
       extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
 
-    private final Cache<ClassLoader, Boolean> cache =
-        CacheBuilder.newBuilder()
-            .weakKeys()
-            .maximumSize(CACHE_MAX_SIZE)
-            .concurrencyLevel(CACHE_CONCURRENCY)
-            .build();
+    private final WeakCache<ClassLoader, Boolean> cache = AgentTooling.newWeakCache(25);
 
     private final String[] resources;
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/GuavaWeakCache.java

@@ -0,0 +1,90 @@
+package datadog.trace.agent.tooling;
+
+import com.google.auto.service.AutoService;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import datadog.trace.bootstrap.WeakCache;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * no null keys nor null values are permitted
+ *
+ * @param <K>
+ * @param <V>
+ */
+@Slf4j
+public final class GuavaWeakCache<K, V> implements WeakCache<K, V> {
+  @AutoService(WeakCache.Provider.class)
+  public static final class Provider<K, V> implements WeakCache.Provider<K, V> {
+    private static final int CACHE_CONCURRENCY =
+        Math.max(8, Runtime.getRuntime().availableProcessors());
+
+    @Override
+    public GuavaWeakCache<K, V> newWeakCache() {
+      return new GuavaWeakCache(
+          CacheBuilder.newBuilder()
+              .weakKeys()
+              .concurrencyLevel(CACHE_CONCURRENCY)
+              .expireAfterAccess(10, TimeUnit.MINUTES)
+              .build());
+    }
+
+    @Override
+    public GuavaWeakCache<K, V> newWeakCache(final long maxSize) {
+      return new GuavaWeakCache(
+          CacheBuilder.newBuilder()
+              .weakKeys()
+              .maximumSize(maxSize)
+              .concurrencyLevel(CACHE_CONCURRENCY)
+              .expireAfterAccess(10, TimeUnit.MINUTES)
+              .build());
+    }
+  }
+
+  private final Cache<K, V> cache;
+
+  private GuavaWeakCache(final Cache<K, V> cache) {
+    this.cache = cache;
+  }
+
+  /**
+   * @return null if key is not present
+   * @param key
+   */
+  @Override
+  public V getIfPresent(final K key) {
+    return cache.getIfPresent(key);
+  }
+
+  @Override
+  public V getIfPresentOrCompute(final K key, final Callable<? extends V> loader) {
+    final V v = cache.getIfPresent(key);
+    if (v != null) {
+      return v;
+    }
+    try {
+      return cache.get(key, loader);
+    } catch (ExecutionException e) {
+      log.error("Can't get value from cache", e);
+    }
+    return null;
+  }
+
+  @Override
+  public V get(final K key, final Callable<? extends V> loader) {
+    try {
+      return cache.get(key, loader);
+    } catch (ExecutionException e) {
+      log.error("Can't get value from cache", e);
+    }
+    return null;
+  }
+
+  @Override
+  public void put(final K key, final V value) {
+    cache.put(key, value);
+  }
+}

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java

@@ -1,19 +1,19 @@
 package datadog.trace.agent.tooling.muzzle;
 
-import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
 import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
 
 import datadog.trace.agent.tooling.AgentTooling;
 import datadog.trace.agent.tooling.Utils;
 import datadog.trace.agent.tooling.muzzle.Reference.Mismatch;
 import datadog.trace.agent.tooling.muzzle.Reference.Source;
-import datadog.trace.bootstrap.WeakMap;
+import datadog.trace.bootstrap.WeakCache;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.Callable;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.description.field.FieldDescription;
 import net.bytebuddy.description.method.MethodDescription;
@@ -22,8 +22,8 @@ import net.bytebuddy.pool.TypePool;
 
 /** Matches a set of references against a classloader. */
 @Slf4j
-public final class ReferenceMatcher implements WeakMap.ValueSupplier<ClassLoader, Boolean> {
-  private final WeakMap<ClassLoader, Boolean> mismatchCache = newWeakMap();
+public final class ReferenceMatcher {
+  private final WeakCache<ClassLoader, Boolean> mismatchCache = AgentTooling.newWeakCache();
   private final Reference[] references;
   private final Set<String> helperClassNames;
 
@@ -50,12 +50,18 @@ public final class ReferenceMatcher implements WeakMap.ValueSupplier<ClassLoader
     if (loader == BOOTSTRAP_LOADER) {
       loader = Utils.getBootstrapProxy();
     }
-
-    return mismatchCache.computeIfAbsent(loader, this);
+    final ClassLoader cl = loader;
+    return mismatchCache.getIfPresentOrCompute(
+        loader,
+        new Callable<Boolean>() {
+          @Override
+          public Boolean call() {
+            return doesMatch(cl);
+          }
+        });
   }
 
-  @Override
-  public Boolean get(final ClassLoader loader) {
+  private boolean doesMatch(final ClassLoader loader) {
     for (final Reference reference : references) {
       // Don't reference-check helper classes.
       // They will be injected by the instrumentation's HelperInjector.

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakCacheTest.groovy

@@ -0,0 +1,79 @@
+package datadog.trace.agent.tooling
+
+import spock.lang.Specification
+
+import java.util.concurrent.Callable
+
+class WeakCacheTest  extends Specification  {
+  def supplier = new CounterSupplier()
+
+  def weakCache = AgentTooling.newWeakCache()
+  def weakCacheFor1elem = AgentTooling.newWeakCache(1)
+
+  def "getOrCreate a value"() {
+    when:
+    def count = weakCache.get('key', supplier)
+
+    then:
+    count == 1
+    supplier.counter == 1
+    weakCache.cache.size() == 1
+  }
+
+  def "getOrCreate a value multiple times same class loader same key"() {
+    when:
+    def count1 = weakCache.get('key', supplier)
+    def count2 = weakCache.get('key', supplier)
+
+    then:
+    count1 == 1
+    count2 == 1
+    supplier.counter == 1
+    weakCache.cache.size() == 1
+  }
+
+  def "getOrCreate a value multiple times same class loader different keys"() {
+    when:
+    def count1 = weakCache.get('key1', supplier)
+    def count2 = weakCache.get('key2', supplier)
+
+    then:
+    count1 == 1
+    count2 == 2
+    supplier.counter == 2
+    weakCache.cache.size() == 2
+  }
+
+  def "max size check"() {
+    when:
+    def sizeBefore =  weakCacheFor1elem.cache.size()
+    def valBefore = weakCacheFor1elem.getIfPresent("key1")
+    def sizeAfter =  weakCacheFor1elem.cache.size()
+    def valAfterGet = weakCacheFor1elem.getIfPresentOrCompute("key1", supplier)
+    def sizeAfterCompute = weakCacheFor1elem.cache.size()
+    weakCacheFor1elem.put("key1", 42)
+    def valAfterPut = weakCacheFor1elem.getIfPresentOrCompute("key1", supplier)
+    def valByKey2 = weakCacheFor1elem.getIfPresentOrCompute("key2", supplier)
+    def valAfterReplace = weakCacheFor1elem.getIfPresent("key1")
+
+    then:
+    valBefore == null
+    valAfterGet == 1
+    sizeBefore == 0
+    sizeAfter == 0
+    sizeAfterCompute == 1
+    valAfterPut == 42
+    valByKey2 == 2
+    valAfterReplace == null
+    weakCacheFor1elem.cache.size() == 1
+  }
+
+  class CounterSupplier implements Callable<Integer> {
+    def counter = 0
+    @Override
+    Integer call() {
+      counter = counter + 1
+      return counter
+    }
+  }
+}

dd-java-agent/dd-java-agent.gradle

@@ -114,8 +114,8 @@ tasks.withType(Test).configureEach {
   jvmArgs "-Ddd.service.name=java-agent-tests"
   jvmArgs "-Ddd.writer.type=LoggingWriter"
   // Multi-threaded logging seems to be causing deadlocks with Gradle's log capture.
-//  jvmArgs "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=debug"
-//  jvmArgs "-Dorg.slf4j.simpleLogger.defaultLogLevel=debug"
+  //  jvmArgs "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=debug"
+  //  jvmArgs "-Dorg.slf4j.simpleLogger.defaultLogLevel=debug"
 
   doFirst {
     // Defining here to allow jacoco to be first on the command line.

dd-java-agent/instrumentation/jax-rs-annotations-1/src/main/java/datadog/trace/instrumentation/jaxrs1/JaxRsAnnotationsDecorator.java

@@ -19,7 +19,7 @@ import javax.ws.rs.Path;
 public class JaxRsAnnotationsDecorator extends BaseDecorator {
   public static JaxRsAnnotationsDecorator DECORATE = new JaxRsAnnotationsDecorator();
 
-  private final WeakMap<Class, Map<Method, String>> resourceNames = newWeakMap();
+  private final WeakMap<Class<?>, Map<Method, String>> resourceNames = newWeakMap();
 
   @Override
   protected String[] instrumentationNames() {
@@ -37,7 +37,7 @@ public class JaxRsAnnotationsDecorator extends BaseDecorator {
   }
 
   public void onControllerStart(
-      final AgentSpan span, final AgentSpan parent, final Class target, final Method method) {
+      final AgentSpan span, final AgentSpan parent, final Class<?> target, final Method method) {
     final String resourceName = getPathResourceName(target, method);
     updateParent(parent, resourceName);
 
@@ -70,9 +70,8 @@ public class JaxRsAnnotationsDecorator extends BaseDecorator {
    *
    * @return The result can be an empty string but will never be {@code null}.
    */
-  private String getPathResourceName(final Class target, final Method method) {
+  private String getPathResourceName(final Class<?> target, final Method method) {
     Map<Method, String> classMap = resourceNames.get(target);
-
     if (classMap == null) {
       resourceNames.putIfAbsent(target, new ConcurrentHashMap<Method, String>());
       classMap = resourceNames.get(target);
@@ -127,7 +126,7 @@ public class JaxRsAnnotationsDecorator extends BaseDecorator {
     return method.getAnnotation(Path.class);
   }
 
-  private Path findClassPath(final Class<Object> target) {
+  private Path findClassPath(final Class<?> target) {
     for (final Class<?> currentClass : new ClassHierarchyIterable(target)) {
       final Path annotation = currentClass.getAnnotation(Path.class);
       if (annotation != null) {

dd-java-agent/instrumentation/jax-rs-annotations-1/src/test/groovy/JaxRsAnnotations1InstrumentationTest.groovy

@@ -14,7 +14,7 @@ import java.lang.reflect.Method
 
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 
-class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
+class JaxRsAnnotations1InstrumentationTest extends AgentTestRunner {
 
   def "instrumentation can be used as root span and resource is set to METHOD PATH"() {
     setup:
@@ -136,7 +136,7 @@ class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
     // TODO: uncomment when we drop support for Java 7
     // "GET /child/invoke"  | new JavaInterfaces.DefaultChildClassOnInterface()
 
-    className = getName(obj.class)
+    className = getClassName(obj.class)
 
     // JavaInterfaces classes are loaded on a different classloader, so we need to find the right cache instance.
     decorator = obj.class.classLoader.loadClass(JaxRsAnnotationsDecorator.name).getField("DECORATE").get(null)
@@ -193,18 +193,4 @@ class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
     void call() {
     }
   }
-
-  def getName(Class clazz) {
-    String className = clazz.getSimpleName()
-    if (className.isEmpty()) {
-      className = clazz.getName()
-      if (clazz.getPackage() != null) {
-        final String pkgName = clazz.getPackage().getName()
-        if (!pkgName.isEmpty()) {
-          className = clazz.getName().replace(pkgName, "").substring(1)
-        }
-      }
-    }
-    return className
-  }
 }

dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsDecorator.java

@@ -1,7 +1,5 @@
 package datadog.trace.instrumentation.jaxrs2;
 
-import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
-
 import datadog.trace.agent.tooling.ClassHierarchyIterable;
 import datadog.trace.api.DDSpanTypes;
 import datadog.trace.api.DDTags;
@@ -27,7 +25,8 @@ public class JaxRsAnnotationsDecorator extends BaseDecorator {
 
   public static final JaxRsAnnotationsDecorator DECORATE = new JaxRsAnnotationsDecorator();
 
-  private final WeakMap<Class, Map<Method, String>> resourceNames = newWeakMap();
+  private final WeakMap<Class<?>, Map<Method, String>> resourceNames =
+      WeakMap.Provider.newWeakMap();
 
   @Override
   protected String[] instrumentationNames() {
@@ -45,7 +44,7 @@ public class JaxRsAnnotationsDecorator extends BaseDecorator {
   }
 
   public void onJaxRsSpan(
-      final AgentSpan span, final AgentSpan parent, final Class target, final Method method) {
+      final AgentSpan span, final AgentSpan parent, final Class<?> target, final Method method) {
 
     final String resourceName = getPathResourceName(target, method);
     updateParent(parent, resourceName);
@@ -81,7 +80,7 @@ public class JaxRsAnnotationsDecorator extends BaseDecorator {
    *
    * @return The result can be an empty string but will never be {@code null}.
    */
-  private String getPathResourceName(final Class target, final Method method) {
+  private String getPathResourceName(final Class<?> target, final Method method) {
     Map<Method, String> classMap = resourceNames.get(target);
 
     if (classMap == null) {
@@ -138,7 +137,7 @@ public class JaxRsAnnotationsDecorator extends BaseDecorator {
     return method.getAnnotation(Path.class);
   }
 
-  private Path findClassPath(final Class<Object> target) {
+  private Path findClassPath(final Class<?> target) {
     for (final Class<?> currentClass : new ClassHierarchyIterable(target)) {
       final Path annotation = currentClass.getAnnotation(Path.class);
       if (annotation != null) {

dd-java-agent/instrumentation/jax-rs-annotations-2/src/test/groovy/JaxRsAnnotations2InstrumentationTest.groovy

@@ -14,7 +14,7 @@ import java.lang.reflect.Method
 
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 
-class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
+class JaxRsAnnotations2InstrumentationTest extends AgentTestRunner {
 
   def "instrumentation can be used as root span and resource is set to METHOD PATH"() {
     setup:
@@ -136,7 +136,7 @@ class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
     // TODO: uncomment when we drop support for Java 7
 //    "GET /child/invoke"         | new JavaInterfaces.DefaultChildClassOnInterface()
 
-    className = getName(obj.class)
+    className = getClassName(obj.class)
 
     // JavaInterfaces classes are loaded on a different classloader, so we need to find the right cache instance.
     decorator = obj.class.classLoader.loadClass(JaxRsAnnotationsDecorator.name).getField("DECORATE").get(null)
@@ -193,18 +193,4 @@ class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
     void call() {
     }
   }
-
-  def getName(Class clazz) {
-    String className = clazz.getSimpleName()
-    if (className.isEmpty()) {
-      className = clazz.getName()
-      if (clazz.getPackage() != null) {
-        final String pkgName = clazz.getPackage().getName()
-        if (!pkgName.isEmpty()) {
-          className = clazz.getName().replace(pkgName, "").substring(1)
-        }
-      }
-    }
-    return className
-  }
 }

dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java

@@ -6,19 +6,17 @@ import datadog.trace.context.TraceScope;
 import datadog.trace.instrumentation.netty40.client.HttpClientTracingHandler;
 import datadog.trace.instrumentation.netty40.server.HttpServerTracingHandler;
 import io.netty.util.AttributeKey;
-import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 public class AttributeKeys {
-
-  private static final WeakMap<ClassLoader, Map<String, AttributeKey<?>>> map =
+  private static final WeakMap<ClassLoader, ConcurrentMap<String, AttributeKey<?>>> map =
       WeakMap.Implementation.DEFAULT.get();
-
-  private static final WeakMap.ValueSupplier<ClassLoader, Map<String, AttributeKey<?>>>
+  private static final WeakMap.ValueSupplier<ClassLoader, ConcurrentMap<String, AttributeKey<?>>>
       mapSupplier =
-          new WeakMap.ValueSupplier<ClassLoader, Map<String, AttributeKey<?>>>() {
+          new WeakMap.ValueSupplier<ClassLoader, ConcurrentMap<String, AttributeKey<?>>>() {
             @Override
-            public Map<String, AttributeKey<?>> get(final ClassLoader ignored) {
+            public ConcurrentMap<String, AttributeKey<?>> get(final ClassLoader ignore) {
               return new ConcurrentHashMap<>();
             }
           };
@@ -44,7 +42,7 @@ public class AttributeKeys {
    * cassandra driver.
    */
   private static <T> AttributeKey<T> attributeKey(final String key) {
-    final Map<String, AttributeKey<?>> classLoaderMap =
+    final ConcurrentMap<String, AttributeKey<?>> classLoaderMap =
         map.computeIfAbsent(AttributeKey.class.getClassLoader(), mapSupplier);
     if (classLoaderMap.containsKey(key)) {
       return (AttributeKey<T>) classLoaderMap.get(key);

dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java

@@ -6,19 +6,17 @@ import datadog.trace.context.TraceScope;
 import datadog.trace.instrumentation.netty41.client.HttpClientTracingHandler;
 import datadog.trace.instrumentation.netty41.server.HttpServerTracingHandler;
 import io.netty.util.AttributeKey;
-import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 public class AttributeKeys {
-
-  private static final WeakMap<ClassLoader, Map<String, AttributeKey<?>>> map =
+  private static final WeakMap<ClassLoader, ConcurrentMap<String, AttributeKey<?>>> map =
       WeakMap.Implementation.DEFAULT.get();
-
-  private static final WeakMap.ValueSupplier<ClassLoader, Map<String, AttributeKey<?>>>
+  private static final WeakMap.ValueSupplier<ClassLoader, ConcurrentMap<String, AttributeKey<?>>>
       mapSupplier =
-          new WeakMap.ValueSupplier<ClassLoader, Map<String, AttributeKey<?>>>() {
+          new WeakMap.ValueSupplier<ClassLoader, ConcurrentMap<String, AttributeKey<?>>>() {
             @Override
-            public Map<String, AttributeKey<?>> get(final ClassLoader ignored) {
+            public ConcurrentMap<String, AttributeKey<?>> get(final ClassLoader ignore) {
               return new ConcurrentHashMap<>();
             }
           };
@@ -48,7 +46,7 @@ public class AttributeKeys {
    * cassandra driver.
    */
   private static <T> AttributeKey<T> attributeKey(final String key) {
-    final Map<String, AttributeKey<?>> classLoaderMap =
+    final ConcurrentMap<String, AttributeKey<?>> classLoaderMap =
         map.computeIfAbsent(AttributeKey.class.getClassLoader(), mapSupplier);
     if (classLoaderMap.containsKey(key)) {
       return (AttributeKey<T>) classLoaderMap.get(key);

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.java

@@ -307,4 +307,18 @@ public abstract class AgentTestRunner extends DDSpecification {
       }
     }
   }
+
+  protected static String getClassName(Class clazz) {
+    String className = clazz.getSimpleName();
+    if (className.isEmpty()) {
+      className = clazz.getName();
+      if (clazz.getPackage() != null) {
+        final String pkgName = clazz.getPackage().getName();
+        if (!pkgName.isEmpty()) {
+          className = clazz.getName().replace(pkgName, "").substring(1);
+        }
+      }
+    }
+    return className;
+  }
 }