Add CommandStopListener

service/config/sample.yml

@@ -432,3 +432,6 @@ registrationService:
 
 turn:
   secret: secret://turn.secret
+
+commandStopListener:
+  path: /example/path

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java

@@ -26,6 +26,7 @@ import org.whispersystems.textsecuregcm.configuration.BadgesConfiguration;
 import org.whispersystems.textsecuregcm.configuration.BraintreeConfiguration;
 import org.whispersystems.textsecuregcm.configuration.CdnConfiguration;
 import org.whispersystems.textsecuregcm.configuration.ClientReleaseConfiguration;
+import org.whispersystems.textsecuregcm.configuration.CommandStopListenerConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DatadogConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DirectoryV2Configuration;
 import org.whispersystems.textsecuregcm.configuration.DynamoDbClientConfiguration;
@@ -294,6 +295,11 @@ public class WhisperServerConfiguration extends Configuration {
   @JsonProperty
   private MessageByteLimitCardinalityEstimatorConfiguration messageByteLimitCardinalityEstimator = new MessageByteLimitCardinalityEstimatorConfiguration(Duration.ofDays(1));
 
+  @Valid
+  @NotNull
+  @JsonProperty
+  private CommandStopListenerConfiguration commandStopListener;
+
   public AdminEventLoggingConfiguration getAdminEventLoggingConfiguration() {
     return adminEventLoggingConfiguration;
   }
@@ -488,4 +494,8 @@ public class WhisperServerConfiguration extends Configuration {
   public MessageByteLimitCardinalityEstimatorConfiguration getMessageByteLimitCardinalityEstimator() {
     return messageByteLimitCardinalityEstimator;
   }
+
+  public CommandStopListenerConfiguration getCommandStopListener() {
+    return commandStopListener;
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/configuration/CommandStopListenerConfiguration.java

@@ -0,0 +1,12 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.configuration;
+
+import javax.validation.constraints.NotNull;
+
+public record CommandStopListenerConfiguration(@NotNull String path) {
+
+}

service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractSinglePassCrawlAccountsCommand.java

@@ -9,6 +9,7 @@ import io.dropwizard.Application;
 import io.dropwizard.cli.Cli;
 import io.dropwizard.cli.EnvironmentCommand;
 import io.dropwizard.setup.Environment;
+import java.util.Objects;
 import net.sourceforge.argparse4j.inf.Namespace;
 import net.sourceforge.argparse4j.inf.Subparser;
 import org.slf4j.Logger;
@@ -17,10 +18,8 @@ import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
 import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.util.logging.UncaughtExceptionHandler;
-import reactor.core.publisher.Flux;
 import reactor.core.publisher.ParallelFlux;
 import reactor.core.scheduler.Schedulers;
-import java.util.Objects;
 
 public abstract class AbstractSinglePassCrawlAccountsCommand extends EnvironmentCommand<WhisperServerConfiguration> {
 
@@ -69,7 +68,14 @@ public abstract class AbstractSinglePassCrawlAccountsCommand extends Environment
         segments,
         Runtime.getRuntime().availableProcessors());
 
-    crawlAccounts(commandDependencies.accountsManager().streamAllFromDynamo(segments, Schedulers.parallel()));
+    final CommandStopListener commandStopListener = new CommandStopListener(configuration.getCommandStopListener());
+    try {
+      commandStopListener.start();
+      crawlAccounts(commandDependencies.accountsManager().streamAllFromDynamo(segments, Schedulers.parallel()));
+    } finally {
+      commandStopListener.stop();
+    }
+
   }
 
   @Override

service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandStopListener.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.workers;
+
+import io.dropwizard.lifecycle.Managed;
+import java.io.FileWriter;
+import java.io.IOException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.configuration.CommandStopListenerConfiguration;
+
+/**
+ * When {@link Managed#stop()} is called, writes to a configured file path
+ * <p>
+ * Useful for coordinating process termination via a shared file-system, such as a Kubernetes volume mount.
+ */
+public class CommandStopListener implements Managed {
+
+  private static final Logger logger = LoggerFactory.getLogger(CommandStopListener.class);
+
+  private final String path;
+
+  CommandStopListener(CommandStopListenerConfiguration config) {
+    this.path = config.path();
+  }
+
+  @Override
+  public void start() {
+
+  }
+
+  @Override
+  public void stop() {
+    try {
+      try (FileWriter writer = new FileWriter(path)) {
+        writer.write("stopped");
+      }
+    } catch (final IOException e) {
+      logger.error("Failed to open file {}", path);
+    }
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/workers/CrawlAccountsCommand.java

@@ -142,6 +142,8 @@ public class CrawlAccountsCommand extends EnvironmentCommand<WhisperServerConfig
       }
     };
 
+    environment.lifecycle().manage(new CommandStopListener(configuration.getCommandStopListener()));
+
     environment.lifecycle().getManagedObjects().forEach(managedObject -> {
       try {
         managedObject.start();