From 13afdbda9748ff1e2bd65800ca730ac1f94458b8 Mon Sep 17 00:00:00 2001
From: Chris Eager <79161849+eager-signal@users.noreply.github.com>
Date: Wed, 14 Jun 2023 16:48:23 -0500
Subject: [PATCH] Report system resource metrics from background tasks

---
 .../textsecuregcm/WhisperServerService.java   | 28 +--------
 .../textsecuregcm/metrics/MetricsUtil.java    | 60 +++++++++++++++++++
 .../workers/CrawlAccountsCommand.java         | 56 ++++++++---------
 .../MessagePersisterServiceCommand.java       |  2 +
 ...nPushNotificationSenderServiceCommand.java |  2 +
 5 files changed, 94 insertions(+), 54 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index fc571bb5f..2e83cb1ea 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -46,7 +46,6 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadPoolExecutor;
-import java.util.concurrent.TimeUnit;
 import javax.servlet.DispatcherType;
 import javax.servlet.FilterRegistration;
 import javax.servlet.ServletRegistration;
@@ -127,18 +126,8 @@ import org.whispersystems.textsecuregcm.mappers.NonNormalizedPhoneNumberExceptio
 import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.RegistrationServiceSenderExceptionMapper;
 import org.whispersystems.textsecuregcm.mappers.ServerRejectedExceptionMapper;
-import org.whispersystems.textsecuregcm.metrics.ApplicationShutdownMonitor;
-import org.whispersystems.textsecuregcm.metrics.BufferPoolGauges;
-import org.whispersystems.textsecuregcm.metrics.CpuUsageGauge;
-import org.whispersystems.textsecuregcm.metrics.FileDescriptorGauge;
-import org.whispersystems.textsecuregcm.metrics.FreeMemoryGauge;
-import org.whispersystems.textsecuregcm.metrics.GarbageCollectionGauges;
-import org.whispersystems.textsecuregcm.metrics.MaxFileDescriptorGauge;
 import org.whispersystems.textsecuregcm.metrics.MetricsApplicationEventListener;
 import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
-import org.whispersystems.textsecuregcm.metrics.NetworkReceivedGauge;
-import org.whispersystems.textsecuregcm.metrics.NetworkSentGauge;
-import org.whispersystems.textsecuregcm.metrics.OperatingSystemMemoryGauge;
 import org.whispersystems.textsecuregcm.metrics.ReportedMessageMetricsListener;
 import org.whispersystems.textsecuregcm.metrics.TrafficSource;
 import org.whispersystems.textsecuregcm.providers.MultiRecipientMessageProvider;
@@ -783,7 +772,6 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
       webSocketEnvironment.jersey().register(controller);
     }
 
-
     WebSocketEnvironment<AuthenticatedAccount> provisioningEnvironment = new WebSocketEnvironment<>(environment,
         webSocketEnvironment.getRequestLog(), 60000);
     provisioningEnvironment.jersey().register(new WebsocketRefreshApplicationEventListener(accountsManager, clientPresenceManager));
@@ -817,21 +805,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
 
     environment.healthChecks().register("cacheCluster", new RedisClusterHealthCheck(cacheCluster));
 
-    environment.lifecycle().manage(new ApplicationShutdownMonitor(Metrics.globalRegistry));
-
-    environment.metrics().register(name(CpuUsageGauge.class, "cpu"), new CpuUsageGauge(3, TimeUnit.SECONDS));
-    environment.metrics().register(name(FreeMemoryGauge.class, "free_memory"), new FreeMemoryGauge());
-    environment.metrics().register(name(NetworkSentGauge.class, "bytes_sent"), new NetworkSentGauge());
-    environment.metrics().register(name(NetworkReceivedGauge.class, "bytes_received"), new NetworkReceivedGauge());
-    environment.metrics().register(name(FileDescriptorGauge.class, "fd_count"), new FileDescriptorGauge());
-    environment.metrics().register(name(MaxFileDescriptorGauge.class, "max_fd_count"), new MaxFileDescriptorGauge());
-    environment.metrics()
-        .register(name(OperatingSystemMemoryGauge.class, "buffers"), new OperatingSystemMemoryGauge("Buffers"));
-    environment.metrics()
-        .register(name(OperatingSystemMemoryGauge.class, "cached"), new OperatingSystemMemoryGauge("Cached"));
-
-    BufferPoolGauges.registerMetrics();
-    GarbageCollectionGauges.registerMetrics();
+    MetricsUtil.registerSystemResourceMetrics(environment);
   }
 
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
index b01177a91..494ccdbe0 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
@@ -5,7 +5,11 @@
 
 package org.whispersystems.textsecuregcm.metrics;
 
