Issue#809: Fix not working garbage collection

Signed-off-by: Drophoff <[email protected]>

Author: Drophoff

simpleclient_hotspot/src/main/java/io/prometheus/client/hotspot/DefaultExports.java

@@ -1,5 +1,10 @@
 package io.prometheus.client.hotspot;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+import io.prometheus.client.Collector;
 import io.prometheus.client.CollectorRegistry;
 
 /**
@@ -18,6 +23,18 @@
 public class DefaultExports {
 
   private static boolean initialized = false;
+  
+  private final static List<Collector> defaultCollectors = Collections.unmodifiableList(Arrays.asList(
+        new BufferPoolsExports(),
+        new ClassLoadingExports(),
+        new CompilationExports(),
+        new GarbageCollectorExports(),
+        new MemoryAllocationExports(),
+        new MemoryPoolsExports(),
+        new StandardExports(),
+        new ThreadExports(),
+        new VersionInfoExports()
+      ));
 
   /**
    * Register the default Hotspot collectors with the default
@@ -30,19 +47,48 @@ public static synchronized void initialize() {
       initialized = true;
     }
   }
+  
+  /**
+   * Release the resources used by the default Hotspot
+   * collectors. It is safe to call this method multiple times,
+   * as this will unregister all collectors once.
+   */
+  public static synchronized void terminate() {
+    if (initialized) {
+      unregister(CollectorRegistry.defaultRegistry);
+      initialized = false;
+    }
+  }
 
   /**
    * Register the default Hotspot collectors with the given registry.
+   * 
+   * @param registry to which the collector is added. Null values ​​are
+   * not allowed.
    */
   public static void register(CollectorRegistry registry) {
-    new BufferPoolsExports().register(registry);
-    new ClassLoadingExports().register(registry);
-    new CompilationExports().register(registry);
-    new GarbageCollectorExports().register(registry);
-    new MemoryAllocationExports().register(registry);
-    new MemoryPoolsExports().register(registry);
-    new StandardExports().register(registry);
-    new ThreadExports().register(registry);
-    new VersionInfoExports().register(registry);
+    for (Collector collector : defaultCollectors) {
+      collector.register(registry);
+      
+      if (collector instanceof MemoryAllocationExports) {
+        ((MemoryAllocationExports)collector).addGarbageCollectorListener();
+      }
+    }
+  }
+  
+  /**
+   * Unregister the default Hotspot collectors from the given registry.
+   * 
+   * @param registry from that the default collectors are removed. Null
+   * values ​​are not allowed.
+   */
+  public static void unregister(CollectorRegistry registry) {
+    for (Collector collector : defaultCollectors) {
+      registry.unregister(collector);
+      
+      if (collector instanceof MemoryAllocationExports) {
+        ((MemoryAllocationExports)collector).removeGarbageCollectorListener();
+      }
+    }
   }
 }

simpleclient_hotspot/src/main/java/io/prometheus/client/hotspot/MemoryAllocationExports.java

@@ -5,16 +5,19 @@
 import io.prometheus.client.Collector;
 import io.prometheus.client.Counter;
 
