Add a command for regenerating account constraint tables

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

@@ -269,6 +269,7 @@ import org.whispersystems.textsecuregcm.workers.IdleDeviceNotificationSchedulerF
 import org.whispersystems.textsecuregcm.workers.MessagePersisterServiceCommand;
 import org.whispersystems.textsecuregcm.workers.NotifyIdleDevicesCommand;
 import org.whispersystems.textsecuregcm.workers.ProcessScheduledJobsServiceCommand;
+import org.whispersystems.textsecuregcm.workers.RegenerateAccountConstraintDataCommand;
 import org.whispersystems.textsecuregcm.workers.RemoveExpiredAccountsCommand;
 import org.whispersystems.textsecuregcm.workers.RemoveExpiredBackupsCommand;
 import org.whispersystems.textsecuregcm.workers.RemoveExpiredLinkedDevicesCommand;
@@ -335,6 +336,8 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     bootstrap.addCommand(new ProcessScheduledJobsServiceCommand("process-idle-device-notification-jobs",
         "Processes scheduled jobs to send notifications to idle devices",
         new IdleDeviceNotificationSchedulerFactory()));
+
+    bootstrap.addCommand(new RegenerateAccountConstraintDataCommand());
   }
 
   @Override

service/src/main/java/org/whispersystems/textsecuregcm/storage/Accounts.java

@@ -41,6 +41,7 @@ import org.apache.commons.lang3.StringUtils;
 import org.signal.libsignal.zkgroup.backups.BackupCredentialType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.util.AsyncTimerUtil;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.ExceptionUtils;
@@ -59,6 +60,7 @@ import software.amazon.awssdk.services.dynamodb.model.Delete;
 import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.Put;
+import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
 import software.amazon.awssdk.services.dynamodb.model.QueryResponse;
 import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
