WeakMap computeIfAbsent
MODIFY
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java
@@ -18,7 +18,7 @@ public interface WeakMap<K, V> {
void putIfAbsent(K key, V value);
- V getOrCreate(K key, ValueSupplier<V> supplier);
+ V computeIfAbsent(K key, ValueSupplier<? super K, ? extends V> supplier);
@Slf4j
class Provider {
@@ -62,8 +62,8 @@ public interface WeakMap<K, V> {
* Supplies the value to be stored and it is called only when a value does not exists yet in the
* registry.
*/
- interface ValueSupplier<V> {
- V get();
+ interface ValueSupplier<K, V> {
+ V get(K key);
}
class MapAdapter<K, V> implements WeakMap<K, V> {
@@ -107,11 +107,12 @@ public interface WeakMap<K, V> {
}
@Override
- public V getOrCreate(final K key, final ValueSupplier<V> supplier) {
+ public V computeIfAbsent(final K key, final ValueSupplier<? super K, ? extends V> supplier) {
+ // We can't use computeIfAbsent since it was added in 1.8.
if (!map.containsKey(key)) {
synchronized (this) {
if (!map.containsKey(key)) {
- map.put(key, supplier.get());
+ map.put(key, supplier.get(key));
}
}
}

MODIFY
dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/WeakMapTest.groovy
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<Integer> {
+ class CounterSupplier implements WeakMap.ValueSupplier<String, Integer> {
def counter = 0
@Override
- Integer get() {
+ Integer get(String ignored) {
counter = counter + 1
return counter
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderMatcher.java
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<ClassLoader> classLoaderHasClassWithField(
- final String className, final String fieldName) {
- return new ClassLoaderHasClassWithFieldMatcher(className, fieldName);
- }
-
- public static ElementMatcher.Junction.AbstractBase<ClassLoader> classLoaderHasClassWithMethod(
- final String className, final String methodName, final String... methodArgs) {
- return new ClassLoaderHasClassWithMethodMatcher(className, methodName, methodArgs);
- }
-
private static class SkipClassLoaderMatcher
- extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
+ extends ElementMatcher.Junction.AbstractBase<ClassLoader>
+ implements WeakMap.ValueSupplier<ClassLoader, Boolean> {
public static final SkipClassLoaderMatcher INSTANCE = new SkipClassLoaderMatcher();
/* Cache of classloader-instance -> (true|false). True = skip instrumentation. False = safe to instrument. */
private static final WeakMap<ClassLoader, Boolean> SKIP_CACHE = newWeakMap();
@@ -74,24 +65,19 @@ 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();
+ 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());
}
- SKIP_CACHE.put(loader, skip);
+
return skip;
}
- }
/**
* TODO: this turns out to be useless with OSGi: {@code
@@ -121,23 +107,9 @@ public class ClassLoaderMatcher {
}
}
- private static class ClassLoaderNameMatcher
- extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
-
- 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<ClassLoader> {
+ extends ElementMatcher.Junction.AbstractBase<ClassLoader>
+ implements WeakMap.ValueSupplier<ClassLoader, Boolean> {
private final WeakMap<ClassLoader, Boolean> 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;
+ return cache.computeIfAbsent(target, this);
}
- 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 false;
- }
- }
-
- public static class ClassLoaderHasClassWithFieldMatcher
- extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
- private final WeakMap<ClassLoader, Boolean> cache = newWeakMap();
-
- private final String className;
- private final String fieldName;
-
- private ClassLoaderHasClassWithFieldMatcher(final String className, final String fieldName) {
- this.className = className;
- this.fieldName = fieldName;
+ return false;
}
@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;
}
}
- public static class ClassLoaderHasClassWithMethodMatcher
- extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
-
- private final WeakMap<ClassLoader, Boolean> 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;
}
}
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/DDCachingPoolStrategy.java
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;
*
* <p>See eviction policy below.
*/
-public class DDCachingPoolStrategy implements PoolStrategy {
+public class DDCachingPoolStrategy
+ implements PoolStrategy, WeakMap.ValueSupplier<ClassLoader, TypePool.CacheProvider> {
+
+ // 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<ClassLoader, TypePool.CacheProvider> typePoolCache =
WeakMap.Provider.newWeakMap();
private final Cleaner cleaner;
@@ -40,27 +47,24 @@ 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)) {
+ 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.
- cache = TypePool.CacheProvider.NoOp.INSTANCE;
+ return TypePool.CacheProvider.NoOp.INSTANCE;
} else {
- cache = EvictingCacheProvider.withObjectType(cleaner, 1, TimeUnit.MINUTES);
- }
- typePoolCache.put(key, cache);
- }
+ return EvictingCacheProvider.withObjectType(cleaner, 1, TimeUnit.MINUTES);
}
}
- return new TypePool.Default.WithLazyResolution(
- cache, classFileLocator, TypePool.Default.ReaderMode.FAST);
- }
private static class EvictingCacheProvider implements TypePool.CacheProvider {

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java
@@ -41,7 +41,7 @@ class WeakMapSuppliers {
public <K, V> WeakMap<K, V> get() {
final WeakConcurrentMap<K, V> 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<WeakConcurrentMap> {
@@ -86,11 +86,11 @@ class WeakMapSuppliers {
}
@Override
- public V getOrCreate(final K key, final ValueSupplier<V> supplier) {
+ public V computeIfAbsent(final K key, final ValueSupplier<? super K, ? extends V> supplier) {
if (!map.containsKey(key)) {
synchronized (this) {
if (!map.containsKey(key)) {
- map.put(key, supplier.get());
+ map.put(key, supplier.get(key));
}
}
}
@@ -103,7 +103,7 @@ class WeakMapSuppliers {
@Override
public <K, V> WeakMap<K, V> get() {
- return new Adapter(new WeakConcurrentMap.WithInlinedExpunction<K, V>());
+ return new Adapter<>(new WeakConcurrentMap.WithInlinedExpunction<K, V>());
}
}
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java
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<ClassLoader, List<Reference.Mismatch>> {
private final WeakMap<ClassLoader, List<Reference.Mismatch>> mismatchCache = newWeakMap();
private final Reference[] references;
private final Set<String> 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,12 +57,14 @@ public class ReferenceMatcher {
if (loader == BOOTSTRAP_LOADER) {
loader = Utils.getBootstrapProxy();
}
- List<Reference.Mismatch> mismatches = mismatchCache.get(loader);
- if (null == mismatches) {
- synchronized (loader) {
- mismatches = mismatchCache.get(loader);
- if (null == mismatches) {
- mismatches = new ArrayList<>(0);
+
+ return mismatchCache.computeIfAbsent(loader, this);
+ }
+
+ @Override
+ public List<Mismatch> get(final ClassLoader loader) {
+ final List<Mismatch> mismatches = new ArrayList<>(0);
+
for (final Reference reference : references) {
// Don't reference-check helper classes.
// They will be injected by the instrumentation's HelperInjector.
@@ -69,10 +72,7 @@ public class ReferenceMatcher {
mismatches.addAll(checkMatch(reference, loader));
}
}
- mismatchCache.put(loader, mismatches);
- }
- }
- }
+
return mismatches;
}

MODIFY
dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java
dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java
@@ -14,10 +14,11 @@ public class AttributeKeys {
private static final WeakMap<ClassLoader, Map<String, AttributeKey<?>>> map =
WeakMap.Implementation.DEFAULT.get();
- private static final WeakMap.ValueSupplier<Map<String, AttributeKey<?>>> mapSupplier =
- new WeakMap.ValueSupplier<Map<String, AttributeKey<?>>>() {
+ private static final WeakMap.ValueSupplier<ClassLoader, Map<String, AttributeKey<?>>>
+ mapSupplier =
+ new WeakMap.ValueSupplier<ClassLoader, Map<String, AttributeKey<?>>>() {
@Override
- public Map<String, AttributeKey<?>> get() {
+ public Map<String, AttributeKey<?>> get(final ClassLoader ignored) {
return new ConcurrentHashMap<>();
}
};
@@ -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 <T> AttributeKey<T> attributeKey(String key) {
- Map<String, AttributeKey<?>> classLoaderMap =
- map.getOrCreate(AttributeKey.class.getClassLoader(), mapSupplier);
+ private static <T> AttributeKey<T> attributeKey(final String key) {
+ final Map<String, AttributeKey<?>> classLoaderMap =
+ map.computeIfAbsent(AttributeKey.class.getClassLoader(), mapSupplier);
if (classLoaderMap.containsKey(key)) {
return (AttributeKey<T>) classLoaderMap.get(key);
}

MODIFY
dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java
dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java
@@ -14,10 +14,11 @@ public class AttributeKeys {
private static final WeakMap<ClassLoader, Map<String, AttributeKey<?>>> map =
WeakMap.Implementation.DEFAULT.get();
- private static final WeakMap.ValueSupplier<Map<String, AttributeKey<?>>> mapSupplier =
- new WeakMap.ValueSupplier<Map<String, AttributeKey<?>>>() {
+ private static final WeakMap.ValueSupplier<ClassLoader, Map<String, AttributeKey<?>>>
+ mapSupplier =
+ new WeakMap.ValueSupplier<ClassLoader, Map<String, AttributeKey<?>>>() {
@Override
- public Map<String, AttributeKey<?>> get() {
+ public Map<String, AttributeKey<?>> get(final ClassLoader ignored) {
return new ConcurrentHashMap<>();
}
};
@@ -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 <T> AttributeKey<T> attributeKey(String key) {
- Map<String, AttributeKey<?>> classLoaderMap =
- map.getOrCreate(AttributeKey.class.getClassLoader(), mapSupplier);
+ private static <T> AttributeKey<T> attributeKey(final String key) {
+ final Map<String, AttributeKey<?>> classLoaderMap =
+ map.computeIfAbsent(AttributeKey.class.getClassLoader(), mapSupplier);
if (classLoaderMap.containsKey(key)) {
return (AttributeKey<T>) classLoaderMap.get(key);
}

WeakMap computeIfAbsent