DataDog · randomanderson · Nov 24, 2020 · Oct 29, 2020 · Oct 29, 2020 · Oct 29, 2020
@@ -1,19 +1,15 @@
 package datadog.trace.bootstrap;
 
-import java.util.concurrent.Callable;
+import datadog.trace.api.Function;
 
 public interface WeakCache<K, V> {
-  interface Provider<K, V> {
-    WeakCache<K, V> newWeakCache();
-
-    WeakCache<K, V> newWeakCache(final long maxSize);
+  interface Provider {
+    <K, V> WeakCache<K, V> newWeakCache(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);
+  V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction);
 
   void put(final K key, final V value);
 }
@@ -1,5 +1,6 @@
 package datadog.trace.bootstrap;
 
+import datadog.trace.api.Function;
 import java.util.Collections;
 import java.util.Map;
 import java.util.WeakHashMap;
@@ -18,7 +19,7 @@ public interface WeakMap<K, V> {
 
   void putIfAbsent(K key, V value);
 
-  V computeIfAbsent(K key, ValueSupplier<? super K, ? extends V> supplier);
+  V computeIfAbsent(K key, Function<? super K, ? extends V> supplier);
 
   @Slf4j
   class Provider {
@@ -58,14 +59,6 @@ public <K, V> WeakMap<K, V> get() {
     }
   }
 
-  /**
-   * Supplies the value to be stored and it is called only when a value does not exists yet in the
-   * registry.
-   */
-  interface ValueSupplier<K, V> {
-    V get(K key);
-  }
-
   class MapAdapter<K, V> implements WeakMap<K, V> {
     private final Object[] locks = new Object[16];
     private final Map<K, V> map;
@@ -111,14 +104,14 @@ public void putIfAbsent(final K key, final V value) {
     }
 
     @Override
-    public V computeIfAbsent(final K key, final ValueSupplier<? super K, ? extends V> supplier) {
+    public V computeIfAbsent(final K key, final Function<? super K, ? extends V> supplier) {
       // We can't use computeIfAbsent since it was added in 1.8.
       V value = map.get(key);
       if (null == value) {
         synchronized (locks[key.hashCode() & (locks.length - 1)]) {
           value = map.get(key);
           if (null == value) {
-            value = supplier.get(key);
+            value = supplier.apply(key);
             map.put(key, value);
           }
         }

@@ -1,5 +1,6 @@
 package datadog.trace.bootstrap
 
+import datadog.trace.api.Function
 import spock.lang.Specification
 
 class WeakMapTest extends Specification {
@@ -39,12 +40,12 @@ class WeakMapTest extends Specification {
     supplier.counter == 2
   }
 
-  class CounterSupplier implements WeakMap.ValueSupplier<String, Integer> {
+  class CounterSupplier implements Function<String, Integer> {
 
     def counter = 0
 
     @Override
-    Integer get(String ignored) {
+    Integer apply(String ignored) {
       counter = counter + 1
       return counter
     }

@@ -3,6 +3,18 @@ apply from: "$rootDir/gradle/java.gradle"
 minimumBranchCoverage = 0.6
 excludedClassesCoverage += ['datadog.trace.agent.tooling.*']
 
+// patch inner class from Caffeine to avoid ForkJoinTask from being loaded too early
+sourceSets {
+  patch {
+    java {}
+  }
+}
+jar {
+  from(sourceSets.patch.output) {
+    include 'com/github/benmanes/caffeine/cache/BoundedLocalCache$PerformCleanupTask.class'
+  }
+}
+
 configurations {
   // classpath used by the instrumentation muzzle plugin
   instrumentationMuzzle
@@ -12,13 +24,11 @@ dependencies {
   compile(project(':dd-java-agent:agent-bootstrap')) {
     exclude group: 'com.datadoghq', module: 'agent-logging'
   }
-  compile group: 'com.blogspot.mydailyjava', name: 'weak-lock-free', version: '0.15'
+  compile group: 'com.blogspot.mydailyjava', name: 'weak-lock-free', version: '0.17'
+  compile group: 'com.googlecode.concurrentlinkedhashmap', name: 'concurrentlinkedhashmap-lru', version: '1.4.2'
+  compile group: 'com.github.ben-manes.caffeine', name: 'caffeine', version: '2.8.6'
   compile deps.bytebuddy
   compile deps.bytebuddyagent
-  compile deps.guava
-
-  annotationProcessor deps.autoserviceProcessor
-  compileOnly deps.autoserviceAnnotation
 
   compile project(':dd-trace-core')
   compile project(':dd-trace-core:jfr-openjdk')

@@ -2,17 +2,18 @@
 
 import datadog.trace.agent.tooling.bytebuddy.DDCachingPoolStrategy;
 import datadog.trace.agent.tooling.bytebuddy.DDLocationStrategy;
+import datadog.trace.api.Platform;
 import datadog.trace.bootstrap.WeakCache;
 import datadog.trace.bootstrap.WeakCache.Provider;
 import datadog.trace.bootstrap.WeakMap;
-import java.util.Iterator;
-import java.util.ServiceLoader;
+import lombok.extern.slf4j.Slf4j;
 
 /**
  * This class contains class references for objects shared by the agent installer as well as muzzle
  * (both compile and runtime). Extracted out from AgentInstaller to begin separating some of the
  * logic out.
  */
+@Slf4j
 public class AgentTooling {
 
   static {
@@ -27,27 +28,37 @@ static void registerWeakMapProvider() {
     }
   }
 
-  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");
+  private static Provider loadWeakCacheProvider() {
+    ClassLoader classLoader = AgentInstaller.class.getClassLoader();
+    Class<Provider> providerClass;
+
+    try {
+      if (Platform.isJavaVersionAtLeast(8)) {
+        providerClass =
+            (Class<Provider>)
+                classLoader.loadClass("datadog.trace.agent.tooling.CaffeineWeakCache$Provider");
+        log.debug("Using CaffeineWeakCache Provider");
+      } else {
+        providerClass =
+            (Class<Provider>)
+                classLoader.loadClass("datadog.trace.agent.tooling.CLHMWeakCache$Provider");
+        log.debug("Using CLHMWeakCache Provider");
       }
-      return provider;
+
+      return providerClass.getDeclaredConstructor().newInstance();
+    } catch (ReflectiveOperationException e) {
+      throw new IllegalStateException("Can't load implementation of WeakCache.Provider", e);
     }
-    throw new IllegalStateException("Can't load implementation of WeakCache.Provider");
   }
 
+  private static final long DEFAULT_CACHE_CAPACITY = 32;
   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();
+    return newWeakCache(DEFAULT_CACHE_CAPACITY);
   }
 
   public static <K, V> WeakCache<K, V> newWeakCache(final long maxSize) {

@@ -0,0 +1,68 @@
+package datadog.trace.agent.tooling;
+
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
+import com.googlecode.concurrentlinkedhashmap.ConcurrentLinkedHashMap;
+import datadog.trace.api.Function;
+import datadog.trace.bootstrap.WeakCache;
+import java.util.concurrent.ConcurrentMap;
+
+public class CLHMWeakCache<K, V> implements WeakCache<K, V> {
+  public static final class Provider implements WeakCache.Provider {
+    @Override
+    public <K, V> WeakCache<K, V> newWeakCache(long maxSize) {
+      return new CLHMWeakCache<>(maxSize);
+    }
+  }
+
+  private static final int CACHE_CONCURRENCY =
+      Math.max(8, Runtime.getRuntime().availableProcessors());
+  private final WeakConcurrentMap<K, V> weakMap;
+  private final long maxSize;
+
+  public CLHMWeakCache(long maxSize) {
+    // No parameterization because WeakKey isn't visible
+    ConcurrentMap linkedMap =
+        new ConcurrentLinkedHashMap.Builder()
+            .maximumWeightedCapacity(maxSize)
+            .concurrencyLevel(CACHE_CONCURRENCY)
+            .build();
+
+    weakMap = new WeakConcurrentMap<>(false, true, linkedMap);
+
+    this.maxSize = maxSize;
+  }
+
+  @Override
+  public V getIfPresent(K key) {
+    return weakMap.getIfPresent(key);
+  }
+
+  @Override
+  public V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction) {
+    V value = weakMap.getIfPresent(key);
+    if (value == null) {
+      value = mappingFunction.apply(key);
+
+      expungeIfNecessary();
+      V oldValue = weakMap.putIfProbablyAbsent(key, value);
+      if (oldValue != null) {
+        value = oldValue;
+      }
+    }
+
+    return value;
+  }
+
+  @Override
+  public void put(K key, V value) {
+    expungeIfNecessary();
+
+    weakMap.put(key, value);
+  }
+
+  private void expungeIfNecessary() {
+    if (weakMap.approximateSize() >= maxSize) {
+      weakMap.expungeStaleEntries();
+    }
+  }
+}
@@ -0,0 +1,53 @@
+package datadog.trace.agent.tooling;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import datadog.trace.api.Function;
+import datadog.trace.bootstrap.WeakCache;
+import datadog.trace.util.AgentTaskScheduler;
+import java.util.concurrent.TimeUnit;
+
+public class CaffeineWeakCache<K, V> implements WeakCache<K, V> {
+  public static final class Provider implements WeakCache.Provider {
+    @Override
+    public <K, V> WeakCache<K, V> newWeakCache(long maxSize) {
+      return new CaffeineWeakCache<>(maxSize);
+    }
+  }
+
+  private final Cache<K, V> cache;
+
+  public CaffeineWeakCache(long maxSize) {
+    cache =
+        Caffeine.newBuilder()
+            .weakKeys()
+            .maximumSize(maxSize)
+            .expireAfterAccess(10, TimeUnit.MINUTES)
+            .executor(AgentTaskScheduler.INSTANCE)
+            .build();
+  }
+
+  @Override
+  public V getIfPresent(K key) {
+    return cache.getIfPresent(key);
+  }
+
+  @Override
+  public V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction) {
+    // Unable to use cache.get() directly because it relies on java.util.function.Function which is
+    // only available in Java8+.  This is less efficient.  The raciness is unimportant because this
+    // is a cache
+    V value = cache.getIfPresent(key);
+    if (value == null) {
+      value = mappingFunction.apply(key);
+
+      cache.put(key, value);
+    }
+    return value;
+  }
+
+  @Override
+  public void put(K key, V value) {
+    cache.put(key, value);
+  }
+}