001package io.prometheus.client.hotspot;
002
003import com.sun.management.GarbageCollectionNotificationInfo;
004import com.sun.management.GcInfo;
005import io.prometheus.client.Collector;
006import io.prometheus.client.Counter;
007
008import javax.management.Notification;
009import javax.management.NotificationEmitter;
010import javax.management.NotificationListener;
011import javax.management.openmbean.CompositeData;
012import java.lang.management.GarbageCollectorMXBean;
013import java.lang.management.ManagementFactory;
014import java.lang.management.MemoryUsage;
015import java.util.HashMap;
016import java.util.List;
017import java.util.Map;
018
019public class MemoryAllocationExports extends Collector {
020  private final Counter allocatedCounter = Counter.build()
021          .name("jvm_memory_pool_allocated_bytes_total")
022          .help("Total bytes allocated in a given JVM memory pool. Only updated after GC, not continuously.")
023          .labelNames("pool")
024          .create();
025
026  public MemoryAllocationExports() {
027    AllocationCountingNotificationListener listener = new AllocationCountingNotificationListener(allocatedCounter);
028    for (GarbageCollectorMXBean garbageCollectorMXBean : ManagementFactory.getGarbageCollectorMXBeans()) {
029      ((NotificationEmitter) garbageCollectorMXBean).addNotificationListener(listener, null, null);
030    }
031  }
032
033  @Override
034  public List<MetricFamilySamples> collect() {
035    return allocatedCounter.collect();
036  }
037
038  static class AllocationCountingNotificationListener implements NotificationListener {
039    private final Map<String, Long> lastMemoryUsage = new HashMap<String, Long>();
040    private final Counter counter;
041
042    AllocationCountingNotificationListener(Counter counter) {
043      this.counter = counter;
044    }
045
046    @Override
047    public synchronized void handleNotification(Notification notification, Object handback) {
048      GarbageCollectionNotificationInfo info = GarbageCollectionNotificationInfo.from((CompositeData) notification.getUserData());
049      GcInfo gcInfo = info.getGcInfo();
050      Map<String, MemoryUsage> memoryUsageBeforeGc = gcInfo.getMemoryUsageBeforeGc();
051      Map<String, MemoryUsage> memoryUsageAfterGc = gcInfo.getMemoryUsageAfterGc();
052      for (Map.Entry<String, MemoryUsage> entry : memoryUsageBeforeGc.entrySet()) {
053        String memoryPool = entry.getKey();
054        long before = entry.getValue().getUsed();
055        long after = memoryUsageAfterGc.get(memoryPool).getUsed();
056        handleMemoryPool(memoryPool, before, after);
057      }
058    }
059
060    // Visible for testing
061    void handleMemoryPool(String memoryPool, long before, long after) {
062      /*
063       * Calculate increase in the memory pool by comparing memory used
064       * after last GC, before this GC, and after this GC.
065       * See ascii illustration below.
066       * Make sure to count only increases and ignore decreases.
067       * (Typically a pool will only increase between GCs or during GCs, not both.
068       * E.g. eden pools between GCs. Survivor and old generation pools during GCs.)
069       *
070       *                         |<-- diff1 -->|<-- diff2 -->|
071       * Timeline: |-- last GC --|             |---- GC -----|
072       *                      ___^__        ___^____      ___^___
073       * Mem. usage vars:    / last \      / before \    / after \
074       */
075
076      // Get last memory usage after GC and remember memory used after for next time
077      long last = getAndSet(lastMemoryUsage, memoryPool, after);
078      // Difference since last GC
079      long diff1 = before - last;
080      // Difference during this GC
081      long diff2 = after - before;
082      // Make sure to only count increases
083      if (diff1 < 0) {
084        diff1 = 0;
085      }
086      if (diff2 < 0) {
087        diff2 = 0;
088      }
089      long increase = diff1 + diff2;
090      if (increase > 0) {
091        counter.labels(memoryPool).inc(increase);
092      }
093    }
094
095    private static long getAndSet(Map<String, Long> map, String key, long value) {
096      Long last = map.put(key, value);
097      return last == null ? 0 : last;
098    }
099  }
100}