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 : getGarbageCollectorMXBeans()) {
029      if (garbageCollectorMXBean instanceof NotificationEmitter) {
030        ((NotificationEmitter) garbageCollectorMXBean).addNotificationListener(listener, null, null);
031      }
032    }
033  }
034
035  @Override
036  public List<MetricFamilySamples> collect() {
037    return allocatedCounter.collect();
038  }
039
040  static class AllocationCountingNotificationListener implements NotificationListener {
041    private final Map<String, Long> lastMemoryUsage = new HashMap<String, Long>();
042    private final Counter counter;
043
044    AllocationCountingNotificationListener(Counter counter) {
045      this.counter = counter;
046    }
047
048    @Override
049    public synchronized void handleNotification(Notification notification, Object handback) {
050      GarbageCollectionNotificationInfo info = GarbageCollectionNotificationInfo.from((CompositeData) notification.getUserData());
051      GcInfo gcInfo = info.getGcInfo();
052      Map<String, MemoryUsage> memoryUsageBeforeGc = gcInfo.getMemoryUsageBeforeGc();
053      Map<String, MemoryUsage> memoryUsageAfterGc = gcInfo.getMemoryUsageAfterGc();
054      for (Map.Entry<String, MemoryUsage> entry : memoryUsageBeforeGc.entrySet()) {
055        String memoryPool = entry.getKey();
056        long before = entry.getValue().getUsed();
057        long after = memoryUsageAfterGc.get(memoryPool).getUsed();
058        handleMemoryPool(memoryPool, before, after);
059      }
060    }
061
062    // Visible for testing
063    void handleMemoryPool(String memoryPool, long before, long after) {
064      /*
065       * Calculate increase in the memory pool by comparing memory used
066       * after last GC, before this GC, and after this GC.
067       * See ascii illustration below.
068       * Make sure to count only increases and ignore decreases.
069       * (Typically a pool will only increase between GCs or during GCs, not both.
070       * E.g. eden pools between GCs. Survivor and old generation pools during GCs.)
071       *
072       *                         |<-- diff1 -->|<-- diff2 -->|
073       * Timeline: |-- last GC --|             |---- GC -----|
074       *                      ___^__        ___^____      ___^___
075       * Mem. usage vars:    / last \      / before \    / after \
076       */
077
078      // Get last memory usage after GC and remember memory used after for next time
079      long last = getAndSet(lastMemoryUsage, memoryPool, after);
080      // Difference since last GC
081      long diff1 = before - last;
082      // Difference during this GC
083      long diff2 = after - before;
084      // Make sure to only count increases
085      if (diff1 < 0) {
086        diff1 = 0;
087      }
088      if (diff2 < 0) {
089        diff2 = 0;
090      }
091      long increase = diff1 + diff2;
092      if (increase > 0) {
093        counter.labels(memoryPool).inc(increase);
094      }
095    }
096
097    private static long getAndSet(Map<String, Long> map, String key, long value) {
098      Long last = map.put(key, value);
099      return last == null ? 0 : last;
100    }
101  }
102
103  protected List<GarbageCollectorMXBean> getGarbageCollectorMXBeans() {
104    return ManagementFactory.getGarbageCollectorMXBeans();
105  }
106}