diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/JDBCMaps.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/JDBCMaps.java index 0e3a17234f..99b8193278 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/JDBCMaps.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/JDBCMaps.java @@ -1,10 +1,9 @@ package datadog.trace.bootstrap; +import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap; + import java.sql.Connection; import java.sql.PreparedStatement; -import java.util.Collections; -import java.util.Map; -import java.util.WeakHashMap; import lombok.Data; /** @@ -13,10 +12,8 @@ import lombok.Data; *

In the bootstrap project to ensure visibility by all classes. */ public class JDBCMaps { - public static final Map connectionInfo = - Collections.synchronizedMap(new WeakHashMap()); - public static final Map preparedStatements = - Collections.synchronizedMap(new WeakHashMap()); + public static final WeakMap connectionInfo = newWeakMap(); + public static final WeakMap preparedStatements = newWeakMap(); public static final String DB_QUERY = "DB Query"; 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 new file mode 100644 index 0000000000..d1fdc6aff8 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java @@ -0,0 +1,82 @@ +package datadog.trace.bootstrap; + +import java.util.Collections; +import java.util.Map; +import java.util.WeakHashMap; +import java.util.concurrent.atomic.AtomicReference; +import lombok.extern.slf4j.Slf4j; + +public interface WeakMap { + + int size(); + + boolean containsKey(K target); + + V get(K key); + + void put(K key, V value); + + class Provider { + private static final AtomicReference provider = + new AtomicReference<>(Supplier.DEFAULT); + + public static void registerIfAbsent(final Supplier provider) { + if (provider != null && provider != Supplier.DEFAULT) { + Provider.provider.compareAndSet(Supplier.DEFAULT, provider); + } + } + + public static WeakMap newWeakMap() { + return provider.get().get(); + } + } + + interface Supplier { + WeakMap get(); + + Supplier DEFAULT = new Default(); + + @Slf4j + class Default implements Supplier { + + @Override + public WeakMap get() { + log.warn("WeakMap.Supplier not registered. Returning a synchronized WeakHashMap."); + return new MapAdapter<>(Collections.synchronizedMap(new WeakHashMap())); + } + } + } + + class MapAdapter implements WeakMap { + private final Map map; + + public MapAdapter(final Map map) { + this.map = map; + } + + @Override + public int size() { + return map.size(); + } + + @Override + public boolean containsKey(final K key) { + return map.containsKey(key); + } + + @Override + public V get(final K key) { + return map.get(key); + } + + @Override + public void put(final K key, final V value) { + map.put(key, value); + } + + @Override + public String toString() { + return map.toString(); + } + } +} diff --git a/dd-java-agent/agent-tooling/agent-tooling.gradle b/dd-java-agent/agent-tooling/agent-tooling.gradle index b44748192e..0c28a3bd4f 100644 --- a/dd-java-agent/agent-tooling/agent-tooling.gradle +++ b/dd-java-agent/agent-tooling/agent-tooling.gradle @@ -9,6 +9,7 @@ configurations { dependencies { compile project(':dd-java-agent:agent-bootstrap') + compile group: 'com.blogspot.mydailyjava', name: 'weak-lock-free', version: '0.13' compile deps.bytebuddy compile deps.bytebuddyagent annotationProcessor deps.autoservice @@ -17,6 +18,7 @@ dependencies { compileOnly project(':dd-trace-ot') testCompile deps.opentracing + testCompile project(':dd-java-agent:testing') instrumentationMuzzle sourceSets.main.output instrumentationMuzzle configurations.compile diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java index 90d4cd1653..c8cce78f1e 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java @@ -8,6 +8,7 @@ import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith; import static net.bytebuddy.matcher.ElementMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.not; +import datadog.trace.bootstrap.WeakMap; import java.lang.instrument.Instrumentation; import java.util.ServiceLoader; import lombok.extern.slf4j.Slf4j; @@ -38,6 +39,8 @@ public class AgentInstaller { public static ResettableClassFileTransformer installBytebuddyAgent( final Instrumentation inst, final AgentBuilder.Listener... listeners) { INSTRUMENTATION = inst; + registerWeakMapProvider(); + AgentBuilder agentBuilder = new AgentBuilder.Default() .disableClassFormatChanges() @@ -86,6 +89,12 @@ public class AgentInstaller { return agentBuilder.installOn(inst); } + private static void registerWeakMapProvider() { + WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent()); + // WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent.Inline()); + // WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.Guava()); + } + @Slf4j static class LoggingListener implements AgentBuilder.Listener { 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 1dc632faa3..068916e882 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 @@ -1,13 +1,14 @@ package datadog.trace.agent.tooling; +import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap; + import datadog.trace.bootstrap.DatadogClassLoader; import datadog.trace.bootstrap.PatchLogger; +import datadog.trace.bootstrap.WeakMap; import io.opentracing.util.GlobalTracer; import java.util.Collections; import java.util.HashSet; -import java.util.Map; import java.util.Set; -import java.util.WeakHashMap; import lombok.extern.slf4j.Slf4j; import net.bytebuddy.matcher.ElementMatcher; @@ -43,8 +44,7 @@ public class ClassLoaderMatcher { extends ElementMatcher.Junction.AbstractBase { public static final SkipClassLoaderMatcher INSTANCE = new SkipClassLoaderMatcher(); /* Cache of classloader-instance -> (true|false). True = skip instrumentation. False = safe to instrument. */ - private static final Map SKIP_CACHE = - Collections.synchronizedMap(new WeakHashMap()); + private static final WeakMap SKIP_CACHE = newWeakMap(); private static final Set CLASSLOADER_CLASSES_TO_SKIP; static { @@ -131,8 +131,7 @@ public class ClassLoaderMatcher { public static class ClassLoaderHasClassMatcher extends ElementMatcher.Junction.AbstractBase { - private final Map cache = - Collections.synchronizedMap(new WeakHashMap()); + private final WeakMap cache = newWeakMap(); private final String[] names; @@ -164,8 +163,7 @@ public class ClassLoaderMatcher { public static class ClassLoaderHasClassWithFieldMatcher extends ElementMatcher.Junction.AbstractBase { - private final Map cache = - Collections.synchronizedMap(new WeakHashMap()); + private final WeakMap cache = newWeakMap(); private final String className; private final String fieldName; @@ -203,8 +201,7 @@ public class ClassLoaderMatcher { public static class ClassLoaderHasClassWithMethodMatcher extends ElementMatcher.Junction.AbstractBase { - private final Map cache = - Collections.synchronizedMap(new WeakHashMap()); + private final WeakMap cache = newWeakMap(); private final String className; private final String methodName; diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Utils.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Utils.java index f2ad09cd6b..6ce61707c4 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Utils.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Utils.java @@ -22,6 +22,8 @@ public class Utils { "com.google.auto", "com.google.common", "com.google.thirdparty.publicsuffix", + // WeakConcurrentMap + "com.blogspot.mydailyjava.weaklockfree", // bytebuddy "net.bytebuddy", // OT contribs for dd trace resolver @@ -34,13 +36,13 @@ public class Utils { private static Method findLoadedClassMethod = null; - private static BootstrapClassLoaderProxy unitTestBootstrapProxy = + private static final BootstrapClassLoaderProxy unitTestBootstrapProxy = new BootstrapClassLoaderProxy(new URL[0], null); static { try { findLoadedClassMethod = ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class); - } catch (NoSuchMethodException | SecurityException e) { + } catch (final NoSuchMethodException | SecurityException e) { throw new IllegalStateException(e); } } 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 new file mode 100644 index 0000000000..9f948e2228 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java @@ -0,0 +1,158 @@ +package datadog.trace.agent.tooling; + +import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.MapMaker; +import datadog.trace.bootstrap.WeakMap; +import java.lang.ref.WeakReference; +import java.util.Iterator; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; + +class WeakMapSuppliers { + // Comparison with using WeakConcurrentMap vs Guava's implementation: + // Cleaning: + // * `WeakConcurrentMap`: centralized but we have to maintain out own code and thread for it + // * `Guava`: inline on application's thread, with constant max delay + // Jar Size: + // * `WeakConcurrentMap`: small + // * `Guava`: large, but we may use other features, like immutable collections - and we already + // ship Guava as part of distribution now, so using Guava for this doesn’t increase size. + // Must go on bootstrap classpath: + // * `WeakConcurrentMap`: version conflict is unlikely, so we can directly inject for now + // * `Guava`: need to implement shadow copy (might eventually be necessary for other dependencies) + // Used by other javaagents for similar purposes: + // * `WeakConcurrentMap`: anecdotally used by other agents + // * `Guava`: specifically agent use is unknown at the moment, but Guava is a well known library + // backed by big company with many-many users + + /** + * Provides instances of {@link WeakConcurrentMap} and retains weak reference to them to allow a + * single thread to clean void weak references out for all instances. Cleaning is done every + * second. + */ + static class WeakConcurrent implements WeakMap.Supplier { + + private static final long SHUTDOWN_WAIT_SECONDS = 5; + + @VisibleForTesting static final long CLEAN_FREQUENCY_SECONDS = 1; + + private static final ThreadFactory THREAD_FACTORY = + new ThreadFactory() { + @Override + public Thread newThread(final Runnable r) { + final Thread thread = new Thread(r, "dd-weak-ref-cleaner"); + thread.setDaemon(true); + thread.setPriority(Thread.MIN_PRIORITY); + return thread; + } + }; + + private final ScheduledExecutorService cleanerExecutorService; + + private final Queue> suppliedMaps = + new ConcurrentLinkedQueue<>(); + + private final Runnable runnable = + new Runnable() { + @Override + public void run() { + for (final Iterator> iterator = + suppliedMaps.iterator(); + iterator.hasNext(); ) { + final WeakConcurrentMap map = iterator.next().get(); + if (map == null) { + iterator.remove(); + } else { + map.expungeStaleEntries(); + } + } + } + }; + + WeakConcurrent() { + cleanerExecutorService = Executors.newScheduledThreadPool(1, THREAD_FACTORY); + cleanerExecutorService.scheduleAtFixedRate( + runnable, CLEAN_FREQUENCY_SECONDS, CLEAN_FREQUENCY_SECONDS, TimeUnit.SECONDS); + + try { + Runtime.getRuntime() + .addShutdownHook( + new Thread() { + @Override + public void run() { + try { + cleanerExecutorService.shutdownNow(); + cleanerExecutorService.awaitTermination( + SHUTDOWN_WAIT_SECONDS, TimeUnit.SECONDS); + } catch (final InterruptedException e) { + // Don't bother waiting then... + } + } + }); + } catch (final IllegalStateException ex) { + // The JVM is already shutting down. + } + } + + @Override + public WeakMap get() { + final WeakConcurrentMap map = new WeakConcurrentMap<>(false); + suppliedMaps.add(new WeakReference(map)); + return new Adapter<>(map); + } + + private static class Adapter implements WeakMap { + private final WeakConcurrentMap map; + + private Adapter(final WeakConcurrentMap map) { + this.map = map; + } + + @Override + public int size() { + return map.approximateSize(); + } + + @Override + public boolean containsKey(final K key) { + return map.containsKey(key); + } + + @Override + public V get(final K key) { + return map.get(key); + } + + @Override + public void put(final K key, final V value) { + map.put(key, value); + } + } + + static class Inline implements WeakMap.Supplier { + + @Override + public WeakMap get() { + return new Adapter<>(new WeakConcurrentMap.WithInlinedExpunction()); + } + } + } + + static class Guava implements WeakMap.Supplier { + + @Override + public WeakMap get() { + return new WeakMap.MapAdapter<>(new MapMaker().weakKeys().makeMap()); + } + + public WeakMap get(int concurrencyLevel) { + return new WeakMap.MapAdapter<>( + new MapMaker().concurrencyLevel(concurrencyLevel).weakKeys().makeMap()); + } + } +} 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 3a6a25bf5f..e5ed74906d 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 @@ -1,23 +1,21 @@ package datadog.trace.agent.tooling.muzzle; +import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap; import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER; import datadog.trace.agent.tooling.Utils; +import datadog.trace.bootstrap.WeakMap; import java.util.ArrayList; import java.util.Arrays; -import java.util.Collections; import java.util.HashSet; import java.util.List; -import java.util.Map; import java.util.Set; -import java.util.WeakHashMap; import lombok.extern.slf4j.Slf4j; /** Matches a set of references against a classloader. */ @Slf4j public class ReferenceMatcher { - private final Map> mismatchCache = - Collections.synchronizedMap(new WeakHashMap>()); + private final WeakMap> mismatchCache = newWeakMap(); private final Reference[] references; private final Set helperClassNames; diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakConcurrentSupplierTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakConcurrentSupplierTest.groovy new file mode 100644 index 0000000000..a8245c5570 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakConcurrentSupplierTest.groovy @@ -0,0 +1,106 @@ +package datadog.trace.agent.tooling + +import datadog.trace.agent.test.TestUtils +import datadog.trace.bootstrap.WeakMap +import spock.lang.Shared +import spock.lang.Specification + +import java.lang.ref.WeakReference +import java.util.concurrent.TimeUnit + +class WeakConcurrentSupplierTest extends Specification { + @Shared + def weakConcurrentSupplier = new WeakMapSuppliers.WeakConcurrent() + @Shared + def weakInlineSupplier = new WeakMapSuppliers.WeakConcurrent.Inline() + @Shared + def guavaSupplier = new WeakMapSuppliers.Guava() + + def "Calling newWeakMap on #name creates independent maps"() { + setup: + WeakMap.Provider.provider.set(supplier) + def key = new Object() + def map1 = WeakMap.Provider.newWeakMap() + def map2 = WeakMap.Provider.newWeakMap() + + when: + map1.put(key, "value1") + map2.put(key, "value2") + + then: + map1.get(key) == "value1" + map2.get(key) == "value2" + + where: + name | supplier + "WeakConcurrent" | weakConcurrentSupplier + "WeakInline" | weakInlineSupplier + "Guava" | guavaSupplier + } + + def "Unreferenced map get cleaned up on #name"() { + setup: + WeakMap.Provider.provider.set(supplier) + def map = WeakMap.Provider.newWeakMap() + def ref = new WeakReference(map) + + when: + map = null + TestUtils.awaitGC() + + then: + ref.get() == null + + where: + name | supplier + "WeakConcurrent" | weakConcurrentSupplier + "WeakInline" | weakInlineSupplier + "Guava" | guavaSupplier + } + + def "Unreferenced keys get cleaned up on #name"() { + setup: + def key = new Object() + map.put(key, "value") + TestUtils.awaitGC() + + expect: + map.size() == 1 + + when: + key = null + TestUtils.awaitGC() + + if (name == "WeakConcurrent") { + // Sleep enough time for cleanup thread to get scheduled. + // But on a very slow box (or high load) scheduling may not be exactly predictable + // so we try a few times. + int count = 0 + while (map.size() != 0 && count < 10) { + Thread.sleep(TimeUnit.SECONDS.toMillis(WeakMapSuppliers.WeakConcurrent.CLEAN_FREQUENCY_SECONDS)) + count++ + } + } + + // Hit map a few times to trigger unreferenced entries cleanup. + // Exact number of times that we need to hit map is implementation dependent. + // For Guava it i specified in + // com.google.common.collect.MapMakerInternalMap.DRAIN_THRESHOLD = 0x3F + if (name == "Guava" || name == "WeakInline") { + for (int i = 0; i <= 0x3F; i++) { + map.get("test") + } + } + + then: + map.size() == 0 + + where: + name | map + "WeakConcurrent" | weakConcurrentSupplier.get() + "WeakInline" | weakInlineSupplier.get() + // Guava's cleanup process depends on concurrency level, + // and in order to be able to test it we need to set concurrency to 1 + "Guava" | guavaSupplier.get(1) + } +} diff --git a/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java b/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java index 8657283799..2061aff504 100644 --- a/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java +++ b/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.http_url_connection; import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType; +import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap; import static io.opentracing.log.Fields.ERROR_OBJECT; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; @@ -12,6 +13,7 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.api.DDSpanTypes; import datadog.trace.api.DDTags; import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import datadog.trace.bootstrap.WeakMap; import io.opentracing.Scope; import io.opentracing.Span; import io.opentracing.Tracer; @@ -21,7 +23,6 @@ import java.net.HttpURLConnection; import java.net.URL; import java.util.Collections; import java.util.Map; -import java.util.WeakHashMap; import javax.net.ssl.HttpsURLConnection; import net.bytebuddy.asm.Advice; import net.bytebuddy.matcher.ElementMatcher; @@ -175,8 +176,7 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default { } public static class HttpURLState { - private static final Map STATE_MAP = - Collections.synchronizedMap(new WeakHashMap()); + private static final WeakMap STATE_MAP = newWeakMap(); public static HttpURLState get(final HttpURLConnection connection) { HttpURLState state = STATE_MAP.get(connection); diff --git a/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/TestUtils.java b/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/TestUtils.java index 8db4c4a094..c4d4cd9356 100644 --- a/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/TestUtils.java +++ b/dd-java-agent/testing/src/main/java/datadog/trace/agent/test/TestUtils.java @@ -13,6 +13,7 @@ import java.io.File; import java.io.FileOutputStream; import java.io.IOException; import java.io.InputStream; +import java.lang.ref.WeakReference; import java.lang.reflect.Field; import java.net.ServerSocket; import java.net.URL; @@ -204,4 +205,13 @@ public class TestUtils { return -1; } } + + public static void awaitGC() { + Object obj = new Object(); + final WeakReference ref = new WeakReference<>(obj); + obj = null; + while (ref.get() != null) { + System.gc(); + } + } } diff --git a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java index badd45c715..e3a43fc2a8 100644 --- a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java +++ b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java @@ -29,6 +29,7 @@ import io.opentracing.ScopeManager; import io.opentracing.Span; import io.opentracing.SpanContext; import io.opentracing.propagation.Format; +import java.io.Closeable; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; @@ -48,7 +49,7 @@ import lombok.extern.slf4j.Slf4j; /** DDTracer makes it easy to send traces and span to DD using the OpenTracing API. */ @Slf4j -public class DDTracer implements io.opentracing.Tracer { +public class DDTracer implements io.opentracing.Tracer, Closeable { public static final String UNASSIGNED_DEFAULT_SERVICE_NAME = "unnamed-java-app"; @@ -302,6 +303,7 @@ public class DDTracer implements io.opentracing.Tracer { } } + @Override public void close() { PendingTrace.close(); writer.close();