Extract logic out for `Cleaner` class.
MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
@@ -29,14 +29,15 @@ import net.bytebuddy.utility.JavaModule;
@Slf4j
public class AgentInstaller {
private static final Map<String, Runnable> classLoadCallbacks = new ConcurrentHashMap<>();
+ public static final Cleaner CLEANER = new Cleaner();
static {
// WeakMap is used by other classes below, so we need to register the provider first.
- registerWeakMapProvider();
+ registerWeakMapProvider(CLEANER);
}
public static final DDLocationStrategy LOCATION_STRATEGY = new DDLocationStrategy();
- public static final AgentBuilder.PoolStrategy POOL_STRATEGY = new DDCachingPoolStrategy();
+ public static final AgentBuilder.PoolStrategy POOL_STRATEGY = new DDCachingPoolStrategy(CLEANER);
private static volatile Instrumentation INSTRUMENTATION;
public static Instrumentation getInstrumentation() {
@@ -152,6 +153,7 @@ public class AgentInstaller {
}
log.debug("Installed {} instrumenter(s)", numInstrumenters);
+ CLEANER.start();
return agentBuilder.installOn(inst);
}
@@ -171,13 +173,13 @@ public class AgentInstaller {
return matcher;
}
- private static void registerWeakMapProvider() {
+ private static void registerWeakMapProvider(final Cleaner cleaner) {
if (!WeakMap.Provider.isProviderRegistered()) {
- WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent());
- }
+ WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent(cleaner));
// WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent.Inline());
// WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.Guava());
}
+ }
@Slf4j
static class RedefinitionLoggingListener implements AgentBuilder.RedefinitionStrategy.Listener {

ADD
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Cleaner.java
+package datadog.trace.agent.tooling;
+
+import java.lang.ref.WeakReference;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+class Cleaner {
+ // TODO: add unit tests for this class.
+
+ private static final long SHUTDOWN_WAIT_SECONDS = 5;
+
+ private static final ThreadFactory THREAD_FACTORY =
+ new ThreadFactory() {
+ @Override
+ public Thread newThread(final Runnable r) {
+ final Thread thread = new Thread(r, "dd-cleaner");
+ thread.setDaemon(true);
+ thread.setPriority(Thread.MIN_PRIORITY);
+ return thread;
+ }
+ };
+
+ private volatile ScheduledThreadPoolExecutor cleanerService = null;
+ private volatile Thread shutdownCallback = null;
+
+ <T> void scheduleCleaning(
+ final T target, final Adapter<T> adapter, final long frequency, final TimeUnit unit) {
+ final CleanupRunnable<T> command = new CleanupRunnable<>(target, adapter);
+ if (cleanerService == null) {
+ log.warn("Cleaning scheduled before starting cleaner. Target won't be cleaned {}", target);
+ } else {
+ command.setFuture(cleanerService.scheduleAtFixedRate(command, frequency, frequency, unit));
+ }
+ }
+
+ public void start() {
+ if (cleanerService == null) {
+ synchronized (this) {
+ if (cleanerService == null) {
+ cleanerService = new ScheduledThreadPoolExecutor(1, THREAD_FACTORY);
+ cleanerService.setRemoveOnCancelPolicy(true);
+ shutdownCallback = new ShutdownCallback(cleanerService);
+ try {
+ Runtime.getRuntime().addShutdownHook(shutdownCallback);
+ } catch (final IllegalStateException ex) {
+ // The JVM is already shutting down.
+ }
+ }
+ }
+ }
+ }
+
+ public void stop() {
+ if (cleanerService == null) {
+ synchronized (this) {
+ if (cleanerService == null) {
+ cleanerService.shutdown();
+ Runtime.getRuntime().removeShutdownHook(shutdownCallback);
+ cleanerService = null;
+ shutdownCallback = null;
+ }
+ }
+ }
+ }
+
+ @Override
+ public void finalize() {
+ // Do we really want this?
+ stop();
+ }
+
+ public interface Adapter<T> {
+ void clean(T target);
+ }
+
+ private static class CleanupRunnable<T> implements Runnable {
+ private final WeakReference<T> target;
+ private final Adapter<T> adapter;
+ private ScheduledFuture<?> future = null;
+
+ private CleanupRunnable(final T target, final Adapter<T> adapter) {
+ this.target = new WeakReference<>(target);
+ this.adapter = adapter;
+ }
+
+ @Override
+ public void run() {
+ final T t = target.get();
+ if (t != null) {
+ adapter.clean(t);
+ } else if (future != null) {
+ future.cancel(false);
+ }
+ }
+
+ public void setFuture(final ScheduledFuture<?> future) {
+ this.future = future;
+ }
+ }
+
+ private static final class ShutdownCallback extends Thread {
+
+ private final ScheduledExecutorService executorService;
+
+ public ShutdownCallback(final ScheduledExecutorService executorService) {
+ this.executorService = executorService;
+ }
+
+ @Override
+ public void run() {
+ try {
+ executorService.shutdownNow();
+ executorService.awaitTermination(SHUTDOWN_WAIT_SECONDS, TimeUnit.SECONDS);
+ } catch (final InterruptedException e) {
+ // Don't bother waiting then...
+ }
+ }
+ }
+}

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
@@ -8,9 +8,6 @@ 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.ScheduledFuture;
-import java.util.concurrent.ScheduledThreadPoolExecutor;
-import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.ClassFileLocator;
@@ -33,6 +30,11 @@ import net.bytebuddy.pool.TypePool;
public class DDCachingPoolStrategy implements PoolStrategy {
private static final WeakMap<ClassLoader, TypePool.CacheProvider> typePoolCache =
WeakMap.Provider.newWeakMap();
+ private final Cleaner cleaner;
+
+ public DDCachingPoolStrategy(final Cleaner cleaner) {
+ this.cleaner = cleaner;
+ }
@Override
public TypePool typePool(final ClassFileLocator classFileLocator, final ClassLoader classLoader) {
@@ -43,7 +45,7 @@ public class DDCachingPoolStrategy implements PoolStrategy {
synchronized (key) {
cache = typePoolCache.get(key);
if (null == cache) {
- cache = EvictingCacheProvider.withObjectType(10, TimeUnit.SECONDS);
+ cache = EvictingCacheProvider.withObjectType(cleaner, 10, TimeUnit.SECONDS);
typePoolCache.put(key, cache);
}
}
@@ -53,30 +55,13 @@ public class DDCachingPoolStrategy implements PoolStrategy {
}
private static class EvictingCacheProvider implements TypePool.CacheProvider {
- private static final ScheduledThreadPoolExecutor cleanerExecutorService =
- new ScheduledThreadPoolExecutor(
- 1,
- new ThreadFactory() {
- @Override
- public Thread newThread(final Runnable r) {
- final Thread thread = new Thread(r, "dd-type-cache-cleaner");
- thread.setDaemon(true);
- thread.setPriority(Thread.MIN_PRIORITY);
- return thread;
- }
- });
-
- static {
- cleanerExecutorService.setRemoveOnCancelPolicy(true);
- }
/** A map containing all cached resolutions by their names. */
private final Cache<String, TypePool.Resolution> cache;
- private final ScheduledFuture<?> cancelFuture;
-
/** Creates a new simple cache. */
- private EvictingCacheProvider(final long expireDuration, final TimeUnit unit) {
+ private EvictingCacheProvider(
+ final Cleaner cleaner, final long expireDuration, final TimeUnit unit) {
cache =
CacheBuilder.newBuilder()
.initialCapacity(500)
@@ -89,15 +74,14 @@ public class DDCachingPoolStrategy implements PoolStrategy {
* We want to ensure this happens more regularly, so we schedule a thread to do run cleanup manually.
*/
final long cleanFrequency = expireDuration / 2;
- cancelFuture =
- cleanerExecutorService.scheduleAtFixedRate(
- new CacheCleaner(), cleanFrequency, cleanFrequency, unit);
+ cleaner.scheduleCleaning(cache, CacheCleaner.CLEANER, cleanFrequency, unit);
}
private static EvictingCacheProvider withObjectType(
- final long expireDuration, final TimeUnit unit) {
+ final Cleaner cleaner, final long expireDuration, final TimeUnit unit) {
assert expireDuration % 2 == 0 : "expireDuration must be even";
- final EvictingCacheProvider cacheProvider = new EvictingCacheProvider(expireDuration, unit);
+ final EvictingCacheProvider cacheProvider =
+ new EvictingCacheProvider(cleaner, expireDuration, unit);
cacheProvider.register(
Object.class.getName(), new TypePool.Resolution.Simple(TypeDescription.OBJECT));
return cacheProvider;
@@ -126,15 +110,12 @@ public class DDCachingPoolStrategy implements PoolStrategy {
return cache.size();
}
- @Override
- public void finalize() {
- cancelFuture.cancel(true);
- }
+ private static class CacheCleaner implements Cleaner.Adapter<Cache> {
+ private static final CacheCleaner CLEANER = new CacheCleaner();
- private class CacheCleaner implements Runnable {
@Override
- public void run() {
- cache.cleanUp();
+ public void clean(final Cache target) {
+ target.cleanUp();
}
}

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
@@ -4,15 +4,7 @@ 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;
-import java.util.concurrent.atomic.AtomicBoolean;
class WeakMapSuppliers {
// Comparison with using WeakConcurrentMap vs Guava's implementation:
@@ -38,109 +30,26 @@ class WeakMapSuppliers {
*/
static class WeakConcurrent implements WeakMap.Implementation {
- private static final long SHUTDOWN_WAIT_SECONDS = 5;
-
@VisibleForTesting static final long CLEAN_FREQUENCY_SECONDS = 1;
+ private final Cleaner cleaner;
- 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 Thread shutdownCallback;
- private final Queue<WeakReference<WeakConcurrentMap>> suppliedMaps =
- new ConcurrentLinkedQueue<>();
-
- private final AtomicBoolean finalized = new AtomicBoolean(false);
-
- WeakConcurrent() {
- cleanerExecutorService = Executors.newScheduledThreadPool(1, THREAD_FACTORY);
- shutdownCallback = new ShutdownCallback(cleanerExecutorService);
-
- cleanerExecutorService.scheduleAtFixedRate(
- new CleanupRunnable(cleanerExecutorService, shutdownCallback, suppliedMaps, finalized),
- CLEAN_FREQUENCY_SECONDS,
- CLEAN_FREQUENCY_SECONDS,
- TimeUnit.SECONDS);
-
- try {
- Runtime.getRuntime().addShutdownHook(shutdownCallback);
- } catch (final IllegalStateException ex) {
- // The JVM is already shutting down.
- }
- }
-
- @Override
- public void finalize() {
- finalized.set(true);
+ WeakConcurrent(final Cleaner cleaner) {
+ this.cleaner = cleaner;
}
@Override
public <K, V> WeakMap<K, V> get() {
final WeakConcurrentMap<K, V> map = new WeakConcurrentMap<>(false);
- suppliedMaps.add(new WeakReference<WeakConcurrentMap>(map));
- return new Adapter<>(map);
+ cleaner.scheduleCleaning(map, MapCleaner.CLEANER, CLEAN_FREQUENCY_SECONDS, TimeUnit.SECONDS);
+ return new Adapter(map);
}
- private static class CleanupRunnable implements Runnable {
-
- private final ScheduledExecutorService executorService;
- private final Thread shutdownCallback;
- private final Queue<WeakReference<WeakConcurrentMap>> suppliedMaps;
- private final AtomicBoolean finalized;
-
- public CleanupRunnable(
- final ScheduledExecutorService executorService,
- final Thread shutdownCallback,
- final Queue<WeakReference<WeakConcurrentMap>> suppliedMaps,
- final AtomicBoolean finalized) {
- this.executorService = executorService;
- this.shutdownCallback = shutdownCallback;
- this.suppliedMaps = suppliedMaps;
- this.finalized = finalized;
- }
+ private static class MapCleaner implements Cleaner.Adapter<WeakConcurrentMap> {
+ private static final MapCleaner CLEANER = new MapCleaner();
@Override
- public void run() {
- for (final Iterator<WeakReference<WeakConcurrentMap>> iterator = suppliedMaps.iterator();
- iterator.hasNext(); ) {
- final WeakConcurrentMap map = iterator.next().get();
- if (map == null) {
- iterator.remove();
- } else {
- map.expungeStaleEntries();
- }
- }
- if (finalized.get() && suppliedMaps.isEmpty()) {
- executorService.shutdown();
- Runtime.getRuntime().removeShutdownHook(shutdownCallback);
- }
- }
- }
-
- private static final class ShutdownCallback extends Thread {
-
- private final ScheduledExecutorService executorService;
-
- public ShutdownCallback(final ScheduledExecutorService executorService) {
- this.executorService = executorService;
- }
-
- @Override
- public void run() {
- try {
- executorService.shutdownNow();
- executorService.awaitTermination(SHUTDOWN_WAIT_SECONDS, TimeUnit.SECONDS);
- } catch (final InterruptedException e) {
- // Don't bother waiting then...
- }
+ public void clean(final WeakConcurrentMap target) {
+ target.expungeStaleEntries();
}
}
@@ -177,7 +86,7 @@ class WeakMapSuppliers {
}
@Override
- public V getOrCreate(K key, ValueSupplier<V> supplier) {
+ public V getOrCreate(final K key, final ValueSupplier<V> supplier) {
if (!map.containsKey(key)) {
synchronized (this) {
if (!map.containsKey(key)) {
@@ -194,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/test/groovy/datadog/trace/agent/tooling/EvictingCacheProviderTest.groovy
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/EvictingCacheProviderTest.groovy
@@ -4,17 +4,22 @@ import datadog.trace.util.gc.GCUtils
import net.bytebuddy.description.type.TypeDescription
import net.bytebuddy.pool.TypePool
import spock.lang.Specification
-import spock.lang.Timeout
import java.lang.ref.WeakReference
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicReference
-@Timeout(5)
+import static datadog.trace.agent.tooling.AgentInstaller.CLEANER
+
+//@Timeout(5)
class EvictingCacheProviderTest extends Specification {
+ static {
+ CLEANER.start()
+ }
+
def "test provider"() {
setup:
- def provider = new DDCachingPoolStrategy.EvictingCacheProvider(2, TimeUnit.MINUTES)
+ def provider = new DDCachingPoolStrategy.EvictingCacheProvider(CLEANER, 2, TimeUnit.MINUTES)
expect:
provider.size() == 0
@@ -40,16 +45,28 @@ class EvictingCacheProviderTest extends Specification {
def "test timeout eviction"() {
setup:
- def provider = new DDCachingPoolStrategy.EvictingCacheProvider(10, TimeUnit.MILLISECONDS)
+ def provider = new DDCachingPoolStrategy.EvictingCacheProvider(CLEANER, timeout, TimeUnit.MILLISECONDS)
def resolutionRef = new AtomicReference<TypePool.Resolution>(new TypePool.Resolution.Simple(TypeDescription.VOID))
def weakRef = new WeakReference(resolutionRef.get())
when:
+ def lastAccess = System.nanoTime()
provider.register(className, resolutionRef.get())
then:
- provider.size() == 1
- Thread.sleep(10)
+ // Ensure continued access prevents expiration.
+ for (int i = 0; i < timeout + 10; i++) {
+ assert TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - lastAccess) < timeout: "test took too long on " + i
+ assert provider.find(className) != null
+ assert provider.size() == 1
+ lastAccess = System.nanoTime()
+ Thread.sleep(1)
+ }
+
+ when:
+ Thread.sleep(timeout)
+
+ then:
provider.find(className) == null
when:
@@ -58,16 +75,18 @@ class EvictingCacheProviderTest extends Specification {
GCUtils.awaitGC(weakRef)
then:
+ // Verify properly GC'd
provider.size() == 0
weakRef.get() == null
where:
className = "SomeClass"
+ timeout = 500 // Takes about 50 ms locally, adding an order of magnitude for CI.
}
def "test size limit"() {
setup:
- def provider = new DDCachingPoolStrategy.EvictingCacheProvider(2, TimeUnit.MINUTES)
+ def provider = new DDCachingPoolStrategy.EvictingCacheProvider(CLEANER, 2, TimeUnit.MINUTES)
def typeDef = new TypePool.Resolution.Simple(TypeDescription.VOID)
for (int i = 0; i < 20000; i++) {
provider.register("ClassName$i", typeDef)

MODIFY
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakConcurrentSupplierTest.groovy
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakConcurrentSupplierTest.groovy
@@ -9,11 +9,16 @@ import spock.lang.Specification
import java.lang.ref.WeakReference
import java.util.concurrent.TimeUnit
+import static datadog.trace.agent.tooling.AgentInstaller.CLEANER
+
@Retry
// These tests fail sometimes in CI.
class WeakConcurrentSupplierTest extends Specification {
+ static {
+ CLEANER.start()
+ }
@Shared
- def weakConcurrentSupplier = new WeakMapSuppliers.WeakConcurrent()
+ def weakConcurrentSupplier = new WeakMapSuppliers.WeakConcurrent(CLEANER)
@Shared
def weakInlineSupplier = new WeakMapSuppliers.WeakConcurrent.Inline()
@Shared
@@ -57,7 +62,7 @@ class WeakConcurrentSupplierTest extends Specification {
where:
name | supplierSupplier
- "WeakConcurrent" | { -> new WeakMapSuppliers.WeakConcurrent() }
+ "WeakConcurrent" | { -> new WeakMapSuppliers.WeakConcurrent(CLEANER) }
"WeakInline" | { -> new WeakMapSuppliers.WeakConcurrent.Inline() }
"Guava" | { -> new WeakMapSuppliers.Guava() }
}

Extract logic out for `Cleaner` class.