+import javax.management.ListenerNotFoundException;
 import javax.management.Notification;
 import javax.management.NotificationEmitter;
 import javax.management.NotificationListener;
 import javax.management.openmbean.CompositeData;
 import java.lang.management.GarbageCollectorMXBean;
 import java.lang.management.ManagementFactory;
 import java.lang.management.MemoryUsage;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.logging.Logger;
 
 public class MemoryAllocationExports extends Collector {
   private final Counter allocatedCounter = Counter.build()
@@ -23,18 +26,45 @@ public class MemoryAllocationExports extends Collector {
           .labelNames("pool")
           .create();
 
+  private static final Logger LOGGER = Logger.getLogger(MemoryAllocationExports.class.getName());
+
+  private NotificationListener listener = null;
+
   public MemoryAllocationExports() {
-    AllocationCountingNotificationListener listener = new AllocationCountingNotificationListener(allocatedCounter);
-    for (GarbageCollectorMXBean garbageCollectorMXBean : getGarbageCollectorMXBeans()) {
-      if (garbageCollectorMXBean instanceof NotificationEmitter) {
-        ((NotificationEmitter) garbageCollectorMXBean).addNotificationListener(listener, null, null);
-      }
-    }
+    addGarbageCollectorListener();
   }
 
   @Override
   public List<MetricFamilySamples> collect() {
-    return allocatedCounter.collect();
+    return null == listener ? new ArrayList<Collector.MetricFamilySamples>() : allocatedCounter.collect();
+  }
+
+  void removeGarbageCollectorListener() {
+    if (null != listener ) {
+      for (GarbageCollectorMXBean garbageCollectorMXBean : getGarbageCollectorMXBeans()) {
+        if (garbageCollectorMXBean instanceof NotificationEmitter) {
+          try {
+            ((NotificationEmitter) garbageCollectorMXBean).removeNotificationListener(listener);
+          } catch (ListenerNotFoundException e) {
+            LOGGER.warning("The default notification listener could not be removed from the garbage collector.");
+          }
+        }
+      }
+      
+      listener = null;
+    }
+  }
+  
+  void addGarbageCollectorListener() {
+    if (null == listener) {
+      listener = new AllocationCountingNotificationListener(allocatedCounter);
+      
+      for (GarbageCollectorMXBean garbageCollectorMXBean : getGarbageCollectorMXBeans()) {
+        if (garbageCollectorMXBean instanceof NotificationEmitter) {
+          ((NotificationEmitter) garbageCollectorMXBean).addNotificationListener(listener, null, null);
+        }
+      }
+    }
   }
 
   static class AllocationCountingNotificationListener implements NotificationListener {

simpleclient_hotspot/src/test/java/io/prometheus/client/hotspot/MemoryAllocationExportsTest.java

@@ -1,5 +1,6 @@
 package io.prometheus.client.hotspot;
 
+import io.prometheus.client.Collector.MetricFamilySamples;
 import io.prometheus.client.Counter;
 import org.junit.Test;
 import org.mockito.Mockito;
@@ -11,6 +12,8 @@
 import java.util.List;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.times;
 
 public class MemoryAllocationExportsTest {
 
@@ -62,7 +65,26 @@ protected List<GarbageCollectorMXBean> getGarbageCollectorMXBeans() {
       }
     };
 
-    Mockito.verify((NotificationEmitter) notificationEmitterMXBean).addNotificationListener(Mockito.any(MemoryAllocationExports.AllocationCountingNotificationListener.class), Mockito.isNull(), Mockito.isNull());
+    Mockito.verify((NotificationEmitter) notificationEmitterMXBean).addNotificationListener(Mockito.any(MemoryAllocationExports.AllocationCountingNotificationListener.class), Mockito.nullable(NotificationFilter.class), Mockito.nullable(Object.class));
     Mockito.verifyNoInteractions(notNotificationEmitterMXBean);
   }
+  
+  @Test
+  public void testEmptyMetricsWithRemovedNotificationEmitter() {
+    MemoryAllocationExports allocationExporter = new MemoryAllocationExports();
+    allocationExporter.removeGarbageCollectorListener();
+    
+    List<MetricFamilySamples> metrics = allocationExporter.collect();
+    
+    assertTrue("Metrics should be empty, if no notification emitter is registered", metrics.isEmpty());
+  }
+  
+  @Test
+  public void testAddNotificationEmitterNotCalledMultipleTimesDuringInstantiation() {
+    MemoryAllocationExports allocationExporter = Mockito.mock(MemoryAllocationExports.class);
+    
+    allocationExporter.addGarbageCollectorListener();
+    
+    Mockito.verify(allocationExporter, times(1)).addGarbageCollectorListener();
+  }
 }