@@ -89,7 +91,8 @@ public class Accounts {
 
   static final List<String> ACCOUNT_FIELDS_TO_EXCLUDE_FROM_SERIALIZATION = List.of("uuid", "usernameLinkHandle");
 
-  private static final ObjectWriter ACCOUNT_DDB_JSON_WRITER = SystemMapper.jsonMapper()
+  @VisibleForTesting
+  static final ObjectWriter ACCOUNT_DDB_JSON_WRITER = SystemMapper.jsonMapper()
       .writer(SystemMapper.excludingField(Account.class, ACCOUNT_FIELDS_TO_EXCLUDE_FROM_SERIALIZATION));
 
   private static final Timer CREATE_TIMER = Metrics.timer(name(Accounts.class, "create"));
@@ -1404,8 +1407,7 @@ public class Accounts {
       final String tableName,
       final AttributeValue uuidAttr,
       final String keyName,
-      final AttributeValue keyValue
+      final AttributeValue keyValue) {
-  ) {
     return TransactWriteItem.builder()
         .put(Put.builder()
             .tableName(tableName)
@@ -1470,6 +1472,68 @@ public class Accounts {
         .build();
   }
 
+  public CompletableFuture<Void> regenerateConstraints(final Account account) {
+    final List<CompletableFuture<?>> constraintFutures = new ArrayList<>();
+
+    constraintFutures.add(writeConstraint(phoneNumberConstraintTableName,
+        account.getIdentifier(IdentityType.ACI),
+        ATTR_ACCOUNT_E164,
+        AttributeValues.fromString(account.getNumber())));
+
+    constraintFutures.add(writeConstraint(phoneNumberIdentifierConstraintTableName,
+        account.getIdentifier(IdentityType.ACI),
+        ATTR_PNI_UUID,
+        AttributeValues.fromUUID(account.getPhoneNumberIdentifier())));
+
+    account.getUsernameHash().ifPresent(usernameHash ->
+        constraintFutures.add(writeUsernameConstraint(account.getIdentifier(IdentityType.ACI),
+            usernameHash,
+            Optional.empty())));
+
+    account.getUsernameHolds().forEach(usernameHold ->
+        constraintFutures.add(writeUsernameConstraint(account.getIdentifier(IdentityType.ACI),
+            usernameHold.usernameHash(),
+            Optional.of(Instant.ofEpochSecond(usernameHold.expirationSecs())))));
+
+    return CompletableFuture.allOf(constraintFutures.toArray(CompletableFuture[]::new));
+  }
+
+  private CompletableFuture<Void> writeConstraint(
+      final String tableName,
+      final UUID accountIdentifier,
+      final String keyName,
+      final AttributeValue keyValue) {
+
+    return dynamoDbAsyncClient.putItem(PutItemRequest.builder()
+            .tableName(tableName)
+            .item(Map.of(
+                keyName, keyValue,
+                KEY_ACCOUNT_UUID, AttributeValues.fromUUID(accountIdentifier)))
+        .build())
+        .thenRun(Util.NOOP);
+  }
+
+  private CompletableFuture<Void> writeUsernameConstraint(
+      final UUID accountIdentifier,
+      final byte[] usernameHash,
+      final Optional<Instant> maybeExpiration) {
+
+    final Map<String, AttributeValue> item = new HashMap<>(Map.of(
+        UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash),
+        UsernameTable.ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(accountIdentifier),
+        UsernameTable.ATTR_CONFIRMED, AttributeValues.fromBool(maybeExpiration.isEmpty())
+    ));
+
+    maybeExpiration.ifPresent(expiration ->
+        item.put(UsernameTable.ATTR_TTL, AttributeValues.fromLong(expiration.getEpochSecond())));
+
+    return dynamoDbAsyncClient.putItem(PutItemRequest.builder()
+            .tableName(usernamesConstraintTableName)
+            .item(item)
+        .build())
+        .thenRun(Util.NOOP);
+  }
+
   @Nonnull
   private static String extractCancellationReasonCodes(final TransactionCanceledException exception) {
     return exception.cancellationReasons().stream()

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

@@ -75,6 +75,7 @@ import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
  * Construct utilities commonly used by worker commands
  */
 record CommandDependencies(
+    Accounts accounts,
     AccountsManager accountsManager,
     ProfilesManager profilesManager,
     ReportMessageManager reportMessageManager,
@@ -290,6 +291,7 @@ record CommandDependencies(
     environment.lifecycle().manage(new ManagedAwsCrt());
 
     return new CommandDependencies(
+        accounts,
         accountsManager,
         profilesManager,
         reportMessageManager,

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

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.workers;
+
+import com.google.common.annotations.VisibleForTesting;
+import io.micrometer.core.instrument.Counter;
+import io.micrometer.core.instrument.Metrics;
+import java.time.Duration;
+import net.sourceforge.argparse4j.inf.Subparser;
+import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.Accounts;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.util.retry.Retry;
+
+public class RegenerateAccountConstraintDataCommand extends AbstractSinglePassCrawlAccountsCommand {
+
+  @VisibleForTesting
+  static final String DRY_RUN_ARGUMENT = "dry-run";
+
+  @VisibleForTesting
+  static final String MAX_CONCURRENCY_ARGUMENT = "max-concurrency";
+
+  @VisibleForTesting
+  static final String RETRIES_ARGUMENT = "retries";
+
+  private static final String PROCESSED_ACCOUNTS_COUNTER_NAME =
+      MetricsUtil.name(RegenerateAccountConstraintDataCommand.class, "processedAccounts");
+
+  public RegenerateAccountConstraintDataCommand() {
+    super("regenerate-account-constraint-data", "Regenerates account constraint data from a core account table");
+  }
+
+  @Override
+  public void configure(final Subparser subparser) {
+    super.configure(subparser);
+
+    subparser.addArgument("--dry-run")
+        .type(Boolean.class)
+        .dest(DRY_RUN_ARGUMENT)
+        .required(false)
+        .setDefault(true)
+        .help("If true, don’t actually write constraint data");
+
+    subparser.addArgument("--max-concurrency")
+        .type(Integer.class)
+        .dest(MAX_CONCURRENCY_ARGUMENT)
+        .setDefault(16)
+        .help("Max concurrency for DynamoDB operations");
+
+    subparser.addArgument("--retries")
+        .type(Integer.class)
+        .dest(RETRIES_ARGUMENT)
+        .setDefault(8)
+        .help("Maximum number of DynamoDB retries permitted per account");
+  }
+
+  @Override
+  protected void crawlAccounts(final Flux<Account> accountRecords) {
+    final boolean dryRun = getNamespace().getBoolean(DRY_RUN_ARGUMENT);
+    final int maxConcurrency = getNamespace().getInt(MAX_CONCURRENCY_ARGUMENT);
+    final int maxRetries = getNamespace().getInt(RETRIES_ARGUMENT);
+
+    final Accounts accounts = getCommandDependencies().accounts();
+
+    final Counter processedAccountsCounter = Metrics.counter(PROCESSED_ACCOUNTS_COUNTER_NAME,
+        "dryRun", String.valueOf(dryRun));
+
+    accountRecords
+        .doOnNext(ignored -> processedAccountsCounter.increment())
+        .flatMap(account -> dryRun
+                ? Mono.empty()
+                : Mono.fromFuture(() -> accounts.regenerateConstraints(account))
+                    .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1)).maxBackoff(Duration.ofSeconds(4))
+                        .onRetryExhaustedThrow((spec, rs) -> rs.failure())),
+            maxConcurrency)
+        .then()
+        .block();
+  }
+}

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsTest.java

@@ -25,11 +25,13 @@ import com.fasterxml.jackson.core.JsonProcessingException;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.time.Instant;
+import java.time.temporal.ChronoUnit;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Base64;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -78,6 +80,7 @@ import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.Put;
 import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.ReturnValuesOnConditionCheckFailure;
+import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItem;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsRequest;
 import software.amazon.awssdk.services.dynamodb.model.TransactWriteItemsResponse;
@@ -118,6 +121,9 @@ class AccountsTest {
   private final TestClock clock = TestClock.pinned(Instant.EPOCH);
   private Accounts accounts;
 
+  private record UsernameConstraint(UUID accountIdentifier, boolean confirmed, Optional<Instant> expiration) {
+  }
+
   @BeforeEach
   void setupAccountsDao() {
 
@@ -1685,7 +1691,7 @@ class AccountsTest {
   }
 
   @Test
-  public void testInvalidDeviceIdDeserialization() throws Exception {
+  void testInvalidDeviceIdDeserialization() throws Exception {
     final Account account = generateAccount("+18005551234", UUID.randomUUID(), UUID.randomUUID());
     final Device device2 = generateDevice((byte) 64);
     account.addDevice(device2);
@@ -1727,6 +1733,189 @@ class AccountsTest {
     assertInstanceOf(DeviceIdDeserializer.DeviceIdDeserializationException.class, cause);
   }
 
+  @Test
+  void testRegenerateConstraints() {
+    final Instant usernameHoldExpiration = clock.instant().plus(Accounts.USERNAME_HOLD_DURATION).truncatedTo(ChronoUnit.SECONDS);
+
+    final Account account = nextRandomAccount();
+    account.setUsernameHash(USERNAME_HASH_1);
+    account.setUsernameLinkDetails(UUID.randomUUID(), ENCRYPTED_USERNAME_1);
+    account.setUsernameHolds(List.of(new Account.UsernameHold(USERNAME_HASH_2, usernameHoldExpiration.getEpochSecond())));
+
+    writeAccountRecordWithoutConstraints(account);
+    accounts.regenerateConstraints(account).join();
+
+    // Check that constraints do what they should from a functional perspective
+    {
+      final Account conflictingNumberAccount = nextRandomAccount();
+      conflictingNumberAccount.setNumber(account.getNumber(), account.getIdentifier(IdentityType.PNI));
+
+      assertThrows(AccountAlreadyExistsException.class,
+          () -> accounts.create(conflictingNumberAccount, Collections.emptyList()));
+    }
+
+    {
+      final Account conflictingUsernameAccount = nextRandomAccount();
+      createAccount(conflictingUsernameAccount);
+
+      final CompletionException completionException = assertThrows(CompletionException.class,
+          () -> accounts.reserveUsernameHash(conflictingUsernameAccount, USERNAME_HASH_1, Accounts.USERNAME_HOLD_DURATION).join());
+
+      assertInstanceOf(UsernameHashNotAvailableException.class, completionException.getCause());
+    }
+
+    {
+      final Account conflictingUsernameHoldAccount = nextRandomAccount();
+      createAccount(conflictingUsernameHoldAccount);
+
+      final CompletionException completionException = assertThrows(CompletionException.class,
+          () -> accounts.reserveUsernameHash(conflictingUsernameHoldAccount, USERNAME_HASH_2, Accounts.USERNAME_HOLD_DURATION).join());
+
+      assertInstanceOf(UsernameHashNotAvailableException.class, completionException.getCause());
+    }
+
+    // Check that bare constraint records are written as expected
+    assertEquals(Optional.of(account.getIdentifier(IdentityType.ACI)),
+        getConstraintValue(Tables.NUMBERS.tableName(), Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber())));
+
+    assertEquals(Optional.of(account.getIdentifier(IdentityType.ACI)),
+        getConstraintValue(Tables.PNI_ASSIGNMENTS.tableName(), Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(account.getIdentifier(IdentityType.PNI))));
+
+    assertEquals(Optional.of(new UsernameConstraint(account.getIdentifier(IdentityType.ACI), true, Optional.empty())),
+        getUsernameConstraint(USERNAME_HASH_1));
+
+    assertEquals(Optional.of(new UsernameConstraint(account.getIdentifier(IdentityType.ACI), false, Optional.of(usernameHoldExpiration))),
+        getUsernameConstraint(USERNAME_HASH_2));
+  }
+
+  @Test
+  void testRegeneratedConstraintsMatchOriginalConstraints() {
+    final Instant usernameHoldExpiration = clock.instant().plus(Accounts.USERNAME_HOLD_DURATION).truncatedTo(ChronoUnit.SECONDS);
+
+    final Account account = nextRandomAccount();
+    account.setUsernameHash(USERNAME_HASH_1);
+    account.setUsernameLinkDetails(UUID.randomUUID(), ENCRYPTED_USERNAME_1);
+    account.setUsernameHolds(List.of(new Account.UsernameHold(USERNAME_HASH_2, usernameHoldExpiration.getEpochSecond())));
+
+    createAccount(account);
+    accounts.reserveUsernameHash(account, USERNAME_HASH_2, Accounts.USERNAME_HOLD_DURATION).join();
+    accounts.confirmUsernameHash(account, USERNAME_HASH_2, ENCRYPTED_USERNAME_2).join();
+    accounts.reserveUsernameHash(account, USERNAME_HASH_1, Accounts.USERNAME_HOLD_DURATION).join();
+    accounts.confirmUsernameHash(account, USERNAME_HASH_1, ENCRYPTED_USERNAME_1).join();
+
+    final Map<String, AttributeValue> originalE164ConstraintItem =
+        DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder()
+                .tableName(Tables.NUMBERS.tableName())
+                .key(Map.of(Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber())))
+                .build())
+            .item();
+
+    final Map<String, AttributeValue> originalPniConstraintItem =
+        DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder()
+                .tableName(Tables.PNI_ASSIGNMENTS.tableName())
+                .key(Map.of(Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(account.getIdentifier(IdentityType.PNI))))
+                .build())
+            .item();
+
+    final Set<Map<String, AttributeValue>> originalUsernameConstraints = new HashSet<>(
+        DYNAMO_DB_EXTENSION.getDynamoDbClient().scan(ScanRequest.builder()
+                .tableName(Tables.USERNAMES.tableName())
+                .build())
+            .items());
+
+    accounts.delete(account.getIdentifier(IdentityType.ACI), Collections.emptyList()).join();
+
+    writeAccountRecordWithoutConstraints(account);
+    accounts.regenerateConstraints(account).join();
+
+    final Map<String, AttributeValue> regeneratedE164ConstraintItem =
+        DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder()
+                .tableName(Tables.NUMBERS.tableName())
+                .key(Map.of(Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber())))
+                .build())
+            .item();
+
+    final Map<String, AttributeValue> regeneratedPniConstraintItem =
+        DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder()
+                .tableName(Tables.PNI_ASSIGNMENTS.tableName())
+                .key(Map.of(Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(account.getIdentifier(IdentityType.PNI))))
+                .build())
+            .item();
+
+    final Set<Map<String, AttributeValue>> regeneratedUsernameConstraints = new HashSet<>(
+        DYNAMO_DB_EXTENSION.getDynamoDbClient().scan(ScanRequest.builder()
+                .tableName(Tables.USERNAMES.tableName())
+                .build())
+            .items());
+
+    assertEquals(originalE164ConstraintItem, regeneratedE164ConstraintItem);
+    assertEquals(originalPniConstraintItem, regeneratedPniConstraintItem);
+    assertEquals(originalUsernameConstraints, regeneratedUsernameConstraints);
+  }
+
+  private void writeAccountRecordWithoutConstraints(final Account account) {
+    final AttributeValue accountData;
+
+    try {
+      accountData = AttributeValues.fromByteArray(Accounts.ACCOUNT_DDB_JSON_WRITER.writeValueAsBytes(account));
+    } catch (final JsonProcessingException e) {
+      throw new IllegalArgumentException(e);
+    }
+
+    final Map<String, AttributeValue> item = new HashMap<>(Map.of(
+        Accounts.KEY_ACCOUNT_UUID, AttributeValues.fromUUID(account.getUuid()),
+        Accounts.ATTR_ACCOUNT_E164, AttributeValues.fromString(account.getNumber()),
+        Accounts.ATTR_PNI_UUID, AttributeValues.fromUUID(account.getPhoneNumberIdentifier()),
+        Accounts.ATTR_ACCOUNT_DATA, accountData,
+        Accounts.ATTR_VERSION, AttributeValues.fromInt(account.getVersion()),
+        Accounts.ATTR_CANONICALLY_DISCOVERABLE, AttributeValues.fromBool(account.isDiscoverableByPhoneNumber())));
+
+    account.getUnidentifiedAccessKey()
+        .map(AttributeValues::fromByteArray)
+        .ifPresent(uak -> item.put(Accounts.ATTR_UAK, uak));
+
+    DYNAMO_DB_EXTENSION.getDynamoDbClient().putItem(PutItemRequest.builder()
+            .tableName(Tables.ACCOUNTS.tableName())
+            .item(item)
+        .build());
+  }
+
+  private Optional<UUID> getConstraintValue(final String tableName,
+      final String keyName,
+      final AttributeValue keyValue) {
+
+    final GetItemResponse response = DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder()
+            .tableName(tableName)
+            .key(Map.of(keyName, keyValue))
+        .build());
+
+    return response.hasItem()
+        ? Optional.ofNullable(AttributeValues.getUUID(response.item(), Accounts.KEY_ACCOUNT_UUID, null))
+        : Optional.empty();
+  }
+
+  private Optional<UsernameConstraint> getUsernameConstraint(final byte[] usernameHash) {
+    final GetItemResponse response = DYNAMO_DB_EXTENSION.getDynamoDbClient().getItem(GetItemRequest.builder()
+            .tableName(Tables.USERNAMES.tableName())
+            .key(Map.of(Accounts.UsernameTable.KEY_USERNAME_HASH, AttributeValues.fromByteArray(usernameHash)))
+        .build());
+
+    if (response.hasItem()) {
+      final UUID accountIdentifier =
+          AttributeValues.getUUID(response.item(), Accounts.UsernameTable.ATTR_ACCOUNT_UUID, null);
+
+      final boolean confirmed = AttributeValues.getBool(response.item(), Accounts.UsernameTable.ATTR_CONFIRMED, false);
+
+      final Optional<Instant> expiration = response.item().containsKey(Accounts.UsernameTable.ATTR_TTL)
+          ? Optional.of(Instant.ofEpochSecond(AttributeValues.getLong(response.item(), Accounts.UsernameTable.ATTR_TTL, 0)))
+          : Optional.empty();
+
+      return Optional.of(new UsernameConstraint(accountIdentifier, confirmed, expiration));
+    }
+
+    return Optional.empty();
+  }
+
   private static Device generateDevice(byte id) {
     return DevicesHelper.createDevice(id);
   }

