From 48c7572dd54b19bf1a4b7137e3465c579ef29470 Mon Sep 17 00:00:00 2001 From: Chris Eager Date: Tue, 25 Jul 2023 11:02:56 -0500 Subject: [PATCH] Add CommandStopListener --- service/config/sample.yml | 3 ++ .../WhisperServerConfiguration.java | 10 +++++ .../CommandStopListenerConfiguration.java | 12 +++++ ...bstractSinglePassCrawlAccountsCommand.java | 12 +++-- .../workers/CommandStopListener.java | 45 +++++++++++++++++++ .../workers/CrawlAccountsCommand.java | 2 + 6 files changed, 81 insertions(+), 3 deletions(-) create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/configuration/CommandStopListenerConfiguration.java create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandStopListener.java diff --git a/service/config/sample.yml b/service/config/sample.yml index fc75ef797..1b9a3d2f8 100644 --- a/service/config/sample.yml +++ b/service/config/sample.yml @@ -432,3 +432,6 @@ registrationService: turn: secret: secret://turn.secret + +commandStopListener: + path: /example/path diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java index 69728f76c..1654ef1e6 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java +++ b/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; + } } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/CommandStopListenerConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/CommandStopListenerConfiguration.java new file mode 100644 index 000000000..0b43720aa --- /dev/null +++ b/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) { + +} 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 d50c9bffd..539efe54d 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractSinglePassCrawlAccountsCommand.java +++ b/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 { @@ -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 diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandStopListener.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandStopListener.java new file mode 100644 index 000000000..a6ca42b78 --- /dev/null +++ b/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 + *

+ * 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); + } + } +} 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 7b4775ef2..49ac15112 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CrawlAccountsCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CrawlAccountsCommand.java @@ -142,6 +142,8 @@ public class CrawlAccountsCommand extends EnvironmentCommand { try { managedObject.start();