diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractCommandWithDependencies.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractCommandWithDependencies.java new file mode 100644 index 000000000..d82cff3e6 --- /dev/null +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractCommandWithDependencies.java @@ -0,0 +1,75 @@ +/* + * Copyright 2024 Signal Messenger, LLC + * SPDX-License-Identifier: AGPL-3.0-only + */ +package org.whispersystems.textsecuregcm.workers; + +import io.dropwizard.core.Application; +import io.dropwizard.core.cli.Cli; +import io.dropwizard.core.cli.EnvironmentCommand; +import io.dropwizard.core.setup.Environment; +import net.sourceforge.argparse4j.inf.Namespace; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.whispersystems.textsecuregcm.WhisperServerConfiguration; +import org.whispersystems.textsecuregcm.metrics.MetricsUtil; +import org.whispersystems.textsecuregcm.util.logging.UncaughtExceptionHandler; + +/** + * Base class for one-shot commands that use {@link CommandDependencies}. + *

+ * Override {@link #run(Environment, Namespace, WhisperServerConfiguration, CommandDependencies)} in a child class to + * let the parent class handle common initialization of dependencies, metrics, and logging. + */ +public abstract class AbstractCommandWithDependencies extends EnvironmentCommand { + + private final Logger logger = LoggerFactory.getLogger(getClass()); + + protected AbstractCommandWithDependencies(final Application application, + final String name, final String description) { + super(application, name, description); + } + + /** + * Run the command with the given initialized {@link CommandDependencies} + */ + protected abstract void run(final Environment environment, final Namespace namespace, + final WhisperServerConfiguration configuration, final CommandDependencies commandDependencies) throws Exception; + + @Override + protected void run(final Environment environment, final Namespace namespace, + final WhisperServerConfiguration configuration) throws Exception { + UncaughtExceptionHandler.register(); + final CommandDependencies commandDependencies = CommandDependencies.build(getName(), environment, configuration); + MetricsUtil.configureRegistries(configuration, environment, commandDependencies.dynamicConfigurationManager()); + + try { + logger.info("Starting command dependencies"); + environment.lifecycle().getManagedObjects().forEach(managedObject -> { + try { + managedObject.start(); + } catch (final Exception e) { + logger.error("Failed to start managed object", e); + throw new RuntimeException(e); + } + }); + + run(environment, namespace, configuration, commandDependencies); + + } finally { + logger.info("Stopping command dependencies"); + environment.lifecycle().getManagedObjects().forEach(managedObject -> { + try { + managedObject.stop(); + } catch (final Exception e) { + logger.error("Failed to stop managed object", e); + } + }); + } + } + + @Override + public void onError(final Cli cli, final Namespace namespace, final Throwable throwable) { + logger.error("Unhandled error", throwable); + } +} 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 e7133ed54..30237ca8c 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractSinglePassCrawlAccountsCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/AbstractSinglePassCrawlAccountsCommand.java @@ -20,7 +20,7 @@ import org.whispersystems.textsecuregcm.util.logging.UncaughtExceptionHandler; import reactor.core.publisher.Flux; import reactor.core.scheduler.Schedulers; -public abstract class AbstractSinglePassCrawlAccountsCommand extends EnvironmentCommand { +public abstract class AbstractSinglePassCrawlAccountsCommand extends AbstractCommandWithDependencies { private CommandDependencies commandDependencies; private Namespace namespace; @@ -59,12 +59,9 @@ public abstract class AbstractSinglePassCrawlAccountsCommand extends Environment @Override protected void run(final Environment environment, final Namespace namespace, - final WhisperServerConfiguration configuration) throws Exception { - - UncaughtExceptionHandler.register(); - + final WhisperServerConfiguration configuration, final CommandDependencies commandDependencies) throws Exception { this.namespace = namespace; - this.commandDependencies = CommandDependencies.build(getName(), environment, configuration); + this.commandDependencies = commandDependencies; final int segments = Objects.requireNonNull(namespace.getInt(SEGMENT_COUNT)); @@ -72,31 +69,7 @@ public abstract class AbstractSinglePassCrawlAccountsCommand extends Environment segments, Runtime.getRuntime().availableProcessors()); - try { - 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 { - environment.lifecycle().getManagedObjects().forEach(managedObject -> { - try { - managedObject.stop(); - } catch (final Exception e) { - logger.error("Failed to stop managed object", e); - } - }); - } - } - - @Override - public void onError(final Cli cli, final Namespace namespace, final Throwable throwable) { - logger.error("Unhandled error", throwable); + crawlAccounts(commandDependencies.accountsManager().streamAllFromDynamo(segments, Schedulers.parallel())); } protected abstract void crawlAccounts(final Flux accounts); diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/BackupMetricsCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/BackupMetricsCommand.java index 2e3d9d6c1..96430236f 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/BackupMetricsCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/BackupMetricsCommand.java @@ -29,7 +29,7 @@ import java.util.Objects; import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name; -public class BackupMetricsCommand extends EnvironmentCommand { +public class BackupMetricsCommand extends AbstractCommandWithDependencies { private final Logger logger = LoggerFactory.getLogger(getClass()); @@ -61,70 +61,47 @@ public class BackupMetricsCommand extends EnvironmentCommand { - try { - managedObject.start(); - } catch (final Exception e) { - logger.error("Failed to start managed object", e); - throw new RuntimeException(e); - } - }); + final DistributionSummary numObjectsMediaTier = Metrics.summary(name(getClass(), "numObjects"), + "tier", BackupLevel.MEDIA.name()); + final DistributionSummary bytesUsedMediaTier = Metrics.summary(name(getClass(), "bytesUsed"), + "tier", BackupLevel.MEDIA.name()); + final DistributionSummary numObjectsMessagesTier = Metrics.summary(name(getClass(), "numObjects"), + "tier", BackupLevel.MESSAGES.name()); + final DistributionSummary bytesUsedMessagesTier = Metrics.summary(name(getClass(), "bytesUsed"), + "tier", BackupLevel.MESSAGES.name()); - final DistributionSummary numObjectsMediaTier = Metrics.summary(name(getClass(), "numObjects"), - "tier", BackupLevel.MEDIA.name()); - final DistributionSummary bytesUsedMediaTier = Metrics.summary(name(getClass(), "bytesUsed"), - "tier", BackupLevel.MEDIA.name()); - final DistributionSummary numObjectsMessagesTier = Metrics.summary(name(getClass(), "numObjects"), - "tier", BackupLevel.MESSAGES.name()); - final DistributionSummary bytesUsedMessagesTier = Metrics.summary(name(getClass(), "bytesUsed"), - "tier", BackupLevel.MESSAGES.name()); + final DistributionSummary timeSinceLastRefresh = Metrics.summary(name(getClass(), + "timeSinceLastRefresh")); + final DistributionSummary timeSinceLastMediaRefresh = Metrics.summary(name(getClass(), + "timeSinceLastMediaRefresh")); + final String backupsCounterName = name(getClass(), "backups"); - final DistributionSummary timeSinceLastRefresh = Metrics.summary(name(getClass(), - "timeSinceLastRefresh")); - final DistributionSummary timeSinceLastMediaRefresh = Metrics.summary(name(getClass(), - "timeSinceLastMediaRefresh")); - final String backupsCounterName = name(getClass(), "backups"); - - final BackupManager backupManager = commandDependencies.backupManager(); - final Long backupsExpired = backupManager - .listBackupAttributes(segments, Schedulers.parallel()) - .doOnNext(backupMetadata -> { - final boolean subscribed = backupMetadata.lastMediaRefresh().equals(backupMetadata.lastRefresh()); - if (subscribed) { - numObjectsMediaTier.record(backupMetadata.numObjects()); - bytesUsedMediaTier.record(backupMetadata.bytesUsed()); - } else { - numObjectsMessagesTier.record(backupMetadata.numObjects()); - bytesUsedMessagesTier.record(backupMetadata.bytesUsed()); - } - timeSinceLastRefresh.record(timeSince(backupMetadata.lastRefresh()).getSeconds()); - timeSinceLastMediaRefresh.record(timeSince(backupMetadata.lastMediaRefresh()).getSeconds()); - Metrics.counter(backupsCounterName, "subscribed", String.valueOf(subscribed)).increment(); - }) - .count() - .block(); - logger.info("Crawled {} backups", backupsExpired); - } finally { - environment.lifecycle().getManagedObjects().forEach(managedObject -> { - try { - managedObject.stop(); - } catch (final Exception e) { - logger.error("Failed to stop managed object", e); - } - }); - } + final BackupManager backupManager = commandDependencies.backupManager(); + final Long backupsExpired = backupManager + .listBackupAttributes(segments, Schedulers.parallel()) + .doOnNext(backupMetadata -> { + final boolean subscribed = backupMetadata.lastMediaRefresh().equals(backupMetadata.lastRefresh()); + if (subscribed) { + numObjectsMediaTier.record(backupMetadata.numObjects()); + bytesUsedMediaTier.record(backupMetadata.bytesUsed()); + } else { + numObjectsMessagesTier.record(backupMetadata.numObjects()); + bytesUsedMessagesTier.record(backupMetadata.bytesUsed()); + } + timeSinceLastRefresh.record(timeSince(backupMetadata.lastRefresh()).getSeconds()); + timeSinceLastMediaRefresh.record(timeSince(backupMetadata.lastMediaRefresh()).getSeconds()); + Metrics.counter(backupsCounterName, "subscribed", String.valueOf(subscribed)).increment(); + }) + .count() + .block(); + logger.info("Crawled {} backups", backupsExpired); } private Duration timeSince(Instant t) { @@ -134,9 +111,4 @@ public class BackupMetricsCommand extends EnvironmentCommand { +public class DeleteUserCommand extends AbstractCommandWithDependencies { private final Logger logger = LoggerFactory.getLogger(DeleteUserCommand.class); @@ -36,22 +34,17 @@ public class DeleteUserCommand extends EnvironmentCommand { +public class RemoveExpiredBackupsCommand extends AbstractCommandWithDependencies { private final Logger logger = LoggerFactory.getLogger(getClass()); @@ -89,12 +86,7 @@ public class RemoveExpiredBackupsCommand extends EnvironmentCommand { - try { - managedObject.start(); - } catch (final Exception e) { - logger.error("Failed to start managed object", e); - throw new RuntimeException(e); - } - }); - final BackupManager backupManager = commandDependencies.backupManager(); - final long backupsExpired = backupManager - .getExpiredBackups(segments, Schedulers.parallel(), clock.instant().minus(gracePeriod)) - .flatMap(expiredBackup -> removeExpiredBackup(backupManager, expiredBackup, dryRun), concurrency) - .filter(Boolean.TRUE::equals) - .count() - .block(); - logger.info("Expired {} backups", backupsExpired); - } finally { - environment.lifecycle().getManagedObjects().forEach(managedObject -> { - try { - managedObject.stop(); - } catch (final Exception e) { - logger.error("Failed to stop managed object", e); - } - }); - } + final BackupManager backupManager = commandDependencies.backupManager(); + final long backupsExpired = backupManager + .getExpiredBackups(segments, Schedulers.parallel(), clock.instant().minus(gracePeriod)) + .flatMap(expiredBackup -> removeExpiredBackup(backupManager, expiredBackup, dryRun), concurrency) + .filter(Boolean.TRUE::equals) + .count() + .block(); + logger.info("Expired {} backups", backupsExpired); } private Mono removeExpiredBackup( @@ -162,9 +136,4 @@ public class RemoveExpiredBackupsCommand extends EnvironmentCommand { +public class SetUserDiscoverabilityCommand extends AbstractCommandWithDependencies { public SetUserDiscoverabilityCommand() { @@ -48,12 +46,10 @@ public class SetUserDiscoverabilityCommand extends EnvironmentCommand maybeAccount; diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/UnlinkDeviceCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/UnlinkDeviceCommand.java index c6b9b5669..b73d6cde5 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/UnlinkDeviceCommand.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/UnlinkDeviceCommand.java @@ -19,7 +19,7 @@ import org.whispersystems.textsecuregcm.WhisperServerConfiguration; import org.whispersystems.textsecuregcm.storage.Account; import org.whispersystems.textsecuregcm.storage.Device; -public class UnlinkDeviceCommand extends EnvironmentCommand { +public class UnlinkDeviceCommand extends AbstractCommandWithDependencies { public UnlinkDeviceCommand() { super(new Application<>() { @@ -49,25 +49,22 @@ public class UnlinkDeviceCommand extends EnvironmentCommand deviceIds = namespace.getList("deviceIds"); - final UUID aci = UUID.fromString(namespace.getString("uuid").trim()); - final List deviceIds = namespace.getList("deviceIds"); + Account account = deps.accountsManager().getByAccountIdentifier(aci) + .orElseThrow(() -> new IllegalArgumentException("account id " + aci + " does not exist")); - final CommandDependencies deps = CommandDependencies.build("unlink-device", environment, configuration); + if (deviceIds.contains(Device.PRIMARY_ID)) { + throw new IllegalArgumentException("cannot delete primary device"); + } - Account account = deps.accountsManager().getByAccountIdentifier(aci) - .orElseThrow(() -> new IllegalArgumentException("account id " + aci + " does not exist")); - - if (deviceIds.contains(Device.PRIMARY_ID)) { - throw new IllegalArgumentException("cannot delete primary device"); - } - - for (byte deviceId : deviceIds) { - /** see {@link org.whispersystems.textsecuregcm.controllers.DeviceController#removeDevice} */ - System.out.format("Removing device %s::%d\n", aci, deviceId); - deps.accountsManager().removeDevice(account, deviceId).join(); - } + for (byte deviceId : deviceIds) { + /** see {@link org.whispersystems.textsecuregcm.controllers.DeviceController#removeDevice} */ + System.out.format("Removing device %s::%d\n", aci, deviceId); + deps.accountsManager().removeDevice(account, deviceId).join(); + } } }