diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/DDCachingPoolStrategy.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/DDCachingPoolStrategy.java index 112f9f86ad..1f8fdc2501 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/DDCachingPoolStrategy.java +++ b/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. * - *

This is similar to: AgentBuilder.PoolStrategy.WithTypePoolCache.Simple(new - * MapMaker().weakKeys().makeMap()) + *

Here we are using WeakMap.Provider as the backing ClassLoader -> CacheProvider lookup. We also + * use our bootstrap proxy when matching against the bootstrap loader. * - *

Main differences: - * - *

    - *
  1. Control over the type of the cache. We many not want to use a java.util.ConcurrentMap - *
  2. Use our bootstrap proxy when matching against the bootstrap loader. - *
+ *

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 typePoolCache = - Collections.synchronizedMap(new WeakHashMap()); + private static final WeakMap 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 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 { + private final TypePool.Resolution value; + + private ResolutionProvider(final TypePool.Resolution value) { + this.value = value; + } + + @Override + public TypePool.Resolution call() { + return value; + } + } + } }