From c5147e0c68709ec494f2cb9bc90b94f9b4258d98 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Fri, 9 Oct 2020 20:03:18 -0400
Subject: [PATCH] Report direct memory metrics.

---
 .../textsecuregcm/WhisperServerService.java           | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 158a35af7..2aa400e9b 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -42,6 +42,7 @@ import io.dropwizard.setup.Bootstrap;
 import io.dropwizard.setup.Environment;
 import io.micrometer.core.instrument.Clock;
 import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Tag;
 import io.micrometer.core.instrument.distribution.DistributionStatisticConfig;
 import io.micrometer.wavefront.WavefrontConfig;
 import io.micrometer.wavefront.WavefrontMeterRegistry;
@@ -162,6 +163,8 @@ import org.whispersystems.websocket.setup.WebSocketEnvironment;
 import javax.servlet.DispatcherType;
 import javax.servlet.FilterRegistration;
 import javax.servlet.ServletRegistration;
+import java.lang.management.BufferPoolMXBean;
+import java.lang.management.ManagementFactory;
 import java.security.Security;
 import java.time.Duration;
 import java.util.Collections;
@@ -453,6 +456,14 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     environment.metrics().register(name(GarbageCollectionCountGauge.class, "gc_count"), new GarbageCollectionCountGauge());
     environment.metrics().register(name(GarbageCollectionTimeGauge.class, "gc_time"), new GarbageCollectionTimeGauge());
     environment.metrics().register(name(MaxFileDescriptorGauge.class, "max_fd_count"), new MaxFileDescriptorGauge());
+
+    for (final BufferPoolMXBean bufferPoolMXBean : ManagementFactory.getPlatformMXBeans(BufferPoolMXBean.class)) {
+      final List<Tag> tags = List.of(Tag.of("name", bufferPoolMXBean.getName()));
+
+      Metrics.gauge(name(BufferPoolMXBean.class, "count"), tags, bufferPoolMXBean, BufferPoolMXBean::getCount);
+      Metrics.gauge(name(BufferPoolMXBean.class, "memory_used"), tags, bufferPoolMXBean, BufferPoolMXBean::getMemoryUsed);
+      Metrics.gauge(name(BufferPoolMXBean.class, "total_capacity"), tags, bufferPoolMXBean, BufferPoolMXBean::getTotalCapacity);
+    }
   }
 
   private void registerExceptionMappers(Environment environment, WebSocketEnvironment<Account> webSocketEnvironment, WebSocketEnvironment<Account> provisioningEnvironment) {