Refactor weak hash map to be able to reuse functionalities
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,21 +18,23 @@ public interface WeakMap<K, V> {
void putIfAbsent(K key, V value);
+ V getOrCreate(K key, ValueSupplier<V> supplier);
+
@Slf4j
class Provider {
- private static final AtomicReference<Supplier> provider =
- new AtomicReference<>(Supplier.DEFAULT);
+ private static final AtomicReference<Implementation> provider =
+ new AtomicReference<>(Implementation.DEFAULT);
- public static void registerIfAbsent(final Supplier provider) {
- if (provider != null && provider != Supplier.DEFAULT) {
- if (Provider.provider.compareAndSet(Supplier.DEFAULT, provider)) {
+ public static void registerIfAbsent(final Implementation provider) {
+ if (provider != null && provider != Implementation.DEFAULT) {
+ if (Provider.provider.compareAndSet(Implementation.DEFAULT, provider)) {
log.debug("Weak map provider set to {}", provider);
}
}
}
public static boolean isProviderRegistered() {
- return provider.get() != Supplier.DEFAULT;
+ return provider.get() != Implementation.DEFAULT;
}
public static <K, V> WeakMap<K, V> newWeakMap() {
@@ -40,13 +42,13 @@ public interface WeakMap<K, V> {
}
}
- interface Supplier {
+ interface Implementation {
<K, V> WeakMap<K, V> get();
- Supplier DEFAULT = new Default();
+ Implementation DEFAULT = new Default();
@Slf4j
- class Default implements Supplier {
+ class Default implements Implementation {
@Override
public <K, V> WeakMap<K, V> get() {
@@ -56,6 +58,14 @@ 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();
+ }
+
class MapAdapter<K, V> implements WeakMap<K, V> {
private final Map<K, V> map;
@@ -96,6 +106,16 @@ public interface WeakMap<K, V> {
}
}
+ @Override
+ public synchronized V getOrCreate(K key, ValueSupplier<V> supplier) {
+
+ if (!map.containsKey(key)) {
+ map.put(key, supplier.get());
+ }
+
+ return map.get(key);
+ }
+
@Override
public String toString() {
return map.toString();

ADD
dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/WeakMapTest.groovy
+package datadog.trace.bootstrap
+
+import spock.lang.Specification
+
+class WeakMapTest extends Specification {
+
+
+ def "getOrCreate a value"() {
+ setup:
+ def supplier = new CounterSupplier()
+ def sut = new WeakMap.MapAdapter<String, Integer>(new WeakHashMap<>())
+
+ when:
+ def count = sut.getOrCreate('key', supplier)
+
+ then:
+ count == 1
+ supplier.counter == 1
+ }
+
+ def "getOrCreate a value multiple times same class loader same key"() {
+ setup:
+ def supplier = new CounterSupplier()
+ def sut = new WeakMap.MapAdapter<String, Integer>(new WeakHashMap<>())
+
+ when:
+ def count1 = sut.getOrCreate('key', supplier)
+ def count2 = sut.getOrCreate('key', supplier)
+
+ then:
+ count1 == 1
+ count2 == 1
+ supplier.counter == 1
+ }
+
+ def "getOrCreate a value multiple times same class loader different keys"() {
+ setup:
+ def supplier = new CounterSupplier()
+ def sut = new WeakMap.MapAdapter<String, Integer>(new WeakHashMap<>())
+
+ when:
+ def count1 = sut.getOrCreate('key1', supplier)
+ def count2 = sut.getOrCreate('key2', supplier)
+
+ then:
+ count1 == 1
+ count2 == 2
+ supplier.counter == 2
+ }
+
+ class CounterSupplier implements WeakMap.ValueSupplier<Integer> {
+
+ def counter = 0
+
+ @Override
+ Integer get() {
+ counter = counter + 1
+ return counter
+ }
+ }
+}

DELETE
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderScopedWeakMap.java
-package datadog.trace.agent.tooling;
-
-import datadog.trace.bootstrap.WeakMap;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-
-/** A registry which is scoped per-classloader able to hold key-value pairs with weak keys. */
-public class ClassLoaderScopedWeakMap {
-
- public static final ClassLoaderScopedWeakMap INSTANCE = new ClassLoaderScopedWeakMap();
-
- private final WeakMap<ClassLoader, Map<Object, Object>> map = WeakMap.Supplier.DEFAULT.get();
-
- /**
- * Gets the element registered at the specified key or register as new one retrieved by the
- * provided supplier.
- */
- public synchronized Object getOrCreate(
- ClassLoader classLoader, Object key, Supplier valueSupplier) {
- Map<Object, Object> classLoaderMap;
-
- if (!map.containsKey(classLoader)) {
- classLoaderMap = new ConcurrentHashMap<>();
- map.put(classLoader, classLoaderMap);
- } else {
- classLoaderMap = map.get(classLoader);
- }
-
- if (classLoaderMap.containsKey(key)) {
- return classLoaderMap.get(key);
- }
-
- final Object value = valueSupplier.get();
- classLoaderMap.put(key, value);
- return value;
- }
-
- /**
- * Supplies the value to be stored and it is called only when a value does not exists yet in the
- * registry.
- */
- public interface Supplier {
- Object get();
- }
-}

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
@@ -36,7 +36,7 @@ class WeakMapSuppliers {
* single thread to clean void weak references out for all instances. Cleaning is done every
* second.
*/
- static class WeakConcurrent implements WeakMap.Supplier {
+ static class WeakConcurrent implements WeakMap.Implementation {
private static final long SHUTDOWN_WAIT_SECONDS = 5;
@@ -177,7 +177,7 @@ class WeakMapSuppliers {
}
}
- static class Inline implements WeakMap.Supplier {
+ static class Inline implements WeakMap.Implementation {
@Override
public <K, V> WeakMap<K, V> get() {
@@ -186,7 +186,7 @@ class WeakMapSuppliers {
}
}
- static class Guava implements WeakMap.Supplier {
+ static class Guava implements WeakMap.Implementation {
@Override
public <K, V> WeakMap<K, V> get() {

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java
@@ -25,7 +25,7 @@ public class MuzzleVersionScanPlugin {
static {
// prevent WeakMap from logging warning while plugin is running
WeakMap.Provider.registerIfAbsent(
- new WeakMap.Supplier() {
+ new WeakMap.Implementation() {
@Override
public <K, V> WeakMap<K, V> get() {
return new WeakMap.MapAdapter<>(Collections.synchronizedMap(new WeakHashMap<K, V>()));

DELETE
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/ClassLoaderScopedWeakMapTest.groovy
-package datadog.trace.agent.tooling
-
-import spock.lang.Shared
-import spock.lang.Specification
-
-
-class ClassLoaderScopedWeakMapTest extends Specification {
-
- @Shared
- def classLoader1 = new ClassLoader() {
- @Override
- Class<?> loadClass(String name) throws ClassNotFoundException {
- return super.loadClass(name)
- }
- }
-
- @Shared
- def classLoader2 = new ClassLoader() {
- @Override
- Class<?> loadClass(String name) throws ClassNotFoundException {
- return super.loadClass(name)
- }
- }
-
- def "getOrCreate a value"() {
- setup:
- def supplier = new CounterSupplier()
- def sut = new ClassLoaderScopedWeakMap()
-
- when:
- def count = sut.getOrCreate(classLoader1, 'key', supplier)
-
- then:
- count == 1
- supplier.counter == 1
- }
-
- def "getOrCreate a value multiple times same class loader same key"() {
- setup:
- def supplier = new CounterSupplier()
- def sut = new ClassLoaderScopedWeakMap()
-
- when:
- def count1 = sut.getOrCreate(classLoader1, 'key', supplier)
- def count2 = sut.getOrCreate(classLoader1, 'key', supplier)
-
- then:
- count1 == 1
- count2 == 1
- supplier.counter == 1
- }
-
- def "getOrCreate a value multiple times same class loader different keys"() {
- setup:
- def supplier = new CounterSupplier()
- def sut = new ClassLoaderScopedWeakMap()
-
- when:
- def count1 = sut.getOrCreate(classLoader1, 'key1', supplier)
- def count2 = sut.getOrCreate(classLoader1, 'key2', supplier)
-
- then:
- count1 == 1
- count2 == 2
- supplier.counter == 2
- }
-
- def "getOrCreate a value multiple times different class loader same key"() {
- setup:
- def supplier = new CounterSupplier()
- def sut = new ClassLoaderScopedWeakMap()
-
- when:
- def count1 = sut.getOrCreate(classLoader1, 'key', supplier)
- def count2 = sut.getOrCreate(classLoader2, 'key', supplier)
-
- then:
- count1 == 1
- count2 == 2
- supplier.counter == 2
- }
-
- class CounterSupplier implements ClassLoaderScopedWeakMap.Supplier {
-
- def counter = 0
-
- @Override
- Object get() {
- counter = counter + 1
- return counter
- }
- }
-}

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
package datadog.trace.instrumentation.netty40;
-import datadog.trace.agent.tooling.ClassLoaderScopedWeakMap;
+import datadog.trace.bootstrap.WeakMap;
import datadog.trace.context.TraceScope;
import datadog.trace.instrumentation.netty40.client.HttpClientTracingHandler;
import datadog.trace.instrumentation.netty40.server.HttpServerTracingHandler;
import io.netty.util.AttributeKey;
import io.opentracing.Span;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
public class AttributeKeys {
- private static final String PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY_NAME =
- "datadog.trace.instrumentation.netty40.parent.connect.continuation";
+ private static final WeakMap<ClassLoader, Map<String, AttributeKey<?>>> map =
+ WeakMap.Implementation.DEFAULT.get();
- public static final AttributeKey<TraceScope.Continuation>
- PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY =
- (AttributeKey<TraceScope.Continuation>)
- ClassLoaderScopedWeakMap.INSTANCE.getOrCreate(
- AttributeKey.class.getClassLoader(),
- PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY_NAME,
- new ClassLoaderScopedWeakMap.Supplier() {
+ private static final WeakMap.ValueSupplier<Map<String, AttributeKey<?>>> mapSupplier =
+ new WeakMap.ValueSupplier<Map<String, AttributeKey<?>>>() {
@Override
- public Object get() {
- return new AttributeKey<>(PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY_NAME);
+ public Map<String, AttributeKey<?>> get() {
+ return new ConcurrentHashMap<>();
}
- });
+ };
- private static final String SERVER_ATTRIBUTE_KEY_NAME =
- HttpServerTracingHandler.class.getName() + ".span";
+ public static final AttributeKey<TraceScope.Continuation>
+ PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY =
+ attributeKey("datadog.trace.instrumentation.netty40.parent.connect.continuation");
public static final AttributeKey<Span> SERVER_ATTRIBUTE_KEY =
- (AttributeKey<Span>)
- ClassLoaderScopedWeakMap.INSTANCE.getOrCreate(
- AttributeKey.class.getClassLoader(),
- SERVER_ATTRIBUTE_KEY_NAME,
- new ClassLoaderScopedWeakMap.Supplier() {
- @Override
- public Object get() {
- return new AttributeKey<>(SERVER_ATTRIBUTE_KEY_NAME);
- }
- });
-
- private static final String CLIENT_ATTRIBUTE_KEY_NAME =
- HttpClientTracingHandler.class.getName() + ".span";
+ attributeKey(HttpServerTracingHandler.class.getName() + ".span");
public static final AttributeKey<Span> CLIENT_ATTRIBUTE_KEY =
- (AttributeKey<Span>)
- ClassLoaderScopedWeakMap.INSTANCE.getOrCreate(
- AttributeKey.class.getClassLoader(),
- CLIENT_ATTRIBUTE_KEY_NAME,
- new ClassLoaderScopedWeakMap.Supplier() {
- @Override
- public Object get() {
- return new AttributeKey<>(CLIENT_ATTRIBUTE_KEY_NAME);
+ attributeKey(HttpClientTracingHandler.class.getName() + ".span");
+
+ /**
+ * Generate an attribute key or reuse the one existing in the global app map. This implementation
+ * creates attributes only once even if the current class is loaded by several class loaders and
+ * prevents an issue with Apache Atlas project were this class loaded by multiple class loaders,
+ * 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);
+ if (classLoaderMap.containsKey(key)) {
+ return (AttributeKey<T>) classLoaderMap.get(key);
+ }
+
+ final AttributeKey<T> value = new AttributeKey<>(key);
+ classLoaderMap.put(key, value);
+ return value;
}
- });
}

Refactor weak hash map to be able to reuse functionalities