From 6d81f69785180e419478f6c3717ec08389345abd Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Sat, 17 Jun 2023 10:03:07 -0500
Subject: [PATCH] Start and stop all lifecycle-managed objects in
 `CrawlAccountsCommand`

---
 .../textsecuregcm/metrics/MetricsUtil.java    | 41 -------------------
 .../workers/CrawlAccountsCommand.java         | 21 ++++++++--
 2 files changed, 18 insertions(+), 44 deletions(-)

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 494ccdbe0..c74f11bc6 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
@@ -5,11 +5,7 @@
 
 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;
@@ -90,41 +86,4 @@ public class MetricsUtil {
     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 fae59d68a..c2774d790 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CrawlAccountsCommand.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CrawlAccountsCommand.java
@@ -19,6 +19,7 @@ import net.sourceforge.argparse4j.inf.ArgumentParserException;
 import net.sourceforge.argparse4j.inf.ArgumentType;
 import net.sourceforge.argparse4j.inf.Namespace;
 import net.sourceforge.argparse4j.inf.Subparser;
+import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
@@ -39,12 +40,13 @@ public class CrawlAccountsCommand extends EnvironmentCommand<WhisperServerConfig
   private static final String CRAWL_TYPE = "crawlType";
   private static final String WORKER_COUNT = "workers";
 
+  private static final Logger logger = LoggerFactory.getLogger(CrawlAccountsCommand.class);
+
   public enum CrawlType implements ArgumentType<CrawlType> {
     GENERAL_PURPOSE,
     ACCOUNT_CLEANER,
     ;
 
-
     @Override
     public CrawlType convert(final ArgumentParser parser, final Argument arg, final String value)
         throws ArgumentParserException {
@@ -142,7 +144,14 @@ public class CrawlAccountsCommand extends EnvironmentCommand<WhisperServerConfig
       }
     };
 
-    MetricsUtil.startManagedReporters(environment);
+    environment.lifecycle().getManagedObjects().forEach(managedObject -> {
+      try {
+        managedObject.start();
+      } catch (final Exception e) {
+        logger.error("Failed to start managed object", e);
+        throw new RuntimeException(e);
+      }
+    });
 
     try {
       crawler.crawlAllAccounts();
@@ -150,6 +159,12 @@ public class CrawlAccountsCommand extends EnvironmentCommand<WhisperServerConfig
       LoggerFactory.getLogger(CrawlAccountsCommand.class).error("Error crawling accounts", e);
     }
 
-    MetricsUtil.stopManagedReporters(environment);
+    environment.lifecycle().getManagedObjects().forEach(managedObject -> {
+      try {
+        managedObject.stop();
+      } catch (final Exception e) {
+        logger.error("Failed to stop managed object", e);
+      }
+    });
   }
 }