Add MemoryAllocationExports (#434)

* 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 <[email protected]>

Author: snago

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,100 @@
+package io.prometheus.client.hotspot;
+
+import com.sun.management.GarbageCollectionNotificationInfo;
+import com.sun.management.GcInfo;
+import io.prometheus.client.Collector;
+import io.prometheus.client.Counter;
+
+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.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class MemoryAllocationExports extends Collector {
+  private final Counter allocatedCounter = Counter.build()
+          .name("jvm_memory_pool_allocated_bytes_total")
+          .help("Total bytes allocated in a given JVM memory pool. Only updated after GC, not continuously.")
+          .labelNames("pool")
+          .create();
+
+  public MemoryAllocationExports() {
+    AllocationCountingNotificationListener listener = new AllocationCountingNotificationListener(allocatedCounter);
+    for (GarbageCollectorMXBean garbageCollectorMXBean : ManagementFactory.getGarbageCollectorMXBeans()) {
+      ((NotificationEmitter) garbageCollectorMXBean).addNotificationListener(listener, null, null);
+    }
+  }
+
+  @Override
+  public List<MetricFamilySamples> collect() {
+    return allocatedCounter.collect();
+  }
+
+  static class AllocationCountingNotificationListener implements NotificationListener {
+    private final Map<String, Long> lastMemoryUsage = new HashMap<String, Long>();
+    private final Counter counter;
+
+    AllocationCountingNotificationListener(Counter counter) {
+      this.counter = counter;
+    }
+
+    @Override
+    public synchronized 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) {
+      /*
+       * Calculate increase in the memory pool by comparing memory used
+       * after last GC, before this GC, and after this GC.
+       * See ascii illustration below.
+       * Make sure to count only increases and ignore decreases.
+       * (Typically a pool will only increase between GCs or during GCs, not both.
+       * E.g. eden pools between GCs. Survivor and old generation pools during GCs.)
+       *
+       *                         |<-- diff1 -->|<-- diff2 -->|
+       * Timeline: |-- last GC --|             |---- GC -----|
+       *                      ___^__        ___^____      ___^___
+       * Mem. usage vars:    / last \      / before \    / after \
+       */
+
+      // Get last memory usage after GC and remember memory used after for next time
+      long last = getAndSet(lastMemoryUsage, memoryPool, after);
+      // Difference since last GC
+      long diff1 = before - last;
+      // Difference during this GC
+      long diff2 = after - before;
+      // Make sure to only count increases
+      if (diff1 < 0) {
+        diff1 = 0;
+      }
+      if (diff2 < 0) {
+        diff2 = 0;
+      }
+      long increase = diff1 + diff2;
+      if (increase > 0) {
+        counter.labels(memoryPool).inc(increase);
+      }
+    }
+
+    private static long getAndSet(Map<String, Long> map, String key, long value) {
+      Long last = map.put(key, value);
+      return last == null ? 0 : last;
+    }
+  }
+}

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

@@ -0,0 +1,45 @@
+package io.prometheus.client.hotspot;
+
+import io.prometheus.client.Counter;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class MemoryAllocationExportsTest {
+
+  @Test
+  public void testListenerLogic() {
+    Counter counter = Counter.build("test", "test").labelNames("pool").create();
+    MemoryAllocationExports.AllocationCountingNotificationListener listener =
+            new MemoryAllocationExports.AllocationCountingNotificationListener(counter);
+
+    // Increase by 123
+    listener.handleMemoryPool("TestPool", 0, 123);
+    Counter.Child child = counter.labels("TestPool");
+    assertEquals(123, child.get(), 0.001);
+
+    // No increase
+    listener.handleMemoryPool("TestPool", 123, 123);
+    assertEquals(123, child.get(), 0.001);
+
+    // No increase, then decrease to 0
+    listener.handleMemoryPool("TestPool", 123, 0);
+    assertEquals(123, child.get(), 0.001);
+
+    // No increase, then increase by 7
+    listener.handleMemoryPool("TestPool", 0, 7);
+    assertEquals(130, child.get(), 0.001);
+
+    // Increase by 10, then decrease to 10
+    listener.handleMemoryPool("TestPool", 17, 10);
+    assertEquals(140, child.get(), 0.001);
+
+    // Increase by 7, then increase by 3
+    listener.handleMemoryPool("TestPool", 17, 20);
+    assertEquals(150, child.get(), 0.001);
+
+    // Decrease to 17, then increase by 3
+    listener.handleMemoryPool("TestPool", 17, 20);
+    assertEquals(153, child.get(), 0.001);
+  }
+}