From 4deb68bfd76bc7d910ca9aa0c75888295edacff6 Mon Sep 17 00:00:00 2001 From: Nikolay Martynov Date: Wed, 5 Dec 2018 14:48:07 -0500 Subject: [PATCH] Fix memory leak in WeakMapSuppliers.WeakConcurrent The problem there is that it references JVM shutdown hook that keeps reference to cleanup executor which potentially can keep references to all sort of things - and this doesn't get cleaned up untill JVM shutdown. Solution is to remove shutdown hook when supplier is being GCed. --- .../main/java/datadog/trace/bootstrap/WeakMap.java | 4 ---- .../datadog/trace/agent/tooling/AgentInstaller.java | 8 +++----- .../trace/agent/tooling/WeakMapSuppliers.java | 12 +++++++++--- .../agent/tooling/WeakConcurrentSupplierTest.groovy | 2 +- 4 files changed, 13 insertions(+), 13 deletions(-) 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 index 778426e2a5..b49da78791 100644 --- 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 @@ -29,10 +29,6 @@ public interface WeakMap { } } - public static boolean isProviderRegistered() { - return provider.get() != Supplier.DEFAULT; - } - public static WeakMap newWeakMap() { return provider.get().get(); } 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 0c68905aee..1e84ffd3b2 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 @@ -111,11 +111,9 @@ public class AgentInstaller { } private static void registerWeakMapProvider() { - if (!WeakMap.Provider.isProviderRegistered()) { - WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent()); - } - // WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent.Inline()); - // WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.Guava()); + WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent()); + // WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent.Inline()); + // WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.Guava()); } @Slf4j 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 index 679fde502a..daa050c130 100644 --- 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 @@ -54,7 +54,7 @@ class WeakMapSuppliers { }; private final ScheduledExecutorService cleanerExecutorService; - + private final Thread shutdownCallback; private final Queue> suppliedMaps = new ConcurrentLinkedQueue<>(); @@ -62,14 +62,16 @@ class WeakMapSuppliers { WeakConcurrent() { cleanerExecutorService = Executors.newScheduledThreadPool(1, THREAD_FACTORY); + shutdownCallback = new ShutdownCallback(cleanerExecutorService); + cleanerExecutorService.scheduleAtFixedRate( - new CleanupRunnable(cleanerExecutorService, suppliedMaps, finalized), + new CleanupRunnable(cleanerExecutorService, shutdownCallback, suppliedMaps, finalized), CLEAN_FREQUENCY_SECONDS, CLEAN_FREQUENCY_SECONDS, TimeUnit.SECONDS); try { - Runtime.getRuntime().addShutdownHook(new ShutdownCallback(cleanerExecutorService)); + Runtime.getRuntime().addShutdownHook(shutdownCallback); } catch (final IllegalStateException ex) { // The JVM is already shutting down. } @@ -90,14 +92,17 @@ class WeakMapSuppliers { private static class CleanupRunnable implements Runnable { private final ScheduledExecutorService executorService; + private final Thread shutdownCallback; private final Queue> suppliedMaps; private final AtomicBoolean finalized; public CleanupRunnable( final ScheduledExecutorService executorService, + final Thread shutdownCallback, final Queue> suppliedMaps, final AtomicBoolean finalized) { this.executorService = executorService; + this.shutdownCallback = shutdownCallback; this.suppliedMaps = suppliedMaps; this.finalized = finalized; } @@ -115,6 +120,7 @@ class WeakMapSuppliers { } if (finalized.get() && suppliedMaps.isEmpty()) { executorService.shutdown(); + Runtime.getRuntime().removeShutdownHook(shutdownCallback); } } } 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 index 6c3e37403c..997347ea4e 100644 --- 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 @@ -107,7 +107,7 @@ class WeakConcurrentSupplierTest extends Specification { // 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 + // For Guava it is specified in // com.google.common.collect.MapMakerInternalMap.DRAIN_THRESHOLD = 0x3F if (name == "Guava" || name == "WeakInline") { for (int i = 0; i <= 0x3F; i++) {