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++) { 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 da28a2bf55..98008afca6 100644 --- a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java +++ b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java @@ -60,6 +60,11 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace private final Map defaultSpanTags; /** A configured mapping of service names to update with new values */ private final Map serviceNameMappings; + /** + * JVM shutdown callback, keeping a reference to it to remove this if DDTracer gets destroyed + * earlier + */ + private final Thread shutdownCallback; /** Span context decorators */ private final Map> spanContextDecorators = @@ -163,15 +168,15 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace this.runtimeId = runtimeId; this.serviceNameMappings = serviceNameMappings; + shutdownCallback = + new Thread() { + @Override + public void run() { + close(); + } + }; try { - Runtime.getRuntime() - .addShutdownHook( - new Thread() { - @Override - public void run() { - close(); - } - }); + Runtime.getRuntime().addShutdownHook(shutdownCallback); } catch (final IllegalStateException ex) { // The JVM is already shutting down. } @@ -200,6 +205,12 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace log.info("New instance: {}", this); } + @Override + public void finalize() { + Runtime.getRuntime().removeShutdownHook(shutdownCallback); + shutdownCallback.run(); + } + /** * Returns the list of span context decorators *