From 3d92e5b8a92db4f02a03d809f443293f317fb35f Mon Sep 17 00:00:00 2001
From: Jon Chambers <63609320+jon-signal@users.noreply.github.com>
Date: Tue, 24 Oct 2023 16:50:02 -0400
Subject: [PATCH] Explicitly stop and start managed dependencies

---
 ...bstractSinglePassCrawlAccountsCommand.java | 21 ++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractSinglePassCrawlAccountsCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractSinglePassCrawlAccountsCommand.java
index 0606de558..351705552 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractSinglePassCrawlAccountsCommand.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractSinglePassCrawlAccountsCommand.java
@@ -70,16 +70,31 @@ protected void run(final Environment environment, final Namespace namespace,
 
     final int segments = Objects.requireNonNull(namespace.getInt(SEGMENT_COUNT));
 
+    environment.lifecycle().manage(new CommandStopListener(configuration.getCommandStopListener()));
+
     logger.info("Crawling accounts with {} segments and {} processors",
         segments,
         Runtime.getRuntime().availableProcessors());
 
-    final CommandStopListener commandStopListener = new CommandStopListener(configuration.getCommandStopListener());
     try {
-      commandStopListener.start();
+      environment.lifecycle().getManagedObjects().forEach(managedObject -> {
+        try {
+          managedObject.start();
+        } catch (final Exception e) {
+          logger.error("Failed to start managed object", e);
+          throw new RuntimeException(e);
+        }
+      });
+
       crawlAccounts(commandDependencies.accountsManager().streamAllFromDynamo(segments, Schedulers.parallel()));
     } finally {
-      commandStopListener.stop();
+      environment.lifecycle().getManagedObjects().forEach(managedObject -> {
+        try {
+          managedObject.stop();
+        } catch (final Exception e) {
+          logger.error("Failed to stop managed object", e);
+        }
+      });
     }
   }