Regenerate phone number identifiers when regenerating secondary table data

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

@@ -270,7 +270,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.RegenerateSecondaryDynamoDbTableDataCommand;
 import org.whispersystems.textsecuregcm.workers.RemoveExpiredAccountsCommand;
 import org.whispersystems.textsecuregcm.workers.RemoveExpiredBackupsCommand;
 import org.whispersystems.textsecuregcm.workers.RemoveExpiredLinkedDevicesCommand;
@@ -337,7 +337,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         "Processes scheduled jobs to send notifications to idle devices",
         new IdleDeviceNotificationSchedulerFactory()));
 
-    bootstrap.addCommand(new RegenerateAccountConstraintDataCommand());
+    bootstrap.addCommand(new RegenerateSecondaryDynamoDbTableDataCommand());
   }
 
   @Override

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

@@ -1472,7 +1472,7 @@ public class Accounts {
         .build();
   }
 
-  public CompletableFuture<Void> regenerateConstraints(final Account account) {
+  CompletableFuture<Void> regenerateConstraints(final Account account) {
     final List<CompletableFuture<?>> constraintFutures = new ArrayList<>();
 
     constraintFutures.add(writeConstraint(phoneNumberConstraintTableName,

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

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.storage;
+
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * The DynamoDB recovery manager regenerates data for secondary tables in a disaster recovery scenario. In a disaster
+ * recovery scenario, there is no guarantee that table backups will be consistent, and so we need to derive or update
+ * some tables from a "core" data source to ensure consistency.
+ */
+public class DynamoDbRecoveryManager {
+
+  private final Accounts accounts;
+  private final PhoneNumberIdentifiers phoneNumberIdentifiers;
+
+  public DynamoDbRecoveryManager(final Accounts accounts, final PhoneNumberIdentifiers phoneNumberIdentifiers) {
+    this.accounts = accounts;
+    this.phoneNumberIdentifiers = phoneNumberIdentifiers;
+  }
+
+  /**
+   * Regenerates secondary data (i.e. uniqueness constraints) for a given account.
+   *
+   * @param account the account for which to regenerate secondary data
+   *
+   * @return a future that completes when secondary for the given account has been regenerated
+   */
+  public CompletableFuture<Void> regenerateData(final Account account) {
+    return CompletableFuture.allOf(
+        accounts.regenerateConstraints(account),
+        phoneNumberIdentifiers.regeneratePhoneNumberIdentifierMappings(account));
+  }
+}

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

@@ -10,33 +10,24 @@ import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
 import com.google.common.annotations.VisibleForTesting;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Timer;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.ExceptionUtils;
 import org.whispersystems.textsecuregcm.util.Util;
-import reactor.core.publisher.Flux;
-import reactor.core.scheduler.Scheduler;
-import reactor.util.function.Tuple2;
-import reactor.util.function.Tuples;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
-import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.BatchGetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.CancellationReason;
-import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
 import software.amazon.awssdk.services.dynamodb.model.KeysAndAttributes;
 import software.amazon.awssdk.services.dynamodb.model.QueryRequest;
-import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
 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.TransactionCanceledException;
@@ -93,7 +84,7 @@ public class PhoneNumberIdentifiers {
    * UUID was not previously assigned as a PNI by {@link #getPhoneNumberIdentifier(String)}, the
    * returned list will be empty.
    *
-   * @param UUID a phone number identifier
+   * @param phoneNumberIdentifier a phone number identifier
    * @return the list of all e164s associated with the given phone number identifier
    */
   public CompletableFuture<List<String>> getPhoneNumber(final UUID phoneNumberIdentifier) {
@@ -110,12 +101,9 @@ public class PhoneNumberIdentifiers {
                 ":pni", AttributeValues.fromUUID(phoneNumberIdentifier)
             ))
             .build())
-        .thenApply(response -> {
+        .thenApply(response -> response.items().stream().map(item -> item.get(KEY_E164).s()).toList());
-              return response.items().stream().map(item -> item.get(KEY_E164).s()).toList();
-        });
   }
 
-
   @VisibleForTesting
   static <T, E extends Exception> CompletableFuture<T> retry(
       final int numRetries, final Class<E> exceptionToRetry, final Supplier<CompletableFuture<T>> supplier) {
@@ -256,4 +244,9 @@ public class PhoneNumberIdentifiers {
             item -> AttributeValues.getUUID(item, ATTR_PHONE_NUMBER_IDENTIFIER, null))))
         .whenComplete((ignored, throwable) -> sample.stop(GET_PNI_TIMER));
   }
+
+  CompletableFuture<Void> regeneratePhoneNumberIdentifierMappings(final Account account) {
+    return setPni(account.getNumber(), Util.getAlternateForms(account.getNumber()), account.getIdentifier(IdentityType.PNI))
+        .thenRun(Util.NOOP);
+  }
 }

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

@@ -39,9 +39,9 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.MicrometerAwsSdkMetricPublisher;
 import org.whispersystems.textsecuregcm.push.APNSender;
 import org.whispersystems.textsecuregcm.push.FcmSender;
-import org.whispersystems.textsecuregcm.push.WebSocketConnectionEventManager;
 import org.whispersystems.textsecuregcm.push.PushNotificationManager;
 import org.whispersystems.textsecuregcm.push.PushNotificationScheduler;
+import org.whispersystems.textsecuregcm.push.WebSocketConnectionEventManager;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClient;
 import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClusterClient;
 import org.whispersystems.textsecuregcm.securestorage.SecureStorageClient;
@@ -52,6 +52,7 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.ClientPublicKeys;
 import org.whispersystems.textsecuregcm.storage.ClientPublicKeysManager;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
+import org.whispersystems.textsecuregcm.storage.DynamoDbRecoveryManager;
 import org.whispersystems.textsecuregcm.storage.IssuedReceiptsManager;
 import org.whispersystems.textsecuregcm.storage.KeysManager;
 import org.whispersystems.textsecuregcm.storage.MessagesCache;
@@ -77,7 +78,6 @@ import software.amazon.awssdk.services.s3.S3AsyncClient;
  * Construct utilities commonly used by worker commands
  */
 record CommandDependencies(
-    Accounts accounts,
     AccountsManager accountsManager,
     ProfilesManager profilesManager,
     ReportMessageManager reportMessageManager,
@@ -97,7 +97,8 @@ record CommandDependencies(
     IssuedReceiptsManager issuedReceiptsManager,
     DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
     DynamoDbAsyncClient dynamoDbAsyncClient,
-    PhoneNumberIdentifiers phoneNumberIdentifiers) {
+    PhoneNumberIdentifiers phoneNumberIdentifiers,
+    DynamoDbRecoveryManager dynamoDbRecoveryManager) {
 
   static CommandDependencies build(
       final String name,
@@ -294,13 +295,15 @@ record CommandDependencies(
     WebSocketConnectionEventManager webSocketConnectionEventManager =
         new WebSocketConnectionEventManager(accountsManager, pushNotificationManager, messagesCluster, clientEventExecutor, asyncOperationQueueingExecutor);
 
+    final DynamoDbRecoveryManager dynamoDbRecoveryManager =
+        new DynamoDbRecoveryManager(accounts, phoneNumberIdentifiers);
+
     environment.lifecycle().manage(apnSender);
     environment.lifecycle().manage(disconnectionRequestManager);
     environment.lifecycle().manage(webSocketConnectionEventManager);
     environment.lifecycle().manage(new ManagedAwsCrt());
 
     return new CommandDependencies(
-        accounts,
         accountsManager,
         profilesManager,
         reportMessageManager,
@@ -320,7 +323,8 @@ record CommandDependencies(
         issuedReceiptsManager,
         dynamicConfigurationManager,
         dynamoDbAsyncClient,
-        phoneNumberIdentifiers
+        phoneNumberIdentifiers,
+        dynamoDbRecoveryManager
     );
   }
 

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

@@ -12,12 +12,12 @@ 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 org.whispersystems.textsecuregcm.storage.DynamoDbRecoveryManager;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.util.retry.Retry;
 
-public class RegenerateAccountConstraintDataCommand extends AbstractSinglePassCrawlAccountsCommand {
+public class RegenerateSecondaryDynamoDbTableDataCommand extends AbstractSinglePassCrawlAccountsCommand {
 
   @VisibleForTesting
   static final String DRY_RUN_ARGUMENT = "dry-run";
@@ -29,10 +29,10 @@ public class RegenerateAccountConstraintDataCommand extends AbstractSinglePassCr
   static final String RETRIES_ARGUMENT = "retries";
 
   private static final String PROCESSED_ACCOUNTS_COUNTER_NAME =
-      MetricsUtil.name(RegenerateAccountConstraintDataCommand.class, "processedAccounts");
+      MetricsUtil.name(RegenerateSecondaryDynamoDbTableDataCommand.class, "processedAccounts");
 
-  public RegenerateAccountConstraintDataCommand() {
+  public RegenerateSecondaryDynamoDbTableDataCommand() {
-    super("regenerate-account-constraint-data", "Regenerates account constraint data from a core account table");
+    super("regenerate-secondary-dynamodb-table-data", "Regenerates secondary DynamoDB table data from core tables");
   }
 
   @Override
@@ -65,7 +65,7 @@ public class RegenerateAccountConstraintDataCommand extends AbstractSinglePassCr
     final int maxConcurrency = getNamespace().getInt(MAX_CONCURRENCY_ARGUMENT);
     final int maxRetries = getNamespace().getInt(RETRIES_ARGUMENT);
 
-    final Accounts accounts = getCommandDependencies().accounts();
+    final DynamoDbRecoveryManager dynamoDbRecoveryManager = getCommandDependencies().dynamoDbRecoveryManager();
 
     final Counter processedAccountsCounter = Metrics.counter(PROCESSED_ACCOUNTS_COUNTER_NAME,
         "dryRun", String.valueOf(dryRun));
@@ -74,7 +74,7 @@ public class RegenerateAccountConstraintDataCommand extends AbstractSinglePassCr
         .doOnNext(ignored -> processedAccountsCounter.increment())
         .flatMap(account -> dryRun
                 ? Mono.empty()
-                : Mono.fromFuture(() -> accounts.regenerateConstraints(account))
+                : Mono.fromFuture(() -> dynamoDbRecoveryManager.regenerateData(account))
                     .retryWhen(Retry.backoff(maxRetries, Duration.ofSeconds(1)).maxBackoff(Duration.ofSeconds(4))
                         .onRetryExhaustedThrow((spec, rs) -> rs.failure())),
             maxConcurrency)

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

@@ -8,12 +8,16 @@ package org.whispersystems.textsecuregcm.storage;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import com.google.i18n.phonenumbers.PhoneNumberUtil;
 import java.io.IOException;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -21,6 +25,7 @@ import java.util.function.Supplier;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema.Tables;
 import org.whispersystems.textsecuregcm.util.CompletableFutureTestUtil;
 import software.amazon.awssdk.services.dynamodb.model.TransactionCanceledException;
@@ -200,4 +205,25 @@ class PhoneNumberIdentifiersTest {
     final UUID pni = phoneNumberIdentifiers.getPhoneNumberIdentifier(number).join();
     assertEquals(List.of(number), phoneNumberIdentifiers.getPhoneNumber(pni).join());
   }
+
+  @Test
+  void regeneratePhoneNumberIdentifierMappings() {
+    // libphonenumber 8.13.50 and on generate new-format numbers for Benin
+    final String newFormatBeninE164 = PhoneNumberUtil.getInstance()
+        .format(PhoneNumberUtil.getInstance().getExampleNumber("BJ"), PhoneNumberUtil.PhoneNumberFormat.E164);
+    final String oldFormatBeninE164 = newFormatBeninE164.replaceFirst("01", "");
+
+    final UUID phoneNumberIdentifier = UUID.randomUUID();
+
+    final Account account = mock(Account.class);
+    when(account.getNumber()).thenReturn(newFormatBeninE164);
+    when(account.getIdentifier(IdentityType.PNI)).thenReturn(phoneNumberIdentifier);
+
+    phoneNumberIdentifiers.regeneratePhoneNumberIdentifierMappings(account).join();
+
+    assertEquals(phoneNumberIdentifier, phoneNumberIdentifiers.getPhoneNumberIdentifier(newFormatBeninE164).join());
+    assertEquals(phoneNumberIdentifier, phoneNumberIdentifiers.getPhoneNumberIdentifier(oldFormatBeninE164).join());
+    assertEquals(Set.of(newFormatBeninE164, oldFormatBeninE164),
+        new HashSet<>(phoneNumberIdentifiers.getPhoneNumber(phoneNumberIdentifier).join()));
+  }
 }

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

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

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

@@ -51,7 +51,6 @@ class NotifyIdleDevicesCommandTest {
           null,
           null,
           null,
-          null,
           messagesManager,
           null,
           null,
@@ -66,6 +65,7 @@ class NotifyIdleDevicesCommandTest {
           null,
           null,
           null,
+          null,
           null);
 
       this.idleDeviceNotificationScheduler = idleDeviceNotificationScheduler;

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

@@ -19,19 +19,20 @@ 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 org.whispersystems.textsecuregcm.storage.DynamoDbRecoveryManager;
 import reactor.core.publisher.Flux;
 
-class RegenerateAccountConstraintDataCommandTest {
+class RegenerateSecondaryDynamoDbTableDataCommandTest {
 
-  private Accounts accounts;
+  private DynamoDbRecoveryManager dynamoDbRecoveryManager;
 
-  private static class TestRegenerateAccountConstraintDataCommand extends RegenerateAccountConstraintDataCommand {
+  private static class TestRegenerateSecondaryDynamoDbTableDataCommand extends RegenerateSecondaryDynamoDbTableDataCommand {
 
     private final CommandDependencies commandDependencies;
     private final Namespace namespace;
 
-    TestRegenerateAccountConstraintDataCommand(final Accounts accounts, final boolean dryRun) {
+    TestRegenerateSecondaryDynamoDbTableDataCommand(final DynamoDbRecoveryManager dynamoDbRecoveryManager, final boolean dryRun) {
-      commandDependencies = new CommandDependencies(accounts,
+      commandDependencies = new CommandDependencies(null,
           null,
           null,
           null,
@@ -51,12 +52,12 @@ class RegenerateAccountConstraintDataCommandTest {
           null,
           null,
           null,
-          null);
+          dynamoDbRecoveryManager);
 
       namespace = new Namespace(Map.of(
-          RegenerateAccountConstraintDataCommand.DRY_RUN_ARGUMENT, dryRun,
+          RegenerateSecondaryDynamoDbTableDataCommand.DRY_RUN_ARGUMENT, dryRun,
-          RegenerateAccountConstraintDataCommand.MAX_CONCURRENCY_ARGUMENT, 16,
+          RegenerateSecondaryDynamoDbTableDataCommand.MAX_CONCURRENCY_ARGUMENT, 16,
-          RegenerateAccountConstraintDataCommand.RETRIES_ARGUMENT, 3));
+          RegenerateSecondaryDynamoDbTableDataCommand.RETRIES_ARGUMENT, 3));
     }
 
     @Override
@@ -72,9 +73,9 @@ class RegenerateAccountConstraintDataCommandTest {
 
   @BeforeEach
   void setUp() {
-    accounts = mock(Accounts.class);
+    dynamoDbRecoveryManager = mock(DynamoDbRecoveryManager.class);
 
-    when(accounts.regenerateConstraints(any())).thenReturn(CompletableFuture.completedFuture(null));
+    when(dynamoDbRecoveryManager.regenerateData(any())).thenReturn(CompletableFuture.completedFuture(null));
   }
 
   @ParameterizedTest
@@ -82,15 +83,15 @@ class RegenerateAccountConstraintDataCommandTest {
   void crawlAccounts(final boolean dryRun) {
     final Account account = mock(Account.class);
 
-    final RegenerateAccountConstraintDataCommand regenerateAccountConstraintDataCommand =
+    final RegenerateSecondaryDynamoDbTableDataCommand regenerateSecondaryDynamoDbTableDataCommand =
-        new TestRegenerateAccountConstraintDataCommand(accounts, dryRun);
+        new TestRegenerateSecondaryDynamoDbTableDataCommand(dynamoDbRecoveryManager, dryRun);
 
-    regenerateAccountConstraintDataCommand.crawlAccounts(Flux.just(account));
+    regenerateSecondaryDynamoDbTableDataCommand.crawlAccounts(Flux.just(account));
 
     if (!dryRun) {
-      verify(accounts).regenerateConstraints(account);
+      verify(dynamoDbRecoveryManager).regenerateData(account);
     }
 
-    verifyNoMoreInteractions(accounts);
+    verifyNoMoreInteractions(dynamoDbRecoveryManager);
   }
 }

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

@@ -64,7 +64,6 @@ class StartPushNotificationExperimentCommandTest {
           null,
           null,
           null,
-          null,
           pushNotificationExperimentSamples,
           null,
           null,
@@ -73,6 +72,7 @@ class StartPushNotificationExperimentCommandTest {
           null,
           null,
           null,
+          null,
           null);
     }