Port memory metrics to metric-schema for docs and validation (#1391)

Port memory metrics (`jvm.memory`) to metric-schema for better documentation and validation

Author: Carter Kozak

changelog/@unreleased/pr-1391.v2.yml

@@ -0,0 +1,6 @@
+type: improvement
+improvement:
+  description: Port memory metrics (`jvm.memory`) to metric-schema for better documentation
+    and validation
+  links:
+  - https://github.com/palantir/tritium/pull/1391

tritium-metrics-jvm/src/main/java/com/palantir/tritium/metrics/jvm/JvmMetrics.java

@@ -16,16 +16,20 @@
 
 package com.palantir.tritium.metrics.jvm;
 
+import com.codahale.metrics.Gauge;
+import com.codahale.metrics.RatioGauge;
 import com.codahale.metrics.jvm.BufferPoolMetricSet;
-import com.codahale.metrics.jvm.MemoryUsageGaugeSet;
 import com.codahale.metrics.jvm.ThreadDeadlockDetector;
 import com.google.common.base.Suppliers;
-import com.google.common.collect.Maps;
 import com.palantir.logsafe.Preconditions;
+import com.palantir.logsafe.logger.SafeLogger;
+import com.palantir.logsafe.logger.SafeLoggerFactory;
 import com.palantir.tritium.metrics.MetricRegistries;
 import com.palantir.tritium.metrics.registry.TaggedMetricRegistry;
 import java.lang.management.ClassLoadingMXBean;
 import java.lang.management.ManagementFactory;
+import java.lang.management.MemoryMXBean;
+import java.lang.management.MemoryUsage;
 import java.lang.management.RuntimeMXBean;
 import java.lang.management.ThreadInfo;
 import java.lang.management.ThreadMXBean;
@@ -37,6 +41,7 @@
 
 /** {@link JvmMetrics} provides a standard set of metrics for debugging java services. */
 public final class JvmMetrics {
+    private static final SafeLogger log = SafeLoggerFactory.get(JvmMetrics.class);
 
     /**
      * Registers a default set of metrics.
@@ -58,12 +63,7 @@ public static void register(TaggedMetricRegistry registry) {
         MetricRegistries.registerAll(
                 registry, "jvm.buffers", new BufferPoolMetricSet(ManagementFactory.getPlatformMBeanServer()));
         registerClassLoading(metrics);
-        MetricRegistries.registerAll(
-                registry,
-                "jvm.memory",
-                () -> Maps.filterKeys(
-                        // Memory pool metrics are already provided by MetricRegistries.registerMemoryPools
-                        new MemoryUsageGaugeSet().getMetrics(), name -> !name.startsWith("pools")));
+        registerJvmMemory(registry);
         registerThreads(metrics);
     }
 
@@ -126,6 +126,54 @@ private static Map<Thread.State, Integer> threadsByState(ThreadMXBean threads) {
         return threadsByState;
     }
 
+    private static void registerJvmMemory(TaggedMetricRegistry registry) {
+        MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();
+        if (memoryBean == null) {
+            log.warn("Failed to load the MemoryMXBean, jvm.memory metrics will not be recorded");
+            return;
+        }
+        JvmMemoryMetrics metrics = JvmMemoryMetrics.of(registry);
+        // jvm.memory.total
+        metrics.totalInit((Gauge<Long>) () -> memoryBean.getHeapMemoryUsage().getInit()
+                + memoryBean.getNonHeapMemoryUsage().getInit());
+        metrics.totalUsed((Gauge<Long>) () -> memoryBean.getHeapMemoryUsage().getUsed()
+                + memoryBean.getNonHeapMemoryUsage().getUsed());
+        metrics.totalMax((Gauge<Long>) () -> memoryBean.getHeapMemoryUsage().getMax()
+                + memoryBean.getNonHeapMemoryUsage().getMax());
+        metrics.totalCommitted(
+                (Gauge<Long>) () -> memoryBean.getHeapMemoryUsage().getCommitted()
+                        + memoryBean.getNonHeapMemoryUsage().getCommitted());
+        // jvm.memory.heap
+        metrics.heapInit((Gauge<Long>) () -> memoryBean.getHeapMemoryUsage().getInit());
+        metrics.heapUsed((Gauge<Long>) () -> memoryBean.getHeapMemoryUsage().getUsed());
+        metrics.heapMax((Gauge<Long>) () -> memoryBean.getHeapMemoryUsage().getMax());
+        metrics.heapCommitted(
+                (Gauge<Long>) () -> memoryBean.getHeapMemoryUsage().getCommitted());
+        metrics.heapUsage(new RatioGauge() {
+            @Override
+            protected Ratio getRatio() {
+                MemoryUsage heapMemoryUsage = memoryBean.getHeapMemoryUsage();
+                return Ratio.of(heapMemoryUsage.getUsed(), heapMemoryUsage.getMax());
+            }
+        });
+        // jvm.memory.non-heap
+        metrics.nonHeapInit(
+                (Gauge<Long>) () -> memoryBean.getNonHeapMemoryUsage().getInit());
+        metrics.nonHeapUsed(
+                (Gauge<Long>) () -> memoryBean.getNonHeapMemoryUsage().getUsed());
+        metrics.nonHeapMax(
+                (Gauge<Long>) () -> memoryBean.getNonHeapMemoryUsage().getMax());
+        metrics.nonHeapCommitted(
+                (Gauge<Long>) () -> memoryBean.getNonHeapMemoryUsage().getCommitted());
+        metrics.nonHeapUsage(new RatioGauge() {
+            @Override
+            protected Ratio getRatio() {
+                MemoryUsage nonHeapMemoryUsage = memoryBean.getNonHeapMemoryUsage();
+                return Ratio.of(nonHeapMemoryUsage.getUsed(), nonHeapMemoryUsage.getMax());
+            }
+        });
+    }
+
     private JvmMetrics() {
         throw new UnsupportedOperationException();
     }

tritium-metrics-jvm/src/main/metrics/metrics.yml

@@ -77,3 +77,48 @@ namespaces:
       threads.terminated.count:
         type: gauge
         docs: Number of live threads in the `TERMINATED` state.
+  jvm.memory:
+    docs: Java virtual machine memory usage metrics.
+    metrics:
+      total.init:
+        type: gauge
+        docs: Initial amount of memory requested by the JVM.
+      total.used:
+        type: gauge
+        docs: Total non-native memory used by the JVM.
+      total.max:
+        type: gauge
+        docs: Maximum amount of memory that may be used by the JVM before allocations fail.
+      total.committed:
+        type: gauge
+        docs: Amount of committed memory in bytes for the JVM to use.
+      heap.init:
+        type: gauge
+        docs: Initial amount of heap memory requested by the JVM.
+      heap.used:
+        type: gauge
+        docs: Heap memory used by the JVM.
+      heap.max:
+        type: gauge
+        docs: Maximum amount of heap that may be used by the JVM before allocations fail.
+      heap.committed:
+        type: gauge
+        docs: Amount of committed memory in bytes for the JVM heap.
+      heap.usage:
+        type: gauge
+        docs: Ratio of `jvm.memory.heap.used` to `jvm.memory.heap.max`.
+      non-heap.init:
+        type: gauge
+        docs: Initial amount of non-heap memory requested by the JVM.
+      non-heap.used:
+        type: gauge
+        docs: Non-heap memory used by the JVM.
+      non-heap.max:
+        type: gauge
+        docs: Maximum amount of non-heap that may be used by the JVM before direct allocations fail.
+      non-heap.committed:
+        type: gauge
+        docs: Amount of committed memory in bytes for the JVM non-heap (e.g. direct memory).
+      non-heap.usage:
+        type: gauge
+        docs: Ratio of `jvm.memory.non-heap.used` to `jvm.memory.non-heap.max`.