Merge pull request #619 from DataDog/mar-kolya/fix-memory-leak-in-weak-map

Mar kolya/fix memory leak in weak map

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<WeakReference<WeakConcurrentMap>> 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<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;
       }
@@ -115,6 +120,7 @@ class WeakMapSuppliers {
         }
         if (finalized.get() && suppliedMaps.isEmpty()) {
           executorService.shutdown();
+          Runtime.getRuntime().removeShutdownHook(shutdownCallback);
         }
       }
     }

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++) {

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<String, String> defaultSpanTags;
   /** A configured mapping of service names to update with new values */
   private final Map<String, String> 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<String, List<AbstractDecorator>> 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()
+      Runtime.getRuntime().addShutdownHook(shutdownCallback);
-          .addShutdownHook(
-              new Thread() {
-                @Override
-                public void run() {
-                  close();
-                }
-              });
     } 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
    *