Add MemoryAllocationExports

Counters for total bytes allocated to each memory pool.
Can be used to show allocation rate and promotion rate.

Signed-off-by: Robin Karlsson <snago86@gmail.com>

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

@@ -22,6 +22,7 @@ public static synchronized void initialize() {
     if (!initialized) {
       new StandardExports().register();
       new MemoryPoolsExports().register();
+      new MemoryAllocationExports().register();
       new BufferPoolsExports().register();
       new GarbageCollectorExports().register();
       new ThreadExports().register();

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

@@ -0,0 +1,104 @@
+package io.prometheus.client.hotspot;
+
+import com.sun.management.GarbageCollectionNotificationInfo;
+import com.sun.management.GcInfo;
+import io.prometheus.client.Collector;
+import io.prometheus.client.CounterMetricFamily;
+
+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.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+
+public class MemoryAllocationExports extends Collector {
+  private final Map<String, AtomicLong> allocatedMap;
+
+  public MemoryAllocationExports() {
+    this(ManagementFactory.getGarbageCollectorMXBeans(), new ConcurrentHashMap<String, AtomicLong>());
+  }
+
+  // Visible for testing
+  MemoryAllocationExports(List<GarbageCollectorMXBean> garbageCollectorMXBeans, Map<String, AtomicLong> allocatedMap) {
+    this.allocatedMap = allocatedMap;
+    AllocationCountingNotificationListener listener = new AllocationCountingNotificationListener(allocatedMap);
+    for (GarbageCollectorMXBean garbageCollectorMXBean : garbageCollectorMXBeans) {
+      ((NotificationEmitter) garbageCollectorMXBean).addNotificationListener(listener, null, null);
+    }
+  }
+
+  @Override
+  public List<MetricFamilySamples> collect() {
+    List<MetricFamilySamples> sampleFamilies = new ArrayList<MetricFamilySamples>();
+    CounterMetricFamily allocated = new CounterMetricFamily(
+            "jvm_memory_pool_allocated_bytes",
+            "Total bytes allocated in a given JVM memory pool. Only updated after GC, not continuously.",
+            Collections.singletonList("pool"));
+    sampleFamilies.add(allocated);
+
+    for (Map.Entry<String, AtomicLong> entry : allocatedMap.entrySet()) {
+      String memoryPool = entry.getKey();
+      AtomicLong bytesAllocated = entry.getValue();
+      allocated.addMetric(Collections.singletonList(memoryPool), bytesAllocated.doubleValue());
+    }
+    return sampleFamilies;
+  }
+
+  static class AllocationCountingNotificationListener implements NotificationListener {
+    private final Map<String, AtomicLong> lastMemoryUsage = new ConcurrentHashMap<String, AtomicLong>();
+    private final Map<String, AtomicLong> bytesAllocatedMap;
+
+    AllocationCountingNotificationListener(Map<String, AtomicLong> bytesAllocatedMap) {
+      this.bytesAllocatedMap = bytesAllocatedMap;
+    }
+
+    @Override
+    public void handleNotification(Notification notification, Object handback) {
+      GarbageCollectionNotificationInfo info = GarbageCollectionNotificationInfo.from((CompositeData) notification.getUserData());
+      GcInfo gcInfo = info.getGcInfo();
+      Map<String, MemoryUsage> memoryUsageBeforeGc = gcInfo.getMemoryUsageBeforeGc();
+      Map<String, MemoryUsage> memoryUsageAfterGc = gcInfo.getMemoryUsageAfterGc();
+      for (Map.Entry<String, MemoryUsage> entry : memoryUsageBeforeGc.entrySet()) {
+        String memoryPool = entry.getKey();
+        long before = entry.getValue().getUsed();
+        long after = memoryUsageAfterGc.get(memoryPool).getUsed();
+        handleMemoryPool(memoryPool, before, after);
+      }
+    }
+
+    // Visible for testing
+    void handleMemoryPool(String memoryPool, long before, long after) {
+      AtomicLong last = getOrCreate(lastMemoryUsage, memoryPool);
+      long diff1 = before - last.getAndSet(after);
+      long diff2 = after - before;
+      if (diff1 < 0) {
+        diff1 = 0;
+      }
+      if (diff2 < 0) {
+        diff2 = 0;
+      }
+      long increase = diff1 + diff2;
+      if (increase > 0) {
+        AtomicLong bytesAllocated = getOrCreate(bytesAllocatedMap, memoryPool);
+        bytesAllocated.getAndAdd(increase);
+      }
+    }
+
+    private AtomicLong getOrCreate(Map<String, AtomicLong> map, String key) {
+      AtomicLong result = map.get(key);
+      if (result == null) {
+        result = new AtomicLong(0);
+        map.put(key, result);
+      }
+      return result;
+    }
+  }
+}

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

@@ -0,0 +1,89 @@
+package io.prometheus.client.hotspot;
+
+import io.prometheus.client.Collector;
+import org.junit.Test;
+
+import java.lang.management.GarbageCollectorMXBean;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+public class MemoryAllocationExportsTest {
+
+  @Test
+  public void testCollectEmpty() {
+    MemoryAllocationExports exports = new MemoryAllocationExports(
+            Collections.<GarbageCollectorMXBean>emptyList(),
+            Collections.<String, AtomicLong>emptyMap()
+    );
+
+    List<Collector.MetricFamilySamples> initialCollect = exports.collect();
+    assertNotNull(initialCollect);
+    assertEquals(1, initialCollect.size());
+    Collector.MetricFamilySamples metricFamilySamples = initialCollect.get(0);
+    assertEquals(Collector.Type.COUNTER, metricFamilySamples.type);
+    assertTrue(metricFamilySamples.samples.isEmpty());
+  }
+
+  @Test
+  public void testCollectSimple() {
+    MemoryAllocationExports exports = new MemoryAllocationExports(
+            Collections.<GarbageCollectorMXBean>emptyList(),
+            Collections.singletonMap("TestPool", new AtomicLong(123))
+    );
+
+    List<Collector.MetricFamilySamples> initialCollect = exports.collect();
+    assertNotNull(initialCollect);
+    assertEquals(1, initialCollect.size());
+    Collector.MetricFamilySamples metricFamilySamples = initialCollect.get(0);
+    assertEquals(Collector.Type.COUNTER, metricFamilySamples.type);
+    assertEquals(1, metricFamilySamples.samples.size());
+    Collector.MetricFamilySamples.Sample sample = metricFamilySamples.samples.get(0);
+    assertEquals(1, sample.labelNames.size());
+    assertEquals("pool", sample.labelNames.get(0));
+    assertEquals(1, sample.labelValues.size());
+    assertEquals("TestPool", sample.labelValues.get(0));
+    assertEquals(123, sample.value, 0.001);
+  }
+
+  @Test
+  public void testListenerLogic() {
+    HashMap<String, AtomicLong> allocatedMap = new HashMap<String, AtomicLong>();
+    MemoryAllocationExports.AllocationCountingNotificationListener listener =
+            new MemoryAllocationExports.AllocationCountingNotificationListener(allocatedMap);
+
+    // Increase by 123
+    listener.handleMemoryPool("TestPool", 0, 123);
+    assertTrue(allocatedMap.containsKey("TestPool"));
+    assertEquals(123, allocatedMap.get("TestPool").get());
+
+    // No increase
+    listener.handleMemoryPool("TestPool", 123, 123);
+    assertEquals(123, allocatedMap.get("TestPool").get());
+
+    // No increase, then decrease to 0
+    listener.handleMemoryPool("TestPool", 123, 0);
+    assertEquals(123, allocatedMap.get("TestPool").get());
+
+    // No increase, then increase by 7
+    listener.handleMemoryPool("TestPool", 0, 7);
+    assertEquals(130, allocatedMap.get("TestPool").get());
+
+    // Increase by 10, then decrease to 10
+    listener.handleMemoryPool("TestPool", 17, 10);
+    assertEquals(140, allocatedMap.get("TestPool").get());
+
+    // Increase by 7, then increase by 3
+    listener.handleMemoryPool("TestPool", 17, 20);
+    assertEquals(150, allocatedMap.get("TestPool").get());
+
+    // Decrease to 17, then increase by 3
+    listener.handleMemoryPool("TestPool", 17, 20);
+    assertEquals(153, allocatedMap.get("TestPool").get());
+  }
+}