+import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.SharedMetricRegistries;
+import io.dropwizard.lifecycle.JettyManaged;
+import io.dropwizard.lifecycle.setup.LifecycleEnvironment;
+import io.dropwizard.metrics.ScheduledReporterManager;
 import io.dropwizard.setup.Environment;
 import io.micrometer.core.instrument.Meter;
 import io.micrometer.core.instrument.Metrics;
@@ -13,6 +17,7 @@ import io.micrometer.core.instrument.Tags;
 import io.micrometer.core.instrument.config.MeterFilter;
 import io.micrometer.core.instrument.distribution.DistributionStatisticConfig;
 import io.micrometer.datadog.DatadogMeterRegistry;
+import java.util.concurrent.TimeUnit;
 import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
 import org.whispersystems.textsecuregcm.WhisperServerVersion;
 import org.whispersystems.textsecuregcm.util.Constants;
@@ -66,5 +71,60 @@ public class MetricsUtil {
     }
 
     environment.lifecycle().manage(new MicrometerRegistryManager(Metrics.globalRegistry));
+    environment.lifecycle().manage(new ApplicationShutdownMonitor(Metrics.globalRegistry));
+  }
+
+  public static void registerSystemResourceMetrics(final Environment environment) {
+    environment.metrics().register(name(CpuUsageGauge.class, "cpu"), new CpuUsageGauge(3, TimeUnit.SECONDS));
+    environment.metrics().register(name(FreeMemoryGauge.class, "free_memory"), new FreeMemoryGauge());
+    environment.metrics().register(name(NetworkSentGauge.class, "bytes_sent"), new NetworkSentGauge());
+    environment.metrics().register(name(NetworkReceivedGauge.class, "bytes_received"), new NetworkReceivedGauge());
+    environment.metrics().register(name(FileDescriptorGauge.class, "fd_count"), new FileDescriptorGauge());
+    environment.metrics().register(name(MaxFileDescriptorGauge.class, "max_fd_count"), new MaxFileDescriptorGauge());
+    environment.metrics()
+        .register(name(OperatingSystemMemoryGauge.class, "buffers"), new OperatingSystemMemoryGauge("Buffers"));
+    environment.metrics()
+        .register(name(OperatingSystemMemoryGauge.class, "cached"), new OperatingSystemMemoryGauge("Cached"));
+
+    BufferPoolGauges.registerMetrics();
+    GarbageCollectionGauges.registerMetrics();
+  }
+
+  /**
+   * For use in commands where {@link JettyManaged} doesn't apply
+   *
+   * @see io.dropwizard.metrics.MetricsFactory#configure(LifecycleEnvironment, MetricRegistry)
+   */
+  public static void startManagedReporters(Environment environment) {
+    environment.lifecycle().getManagedObjects().forEach(managedObject -> {
+      if (managedObject instanceof JettyManaged jettyManaged) {
+        if (jettyManaged.getManaged() instanceof ScheduledReporterManager scheduledReporterManager) {
+          try {
+            scheduledReporterManager.start();
+          } catch (final Exception e) {
+            throw new RuntimeException(e);
+          }
+        }
+      }
+    });
+  }
+
+  /**
+   * For use in commands where {@link JettyManaged} doesn't apply
+   *
+   * @see io.dropwizard.metrics.MetricsFactory#configure(LifecycleEnvironment, MetricRegistry)
+   */
+  public static void stopManagedReporters(final Environment environment) {
+    environment.lifecycle().getManagedObjects().forEach(lifeCycle -> {
+      if (lifeCycle instanceof JettyManaged jettyManaged) {
+        if (jettyManaged.getManaged() instanceof ScheduledReporterManager scheduledReporterManager) {
+          try {
+            scheduledReporterManager.stop();
+          } catch (final Exception e) {
+            throw new RuntimeException(e);
+          }
+        }
+      }
+    });
   }
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CrawlAccountsCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CrawlAccountsCommand.java
index bab25843e..4f97347cd 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CrawlAccountsCommand.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CrawlAccountsCommand.java
@@ -97,38 +97,39 @@ public class CrawlAccountsCommand extends EnvironmentCommand<WhisperServerConfig
             DynamicConfiguration.class);
 
     dynamicConfigurationManager.start();
+    MetricsUtil.registerSystemResourceMetrics(environment);
 
