WeakCache instead of 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 @@
 
 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,22 +20,46 @@ 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;
   }
 
   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 @@
 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,18 +1,14 @@
 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;
 
 @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 @@ private static final class SkipClassLoaderMatcher
       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 @@ private static boolean loadsExpectedClass(
   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 @@
 
 /** 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 boolean matches(ClassLoader loader) {
     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.