diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java index 29b3a5477a..4486ff4132 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java @@ -18,7 +18,7 @@ public interface WeakMap { void putIfAbsent(K key, V value); - V getOrCreate(K key, ValueSupplier supplier); + V computeIfAbsent(K key, ValueSupplier supplier); @Slf4j class Provider { @@ -62,8 +62,8 @@ public interface WeakMap { * Supplies the value to be stored and it is called only when a value does not exists yet in the * registry. */ - interface ValueSupplier { - V get(); + interface ValueSupplier { + V get(K key); } class MapAdapter implements WeakMap { @@ -107,16 +107,22 @@ public interface WeakMap { } @Override - public V getOrCreate(final K key, final ValueSupplier supplier) { - if (!map.containsKey(key)) { - synchronized (this) { - if (!map.containsKey(key)) { - map.put(key, supplier.get()); - } - } + public V computeIfAbsent(final K key, final ValueSupplier supplier) { + // We can't use computeIfAbsent since it was added in 1.8. + if (map.containsKey(key)) { + return map.get(key); } - return map.get(key); + synchronized (this) { + if (map.containsKey(key)) { + return map.get(key); + } else { + final V value = supplier.get(key); + + map.put(key, value); + return value; + } + } } @Override diff --git a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/WeakMapTest.groovy b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/WeakMapTest.groovy index 38e3c493e0..77fe229688 100644 --- a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/WeakMapTest.groovy +++ b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/WeakMapTest.groovy @@ -10,7 +10,7 @@ class WeakMapTest extends Specification { def "getOrCreate a value"() { when: - def count = sut.getOrCreate('key', supplier) + def count = sut.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.getOrCreate('key', supplier) - def count2 = sut.getOrCreate('key', supplier) + def count1 = sut.computeIfAbsent('key', supplier) + def count2 = sut.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.getOrCreate('key1', supplier) - def count2 = sut.getOrCreate('key2', supplier) + def count1 = sut.computeIfAbsent('key1', supplier) + def count2 = sut.computeIfAbsent('key2', supplier) then: count1 == 1 @@ -39,12 +39,12 @@ class WeakMapTest extends Specification { supplier.counter == 2 } - class CounterSupplier implements WeakMap.ValueSupplier { + class CounterSupplier implements WeakMap.ValueSupplier { def counter = 0 @Override - Integer get() { + Integer get(String ignored) { counter = counter + 1 return counter } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderMatcher.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderMatcher.java index a53eb694f6..ac776607c0 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderMatcher.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderMatcher.java @@ -30,18 +30,9 @@ public class ClassLoaderMatcher { return new ClassLoaderHasClassMatcher(names); } - public static ElementMatcher.Junction.AbstractBase classLoaderHasClassWithField( - final String className, final String fieldName) { - return new ClassLoaderHasClassWithFieldMatcher(className, fieldName); - } - - public static ElementMatcher.Junction.AbstractBase classLoaderHasClassWithMethod( - final String className, final String methodName, final String... methodArgs) { - return new ClassLoaderHasClassWithMethodMatcher(className, methodName, methodArgs); - } - private static class SkipClassLoaderMatcher - extends ElementMatcher.Junction.AbstractBase { + extends ElementMatcher.Junction.AbstractBase + implements WeakMap.ValueSupplier { public static final SkipClassLoaderMatcher INSTANCE = new SkipClassLoaderMatcher(); /* Cache of classloader-instance -> (true|false). True = skip instrumentation. False = safe to instrument. */ private static final WeakMap SKIP_CACHE = newWeakMap(); @@ -74,23 +65,18 @@ public class ClassLoaderMatcher { } private boolean shouldSkipInstance(final ClassLoader loader) { - Boolean cached = SKIP_CACHE.get(loader); - if (null != cached) { - return cached.booleanValue(); - } - synchronized (this) { - cached = SKIP_CACHE.get(loader); - if (null != cached) { - return cached.booleanValue(); - } - final boolean skip = !delegatesToBootstrap(loader); - if (skip) { - log.debug( - "skipping classloader instance {} of type {}", loader, loader.getClass().getName()); - } - SKIP_CACHE.put(loader, skip); - return skip; + return SKIP_CACHE.computeIfAbsent(loader, this); + } + + @Override + public Boolean get(final ClassLoader loader) { + final boolean skip = !delegatesToBootstrap(loader); + if (skip) { + log.debug( + "skipping classloader instance {} of type {}", loader, loader.getClass().getName()); } + + return skip; } /** @@ -121,23 +107,9 @@ public class ClassLoaderMatcher { } } - private static class ClassLoaderNameMatcher - extends ElementMatcher.Junction.AbstractBase { - - private final String name; - - private ClassLoaderNameMatcher(final String name) { - this.name = name; - } - - @Override - public boolean matches(final ClassLoader target) { - return target != null && name.equals(target.getClass().getName()); - } - } - public static class ClassLoaderHasClassMatcher - extends ElementMatcher.Junction.AbstractBase { + extends ElementMatcher.Junction.AbstractBase + implements WeakMap.ValueSupplier { private final WeakMap cache = newWeakMap(); @@ -150,123 +122,21 @@ public class ClassLoaderMatcher { @Override public boolean matches(final ClassLoader target) { if (target != null) { - Boolean result = cache.get(target); - if (result != null) { - return result; - } - synchronized (target) { - result = cache.get(target); - if (result != null) { - return result; - } - for (final String name : names) { - if (target.getResource(Utils.getResourceName(name)) == null) { - cache.put(target, false); - return false; - } - } - cache.put(target, true); - return true; - } + return cache.computeIfAbsent(target, this); } + return false; } - } - - public static class ClassLoaderHasClassWithFieldMatcher - extends ElementMatcher.Junction.AbstractBase { - - private final WeakMap cache = newWeakMap(); - - private final String className; - private final String fieldName; - - private ClassLoaderHasClassWithFieldMatcher(final String className, final String fieldName) { - this.className = className; - this.fieldName = fieldName; - } @Override - public boolean matches(final ClassLoader target) { - if (target != null) { - Boolean result = cache.get(target); - if (result != null) { - return result; - } - synchronized (target) { - result = cache.get(target); - if (result != null) { - return result; - } - try { - final Class aClass = Class.forName(className, false, target); - aClass.getDeclaredField(fieldName); - cache.put(target, true); - return true; - } catch (final ClassNotFoundException e) { - cache.put(target, false); - return false; - } catch (final NoSuchFieldException e) { - cache.put(target, false); - return false; - } + public Boolean get(final ClassLoader target) { + for (final String name : names) { + if (target.getResource(Utils.getResourceName(name)) == null) { + return false; } } - return false; - } - } - public static class ClassLoaderHasClassWithMethodMatcher - extends ElementMatcher.Junction.AbstractBase { - - private final WeakMap cache = newWeakMap(); - - private final String className; - private final String methodName; - private final String[] methodArgs; - - private ClassLoaderHasClassWithMethodMatcher( - final String className, final String methodName, final String... methodArgs) { - this.className = className; - this.methodName = methodName; - this.methodArgs = methodArgs; - } - - @Override - public boolean matches(final ClassLoader target) { - if (target != null) { - Boolean result = cache.get(target); - if (result != null) { - return result; - } - synchronized (target) { - result = cache.get(target); - if (result != null) { - return result; - } - try { - final Class aClass = Class.forName(className, false, target); - final Class[] methodArgsClasses = new Class[methodArgs.length]; - for (int i = 0; i < methodArgs.length; ++i) { - methodArgsClasses[i] = target.loadClass(methodArgs[i]); - } - if (aClass.isInterface()) { - aClass.getMethod(methodName, methodArgsClasses); - } else { - aClass.getDeclaredMethod(methodName, methodArgsClasses); - } - cache.put(target, true); - return true; - } catch (final ClassNotFoundException e) { - cache.put(target, false); - return false; - } catch (final NoSuchMethodException e) { - cache.put(target, false); - return false; - } - } - } - return false; + return true; } } } 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 6315cabe89..886d742c78 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 @@ -7,6 +7,7 @@ import static net.bytebuddy.agent.builder.AgentBuilder.PoolStrategy; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; import datadog.trace.bootstrap.WeakMap; +import java.security.SecureClassLoader; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; @@ -28,7 +29,13 @@ import net.bytebuddy.pool.TypePool; * *

See eviction policy below. */ -public class DDCachingPoolStrategy implements PoolStrategy { +public class DDCachingPoolStrategy + implements PoolStrategy, WeakMap.ValueSupplier { + + // Need this because we can't put null into the typePoolCache map. + private static final ClassLoader BOOTSTRAP_CLASSLOADER_PLACEHOLDER = + new SecureClassLoader(null) {}; + private final WeakMap typePoolCache = WeakMap.Provider.newWeakMap(); private final Cleaner cleaner; @@ -40,28 +47,25 @@ public class DDCachingPoolStrategy implements PoolStrategy { @Override public TypePool typePool(final ClassFileLocator classFileLocator, final ClassLoader classLoader) { final ClassLoader key = - BOOTSTRAP_CLASSLOADER == classLoader ? Utils.getBootstrapProxy() : classLoader; - TypePool.CacheProvider cache = typePoolCache.get(key); - if (null == cache) { - synchronized (key) { - cache = typePoolCache.get(key); - if (null == cache) { - if (skipClassLoader().matches(classLoader)) { - // Don't bother creating a cache for a classloader that won't match. - // (avoiding a lot of DelegatingClassLoader instances) - // This is primarily an optimization. - cache = TypePool.CacheProvider.NoOp.INSTANCE; - } else { - cache = EvictingCacheProvider.withObjectType(cleaner, 1, TimeUnit.MINUTES); - } - typePoolCache.put(key, cache); - } - } - } + BOOTSTRAP_CLASSLOADER == classLoader ? BOOTSTRAP_CLASSLOADER_PLACEHOLDER : classLoader; + final TypePool.CacheProvider cache = typePoolCache.computeIfAbsent(key, this); + return new TypePool.Default.WithLazyResolution( cache, classFileLocator, TypePool.Default.ReaderMode.FAST); } + @Override + public TypePool.CacheProvider get(final ClassLoader key) { + if (BOOTSTRAP_CLASSLOADER_PLACEHOLDER != key && skipClassLoader().matches(key)) { + // Don't bother creating a cache for a classloader that won't match. + // (avoiding a lot of DelegatingClassLoader instances) + // This is primarily an optimization. + return TypePool.CacheProvider.NoOp.INSTANCE; + } else { + return EvictingCacheProvider.withObjectType(cleaner, 1, TimeUnit.MINUTES); + } + } + private static class EvictingCacheProvider implements TypePool.CacheProvider { /** A map containing all cached resolutions by their names. */ diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java index b5830d5460..32a16a0b7c 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java @@ -41,7 +41,7 @@ class WeakMapSuppliers { public WeakMap get() { final WeakConcurrentMap map = new WeakConcurrentMap<>(false); cleaner.scheduleCleaning(map, MapCleaner.CLEANER, CLEAN_FREQUENCY_SECONDS, TimeUnit.SECONDS); - return new Adapter(map); + return new Adapter<>(map); } private static class MapCleaner implements Cleaner.Adapter { @@ -86,16 +86,21 @@ class WeakMapSuppliers { } @Override - public V getOrCreate(final K key, final ValueSupplier supplier) { - if (!map.containsKey(key)) { - synchronized (this) { - if (!map.containsKey(key)) { - map.put(key, supplier.get()); - } - } + public V computeIfAbsent(final K key, final ValueSupplier supplier) { + if (map.containsKey(key)) { + return map.get(key); } - return map.get(key); + synchronized (this) { + if (map.containsKey(key)) { + return map.get(key); + } else { + final V value = supplier.get(key); + + map.put(key, value); + return value; + } + } } } @@ -103,7 +108,7 @@ class WeakMapSuppliers { @Override public WeakMap get() { - return new Adapter(new WeakConcurrentMap.WithInlinedExpunction()); + return new Adapter<>(new WeakConcurrentMap.WithInlinedExpunction()); } } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java index 4f2ed18bd0..c669646024 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java @@ -22,7 +22,8 @@ import net.bytebuddy.pool.TypePool; /** Matches a set of references against a classloader. */ @Slf4j -public class ReferenceMatcher { +public class ReferenceMatcher + implements WeakMap.ValueSupplier> { private final WeakMap> mismatchCache = newWeakMap(); private final Reference[] references; private final Set helperClassNames; @@ -45,7 +46,7 @@ public class ReferenceMatcher { * @return true if all references match the classpath of loader */ public boolean matches(final ClassLoader loader) { - return getMismatchedReferenceSources(loader).size() == 0; + return getMismatchedReferenceSources(loader).isEmpty(); } /** @@ -56,23 +57,22 @@ public class ReferenceMatcher { if (loader == BOOTSTRAP_LOADER) { loader = Utils.getBootstrapProxy(); } - List mismatches = mismatchCache.get(loader); - if (null == mismatches) { - synchronized (loader) { - mismatches = mismatchCache.get(loader); - if (null == mismatches) { - mismatches = new ArrayList<>(0); - for (final Reference reference : references) { - // Don't reference-check helper classes. - // They will be injected by the instrumentation's HelperInjector. - if (!helperClassNames.contains(reference.getClassName())) { - mismatches.addAll(checkMatch(reference, loader)); - } - } - mismatchCache.put(loader, mismatches); - } + + return mismatchCache.computeIfAbsent(loader, this); + } + + @Override + public List get(final ClassLoader loader) { + final List mismatches = new ArrayList<>(0); + + for (final Reference reference : references) { + // Don't reference-check helper classes. + // They will be injected by the instrumentation's HelperInjector. + if (!helperClassNames.contains(reference.getClassName())) { + mismatches.addAll(checkMatch(reference, loader)); } } + return mismatches; } diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java index 12c3f4d906..2bc752548a 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java @@ -14,13 +14,14 @@ public class AttributeKeys { private static final WeakMap>> map = WeakMap.Implementation.DEFAULT.get(); - private static final WeakMap.ValueSupplier>> mapSupplier = - new WeakMap.ValueSupplier>>() { - @Override - public Map> get() { - return new ConcurrentHashMap<>(); - } - }; + private static final WeakMap.ValueSupplier>> + mapSupplier = + new WeakMap.ValueSupplier>>() { + @Override + public Map> get(final ClassLoader ignored) { + return new ConcurrentHashMap<>(); + } + }; public static final AttributeKey PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY = @@ -42,9 +43,9 @@ public class AttributeKeys { * while the Attribute class is loaded by a third class loader and used internally for the * cassandra driver. */ - private static AttributeKey attributeKey(String key) { - Map> classLoaderMap = - map.getOrCreate(AttributeKey.class.getClassLoader(), mapSupplier); + private static AttributeKey attributeKey(final String key) { + final Map> classLoaderMap = + map.computeIfAbsent(AttributeKey.class.getClassLoader(), mapSupplier); if (classLoaderMap.containsKey(key)) { return (AttributeKey) classLoaderMap.get(key); } diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java index 7b632be122..916b2ef3b7 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java @@ -14,13 +14,14 @@ public class AttributeKeys { private static final WeakMap>> map = WeakMap.Implementation.DEFAULT.get(); - private static final WeakMap.ValueSupplier>> mapSupplier = - new WeakMap.ValueSupplier>>() { - @Override - public Map> get() { - return new ConcurrentHashMap<>(); - } - }; + private static final WeakMap.ValueSupplier>> + mapSupplier = + new WeakMap.ValueSupplier>>() { + @Override + public Map> get(final ClassLoader ignored) { + return new ConcurrentHashMap<>(); + } + }; public static final AttributeKey PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY = @@ -46,9 +47,9 @@ public class AttributeKeys { * while the Attribute class is loaded by a third class loader and used internally for the * cassandra driver. */ - private static AttributeKey attributeKey(String key) { - Map> classLoaderMap = - map.getOrCreate(AttributeKey.class.getClassLoader(), mapSupplier); + private static AttributeKey attributeKey(final String key) { + final Map> classLoaderMap = + map.computeIfAbsent(AttributeKey.class.getClassLoader(), mapSupplier); if (classLoaderMap.containsKey(key)) { return (AttributeKey) classLoaderMap.get(key); }