service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java

@@ -64,7 +64,8 @@ class FinishPushNotificationExperimentCommandTest {
               new PushNotificationExperimentSample<>(accountIdentifier, deviceId, true, "test", "test"));
         });
 
-    commandDependencies = new CommandDependencies(accountsManager,
+    commandDependencies = new CommandDependencies(null,
+        accountsManager,
         null,
         null,
         null,

service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java

@@ -51,6 +51,7 @@ class NotifyIdleDevicesCommandTest {
           null,
           null,
           null,
+          null,
           messagesManager,
           null,
           null,

service/src/test/java/org/whispersystems/textsecuregcm/workers/RegenerateAccountConstraintDataCommandTest.java

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.workers;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import net.sourceforge.argparse4j.inf.Namespace;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.Accounts;
+import reactor.core.publisher.Flux;
+
+class RegenerateAccountConstraintDataCommandTest {
+
+  private Accounts accounts;
+
+  private static class TestRegenerateAccountConstraintDataCommand extends RegenerateAccountConstraintDataCommand {
+
+    private final CommandDependencies commandDependencies;
+    private final Namespace namespace;
+
+    TestRegenerateAccountConstraintDataCommand(final Accounts accounts, final boolean dryRun) {
+      commandDependencies = new CommandDependencies(accounts,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null,
+          null);
+
+      namespace = new Namespace(Map.of(
+          RegenerateAccountConstraintDataCommand.DRY_RUN_ARGUMENT, dryRun,
+          RegenerateAccountConstraintDataCommand.MAX_CONCURRENCY_ARGUMENT, 16,
+          RegenerateAccountConstraintDataCommand.RETRIES_ARGUMENT, 3));
+    }
+
+    @Override
+    protected CommandDependencies getCommandDependencies() {
+      return commandDependencies;
+    }
+
+    @Override
+    protected Namespace getNamespace() {
+      return namespace;
+    }
+  }
+
+  @BeforeEach
+  void setUp() {
+    accounts = mock(Accounts.class);
+
+    when(accounts.regenerateConstraints(any())).thenReturn(CompletableFuture.completedFuture(null));
+  }
+
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void crawlAccounts(final boolean dryRun) {
+    final Account account = mock(Account.class);
+
+    final RegenerateAccountConstraintDataCommand regenerateAccountConstraintDataCommand =
+        new TestRegenerateAccountConstraintDataCommand(accounts, dryRun);
+
+    regenerateAccountConstraintDataCommand.crawlAccounts(Flux.just(account));
+
+    if (!dryRun) {
+      verify(accounts).regenerateConstraints(account);
+    }
+
+    verifyNoMoreInteractions(accounts);
+  }
+}

service/src/test/java/org/whispersystems/textsecuregcm/workers/StartPushNotificationExperimentCommandTest.java

@@ -64,6 +64,7 @@ class StartPushNotificationExperimentCommandTest {
           null,
           null,
           null,
+          null,
           pushNotificationExperimentSamples,
           null,
           null,