Merge pull request #468 from DataDog/tyler/evict-typepool

Add eviction to a TypePool.CacheProvider’s cache

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

@@ -1,33 +1,33 @@
 package datadog.trace.agent.tooling;
 
 import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER;
-import static net.bytebuddy.agent.builder.AgentBuilder.*;
+import static net.bytebuddy.agent.builder.AgentBuilder.PoolStrategy;
 
-import java.util.Collections;
-import java.util.Map;
-import java.util.WeakHashMap;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import datadog.trace.bootstrap.WeakMap;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.ClassFileLocator;
 import net.bytebuddy.pool.TypePool;
 
 /**
  * Custom Pool strategy.
  *
- * <p>This is similar to: AgentBuilder.PoolStrategy.WithTypePoolCache.Simple(new
- * MapMaker().weakKeys().<ClassLoader, TypePool.CacheProvider>makeMap())
+ * <p>Here we are using WeakMap.Provider as the backing ClassLoader -> CacheProvider lookup. We also
+ * use our bootstrap proxy when matching against the bootstrap loader.
  *
- * <p>Main differences:
- *
- * <ol>
- *   <li>Control over the type of the cache. We many not want to use a java.util.ConcurrentMap
- *   <li>Use our bootstrap proxy when matching against the bootstrap loader.
- * </ol>
+ * <p>The CacheProvider is also a custom implementation that uses guava's cache to evict. See
+ * eviction policy below.
  */
 public class DDCachingPoolStrategy implements PoolStrategy {
-  private static final Map<ClassLoader, TypePool.CacheProvider> typePoolCache =
-      Collections.synchronizedMap(new WeakHashMap<ClassLoader, TypePool.CacheProvider>());
+  private static final WeakMap<ClassLoader, TypePool.CacheProvider> typePoolCache =
+      WeakMap.Provider.newWeakMap();
 
   @Override
-  public TypePool typePool(ClassFileLocator classFileLocator, ClassLoader classLoader) {
+  public TypePool typePool(final ClassFileLocator classFileLocator, final ClassLoader classLoader) {
     final ClassLoader key =
         BOOTSTRAP_CLASSLOADER == classLoader ? Utils.getBootstrapProxy() : classLoader;
     TypePool.CacheProvider cache = typePoolCache.get(key);
@@ -35,7 +35,7 @@ public class DDCachingPoolStrategy implements PoolStrategy {
       synchronized (key) {
         cache = typePoolCache.get(key);
         if (null == cache) {
-          cache = TypePool.CacheProvider.Simple.withObjectType();
+          cache = EvictingCacheProvider.withObjectType();
           typePoolCache.put(key, cache);
         }
       }
@@ -43,4 +43,59 @@ public class DDCachingPoolStrategy implements PoolStrategy {
     return new TypePool.Default.WithLazyResolution(
         cache, classFileLocator, TypePool.Default.ReaderMode.FAST);
   }
+
+  private static class EvictingCacheProvider implements TypePool.CacheProvider {
+
+    /** A map containing all cached resolutions by their names. */
+    private final Cache<String, TypePool.Resolution> cache;
+
+    /** Creates a new simple cache. */
+    private EvictingCacheProvider() {
+      cache =
+          CacheBuilder.newBuilder()
+              .initialCapacity(100)
+              .maximumSize(1000)
+              .expireAfterAccess(1, TimeUnit.MINUTES)
+              .build();
+    }
+
+    private static TypePool.CacheProvider withObjectType() {
+      final TypePool.CacheProvider cacheProvider = new EvictingCacheProvider();
+      cacheProvider.register(
+          Object.class.getName(), new TypePool.Resolution.Simple(TypeDescription.OBJECT));
+      return cacheProvider;
+    }
+
+    @Override
+    public TypePool.Resolution find(final String name) {
+      return cache.getIfPresent(name);
+    }
+
+    @Override
+    public TypePool.Resolution register(final String name, final TypePool.Resolution resolution) {
+      try {
+        return cache.get(name, new ResolutionProvider(resolution));
+      } catch (final ExecutionException e) {
+        return resolution;
+      }
+    }
+
+    @Override
+    public void clear() {
+      cache.invalidateAll();
+    }
+
+    private class ResolutionProvider implements Callable<TypePool.Resolution> {
+      private final TypePool.Resolution value;
+
+      private ResolutionProvider(final TypePool.Resolution value) {
+        this.value = value;
+      }
+
+      @Override
+      public TypePool.Resolution call() {
+        return value;
+      }
+    }
+  }
 }