-    final AccountDatabaseCrawler crawler =
+    final AccountDatabaseCrawler crawler = switch ((CrawlType) namespace.get(CRAWL_TYPE)) {
+      case GENERAL_PURPOSE -> {
+        final AccountDatabaseCrawlerCache accountDatabaseCrawlerCache = new AccountDatabaseCrawlerCache(
+            cacheCluster,
+            AccountDatabaseCrawlerCache.GENERAL_PURPOSE_PREFIX);
 
-        switch ((CrawlType) namespace.get(CRAWL_TYPE)) {
-          case GENERAL_PURPOSE -> {
-            final AccountDatabaseCrawlerCache accountDatabaseCrawlerCache = new AccountDatabaseCrawlerCache(
-                cacheCluster,
-                AccountDatabaseCrawlerCache.GENERAL_PURPOSE_PREFIX);
+        yield new AccountDatabaseCrawler("General-purpose account crawler",
+            accountsManager,
+            accountDatabaseCrawlerCache, accountDatabaseCrawlerListeners,
+            configuration.getAccountDatabaseCrawlerConfiguration().getChunkSize(),
+            dynamicConfigurationManager
+        );
+      }
+      case ACCOUNT_CLEANER -> {
+        final ExecutorService accountDeletionExecutor = environment.lifecycle()
+            .executorService(name(getClass(), "accountCleaner-%d")).maxThreads(16).minThreads(16).build();
 
-            yield new AccountDatabaseCrawler("General-purpose account crawler",
-                accountsManager,
-                accountDatabaseCrawlerCache, accountDatabaseCrawlerListeners,
-                configuration.getAccountDatabaseCrawlerConfiguration().getChunkSize(),
-                dynamicConfigurationManager
-            );
-          }
-          case ACCOUNT_CLEANER -> {
-            final ExecutorService accountDeletionExecutor = environment.lifecycle()
-                .executorService(name(getClass(), "accountCleaner-%d")).maxThreads(16).minThreads(16).build();
+        final AccountDatabaseCrawlerCache accountDatabaseCrawlerCache = new AccountDatabaseCrawlerCache(
+            cacheCluster, AccountDatabaseCrawlerCache.ACCOUNT_CLEANER_PREFIX);
 
-            final AccountDatabaseCrawlerCache accountDatabaseCrawlerCache = new AccountDatabaseCrawlerCache(
-                cacheCluster, AccountDatabaseCrawlerCache.ACCOUNT_CLEANER_PREFIX);
+        yield new AccountDatabaseCrawler("Account cleaner crawler",
+            accountsManager,
+            accountDatabaseCrawlerCache,
+            List.of(new AccountCleaner(accountsManager, accountDeletionExecutor)),
+            configuration.getAccountDatabaseCrawlerConfiguration().getChunkSize(),
+            dynamicConfigurationManager
+        );
+      }
+    };
 
-            yield new AccountDatabaseCrawler("Account cleaner crawler",
-                accountsManager,
-                accountDatabaseCrawlerCache,
-                List.of(new AccountCleaner(accountsManager, accountDeletionExecutor)),
-                configuration.getAccountDatabaseCrawlerConfiguration().getChunkSize(),
-                dynamicConfigurationManager
-            );
-          }
-        };
+    MetricsUtil.startManagedReporters(environment);
 
     try {
       crawler.crawlAllAccounts();
@@ -136,5 +137,6 @@ public class CrawlAccountsCommand extends EnvironmentCommand<WhisperServerConfig
       LoggerFactory.getLogger(CrawlAccountsCommand.class).error("Error crawling accounts", e);
     }
 
+    MetricsUtil.stopManagedReporters(environment);
   }
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/MessagePersisterServiceCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/MessagePersisterServiceCommand.java
index 467adfef1..a1be8f06f 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/MessagePersisterServiceCommand.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/MessagePersisterServiceCommand.java
@@ -70,6 +70,8 @@ public class MessagePersisterServiceCommand extends ServerCommand<WhisperServerC
     environment.lifecycle().manage(deps.messagesCache());
     environment.lifecycle().manage(messagePersister);
 
+    MetricsUtil.registerSystemResourceMetrics(environment);
+
     super.run(environment, namespace, configuration);
   }
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/ScheduledApnPushNotificationSenderServiceCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/ScheduledApnPushNotificationSenderServiceCommand.java
index a1c3d2348..b2a95ed02 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/ScheduledApnPushNotificationSenderServiceCommand.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/ScheduledApnPushNotificationSenderServiceCommand.java
@@ -79,6 +79,8 @@ public class ScheduledApnPushNotificationSenderServiceCommand extends ServerComm
     environment.lifecycle().manage(apnSender);
     environment.lifecycle().manage(apnPushNotificationScheduler);
 
+    MetricsUtil.registerSystemResourceMetrics(environment);
+
     super.run(environment, namespace, configuration);
   }