Reorganize WeakConcurrent implementation to allow cleanup thread GC when map is out of scope
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
@@ -12,6 +12,7 @@ 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:
@@ -57,12 +58,53 @@ class WeakMapSuppliers {
private final Queue<WeakReference<WeakConcurrentMap>> suppliedMaps =
new ConcurrentLinkedQueue<>();
- private final Runnable runnable =
- new Runnable() {
+ private final AtomicBoolean finalized = new AtomicBoolean(false);
+
+ WeakConcurrent() {
+ cleanerExecutorService = Executors.newScheduledThreadPool(1, THREAD_FACTORY);
+ cleanerExecutorService.scheduleAtFixedRate(
+ new CleanupRunnable(cleanerExecutorService, suppliedMaps, finalized),
+ CLEAN_FREQUENCY_SECONDS,
+ CLEAN_FREQUENCY_SECONDS,
+ TimeUnit.SECONDS);
+
+ try {
+ Runtime.getRuntime().addShutdownHook(new ShutdownCallback(cleanerExecutorService));
+ } catch (final IllegalStateException ex) {
+ // The JVM is already shutting down.
+ }
+ }
+
+ @Override
+ public void finalize() {
+ finalized.set(true);
+ }
+
+ @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);
+ }
+
+ private static class CleanupRunnable implements Runnable {
+
+ private final ScheduledExecutorService executorService;
+ private final Queue<WeakReference<WeakConcurrentMap>> suppliedMaps;
+ private final AtomicBoolean finalized;
+
+ public CleanupRunnable(
+ final ScheduledExecutorService executorService,
+ final Queue<WeakReference<WeakConcurrentMap>> suppliedMaps,
+ final AtomicBoolean finalized) {
+ this.executorService = executorService;
+ this.suppliedMaps = suppliedMaps;
+ this.finalized = finalized;
+ }
+
@Override
public void run() {
- for (final Iterator<WeakReference<WeakConcurrentMap>> iterator =
- suppliedMaps.iterator();
+ for (final Iterator<WeakReference<WeakConcurrentMap>> iterator = suppliedMaps.iterator();
iterator.hasNext(); ) {
final WeakConcurrentMap map = iterator.next().get();
if (map == null) {
@@ -71,39 +113,29 @@ class WeakMapSuppliers {
map.expungeStaleEntries();
}
}
+ if (finalized.get() && suppliedMaps.isEmpty()) {
+ executorService.shutdown();
+ }
+ }
}
- };
- WeakConcurrent() {
- cleanerExecutorService = Executors.newScheduledThreadPool(1, THREAD_FACTORY);
- cleanerExecutorService.scheduleAtFixedRate(
- runnable, CLEAN_FREQUENCY_SECONDS, CLEAN_FREQUENCY_SECONDS, TimeUnit.SECONDS);
+ private static final class ShutdownCallback extends Thread {
+
+ private final ScheduledExecutorService executorService;
+
+ public ShutdownCallback(final ScheduledExecutorService executorService) {
+ this.executorService = executorService;
+ }
- try {
- Runtime.getRuntime()
- .addShutdownHook(
- new Thread() {
@Override
public void run() {
try {
- cleanerExecutorService.shutdownNow();
- cleanerExecutorService.awaitTermination(
- SHUTDOWN_WAIT_SECONDS, TimeUnit.SECONDS);
+ executorService.shutdownNow();
+ executorService.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 <K, V> WeakMap<K, V> get() {
- final WeakConcurrentMap<K, V> map = new WeakConcurrentMap<>(false);
- suppliedMaps.add(new WeakReference<WeakConcurrentMap>(map));
- return new Adapter<>(map);
}
private static class Adapter<K, V> implements WeakMap<K, V> {
@@ -150,7 +182,7 @@ class WeakMapSuppliers {
return new WeakMap.MapAdapter<>(new MapMaker().weakKeys().<K, V>makeMap());
}
- public <K, V> WeakMap<K, V> get(int concurrencyLevel) {
+ public <K, V> WeakMap<K, V> get(final int concurrencyLevel) {
return new WeakMap.MapAdapter<>(
new MapMaker().concurrencyLevel(concurrencyLevel).weakKeys().<K, V>makeMap());
}

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
@@ -38,14 +38,35 @@ class WeakConcurrentSupplierTest extends Specification {
"Guava" | guavaSupplier
}
- def "Unreferenced map get cleaned up on #name"() {
+ def "Unreferenced supplier gets cleaned up on #name"() {
+ setup:
+ // Note: we use 'double supplier' here because Groovy keeps reference to test data preventing it from being GCed
+ def supplier = supplierSupplier()
+ def ref = new WeakReference(supplier)
+
+ when:
+ def supplierRef = new WeakReference(supplier)
+ supplier = null
+ TestUtils.awaitGC(supplierRef)
+
+ then:
+ ref.get() == null
+
+ where:
+ name | supplierSupplier
+ "WeakConcurrent" | { -> new WeakMapSuppliers.WeakConcurrent() }
+ "WeakInline" | { -> new WeakMapSuppliers.WeakConcurrent.Inline() }
+ "Guava" | { -> new WeakMapSuppliers.Guava() }
+ }
+
+ def "Unreferenced map gets cleaned up on #name"() {
setup:
WeakMap.Provider.provider.set(supplier)
def map = WeakMap.Provider.newWeakMap()
def ref = new WeakReference(map)
when:
- def mapRef = new WeakReference<>(map)
+ def mapRef = new WeakReference(map)
map = null
TestUtils.awaitGC(mapRef)
@@ -69,7 +90,7 @@ class WeakConcurrentSupplierTest extends Specification {
map.size() == 1
when:
- def keyRef = new WeakReference<>(key)
+ def keyRef = new WeakReference(key)
key = null
TestUtils.awaitGC(keyRef)

Reorganize WeakConcurrent implementation to allow cleanup thread GC when map is out of scope