From 819d59cd7924c7a022a55e13c75e5836de28d4e2 Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Thu, 24 Jun 2021 18:52:26 -0500
Subject: [PATCH] Update reconciliation crawler to use secondary index

---
 .../WhisperServerConfiguration.java           |  5 ++-
 .../textsecuregcm/WhisperServerService.java   |  2 +-
 .../DeletedAccountsDynamoDbConfiguration.java | 13 ++++++
 .../storage/DeletedAccounts.java              | 42 +++++++++----------
 .../workers/DeleteUserCommand.java            |  2 +-
 .../storage/DeletedAccountsTest.java          | 22 +++++++++-
 6 files changed, 58 insertions(+), 28 deletions(-)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/configuration/DeletedAccountsDynamoDbConfiguration.java

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java
index 40819a925..de63cbb94 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java
@@ -22,6 +22,7 @@ import org.whispersystems.textsecuregcm.configuration.AwsAttachmentsConfiguratio
 import org.whispersystems.textsecuregcm.configuration.CdnConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DatabaseConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DatadogConfiguration;
+import org.whispersystems.textsecuregcm.configuration.DeletedAccountsDynamoDbConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DirectoryConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DonationConfiguration;
 import org.whispersystems.textsecuregcm.configuration.DynamoDbConfiguration;
@@ -160,7 +161,7 @@ public class WhisperServerConfiguration extends Configuration {
   @Valid
   @NotNull
   @JsonProperty
-  private DynamoDbConfiguration deletedAccountsDynamoDb;
+  private DeletedAccountsDynamoDbConfiguration deletedAccountsDynamoDb;
 
   @Valid
   @NotNull
@@ -386,7 +387,7 @@ public class WhisperServerConfiguration extends Configuration {
     return migrationRetryAccountsDynamoDb;
   }
 
-  public DynamoDbConfiguration getDeletedAccountsDynamoDbConfiguration() {
+  public DeletedAccountsDynamoDbConfiguration getDeletedAccountsDynamoDbConfiguration() {
     return deletedAccountsDynamoDb;
   }
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 402e38df0..b3f5e3791 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -355,7 +355,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     DynamoDbClient pendingDevicesDynamoDbClient = DynamoDbFromConfig.client(config.getPendingDevicesDynamoDbConfiguration(),
         software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
 
-    DeletedAccounts deletedAccounts = new DeletedAccounts(deletedAccountsDynamoDbClient, config.getDeletedAccountsDynamoDbConfiguration().getTableName());
+    DeletedAccounts deletedAccounts = new DeletedAccounts(deletedAccountsDynamoDbClient, config.getDeletedAccountsDynamoDbConfiguration().getTableName(), config.getDeletedAccountsDynamoDbConfiguration().getNeedsReconciliationIndexName());
     MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(recentlyDeletedAccountsDynamoDb, config.getMigrationDeletedAccountsDynamoDbConfiguration().getTableName());
     MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(migrationRetryAccountsDynamoDb, config.getMigrationRetryAccountsDynamoDbConfiguration().getTableName());
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/DeletedAccountsDynamoDbConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/DeletedAccountsDynamoDbConfiguration.java
new file mode 100644
index 000000000..5a125fd13
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/DeletedAccountsDynamoDbConfiguration.java
@@ -0,0 +1,13 @@
+package org.whispersystems.textsecuregcm.configuration;
+
+import javax.validation.constraints.NotNull;
+
+public class DeletedAccountsDynamoDbConfiguration extends DynamoDbConfiguration {
+
+  @NotNull
+  private String needsReconciliationIndexName;
+
+  public String getNeedsReconciliationIndexName() {
+    return needsReconciliationIndexName;
+  }
+}
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccounts.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccounts.java
index 129f594e8..8980e25f7 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccounts.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/DeletedAccounts.java
@@ -14,9 +14,8 @@ import org.whispersystems.textsecuregcm.util.AttributeValues;
 import org.whispersystems.textsecuregcm.util.Pair;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
-import software.amazon.awssdk.services.dynamodb.model.PutRequest;
 import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
-import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
+import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
 
 public class DeletedAccounts extends AbstractDynamoDbStore {
 
@@ -24,16 +23,18 @@ public class DeletedAccounts extends AbstractDynamoDbStore {
   static final String KEY_ACCOUNT_E164 = "P";
   static final String ATTR_ACCOUNT_UUID = "U";
   static final String ATTR_EXPIRES = "E";
-  static final String ATTR_RECONCILED_IN_CDS = "R";
+  static final String ATTR_NEEDS_CDS_RECONCILIATION = "R";
 
   static final Duration TIME_TO_LIVE = Duration.ofDays(30);
 
   private final String tableName;
+  private final String needsReconciliationIndexName;
 
-  public DeletedAccounts(final DynamoDbClient dynamoDb, final String tableName) {
+  public DeletedAccounts(final DynamoDbClient dynamoDb, final String tableName, final String needsReconciliationIndexName) {
 
     super(dynamoDb);
     this.tableName = tableName;
+    this.needsReconciliationIndexName = needsReconciliationIndexName;
   }
 
   public void put(UUID uuid, String e164) {
@@ -43,7 +44,7 @@ public class DeletedAccounts extends AbstractDynamoDbStore {
             KEY_ACCOUNT_E164, AttributeValues.fromString(e164),
             ATTR_ACCOUNT_UUID, AttributeValues.fromUUID(uuid),
             ATTR_EXPIRES, AttributeValues.fromLong(Instant.now().plus(TIME_TO_LIVE).getEpochSecond()),
-            ATTR_RECONCILED_IN_CDS, AttributeValues.fromBoolean(false)))
+            ATTR_NEEDS_CDS_RECONCILIATION, AttributeValues.fromInt(1)))
         .build());
   }
 
@@ -51,13 +52,7 @@ public class DeletedAccounts extends AbstractDynamoDbStore {
 
     final ScanRequest scanRequest = ScanRequest.builder()
         .tableName(tableName)
-        .filterExpression("#reconciled = :reconciled")
-        .expressionAttributeNames(Map.of(
-            "#reconciled", ATTR_RECONCILED_IN_CDS
-        ))
-        .expressionAttributeValues(Map.of(
-            ":reconciled", AttributeValues.fromBoolean(false)
-        ))
+        .indexName(needsReconciliationIndexName)
         .limit(max)
         .build();
 
@@ -70,18 +65,19 @@ public class DeletedAccounts extends AbstractDynamoDbStore {
   }
 
   public void markReconciled(final List<String> phoneNumbersReconciled) {
-    writeInBatches(phoneNumbersReconciled, phoneNumbers -> {
 
-      final List<WriteRequest> updates = phoneNumbers.stream()
-          .map(phoneNumber -> WriteRequest.builder()
-              .putRequest(PutRequest.builder().item(
-                  Map.of(KEY_ACCOUNT_E164, AttributeValues.fromString(phoneNumber),
-                      ATTR_RECONCILED_IN_CDS, AttributeValues.fromBoolean(true))
-              ).build()).build()
-          ).collect(Collectors.toList());
-
-      executeTableWriteItemsUntilComplete(Map.of(tableName, updates));
-    });
+    phoneNumbersReconciled.forEach(number -> db().updateItem(
+        UpdateItemRequest.builder()
+            .tableName(tableName)
+            .key(Map.of(
+                KEY_ACCOUNT_E164, AttributeValues.fromString(number)
+            ))
+            .updateExpression("REMOVE #needs_reconciliation")
+            .expressionAttributeNames(Map.of(
+                "#needs_reconciliation", ATTR_NEEDS_CDS_RECONCILIATION
+            ))
+            .build()
+    ));
   }
 
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java b/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java
index c67c2b051..c31e5b4b5 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java
@@ -133,7 +133,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
       DynamoDbClient migrationRetryAccountsDynamoDb = DynamoDbFromConfig.client(configuration.getMigrationRetryAccountsDynamoDbConfiguration(),
           software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider.create());
 
-      DeletedAccounts deletedAccounts = new DeletedAccounts(deletedAccountsDynamoDbClient, configuration.getDeletedAccountsDynamoDbConfiguration().getTableName());
+      DeletedAccounts deletedAccounts = new DeletedAccounts(deletedAccountsDynamoDbClient, configuration.getDeletedAccountsDynamoDbConfiguration().getTableName(), configuration.getDeletedAccountsDynamoDbConfiguration().getNeedsReconciliationIndexName());
       MigrationDeletedAccounts migrationDeletedAccounts = new MigrationDeletedAccounts(migrationDeletedAccountsDynamoDb, configuration.getMigrationDeletedAccountsDynamoDbConfiguration().getTableName());
       MigrationRetryAccounts migrationRetryAccounts = new MigrationRetryAccounts(migrationRetryAccountsDynamoDb, configuration.getMigrationRetryAccountsDynamoDbConfiguration().getTableName());
 
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsTest.java
index 23e354381..05a6d522f 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/DeletedAccountsTest.java
@@ -13,10 +13,18 @@ import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.whispersystems.textsecuregcm.util.Pair;
 import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
+import software.amazon.awssdk.services.dynamodb.model.GlobalSecondaryIndex;
+import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement;
+import software.amazon.awssdk.services.dynamodb.model.KeyType;
+import software.amazon.awssdk.services.dynamodb.model.Projection;
+import software.amazon.awssdk.services.dynamodb.model.ProjectionType;
+import software.amazon.awssdk.services.dynamodb.model.ProvisionedThroughput;
 import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType;
 
 class DeletedAccountsTest {
 
+  private static final String NEEDS_RECONCILIATION_INDEX_NAME = "needs_reconciliation_test";
+
   @RegisterExtension
   static DynamoDbExtension dynamoDbExtension = DynamoDbExtension.builder()
       .tableName("deleted_accounts_test")
@@ -24,13 +32,25 @@ class DeletedAccountsTest {
       .attributeDefinition(AttributeDefinition.builder()
           .attributeName(DeletedAccounts.KEY_ACCOUNT_E164)
           .attributeType(ScalarAttributeType.S).build())
+      .attributeDefinition(AttributeDefinition.builder()
+          .attributeName(DeletedAccounts.ATTR_NEEDS_CDS_RECONCILIATION)
+          .attributeType(ScalarAttributeType.N)
+          .build())
+      .globalSecondaryIndex(GlobalSecondaryIndex.builder()
+          .indexName(NEEDS_RECONCILIATION_INDEX_NAME)
+          .keySchema(KeySchemaElement.builder().attributeName(DeletedAccounts.KEY_ACCOUNT_E164).keyType(KeyType.HASH).build(),
+              KeySchemaElement.builder().attributeName(DeletedAccounts.ATTR_NEEDS_CDS_RECONCILIATION).keyType(KeyType.RANGE).build())
+          .projection(Projection.builder().projectionType(ProjectionType.INCLUDE).nonKeyAttributes(DeletedAccounts.ATTR_ACCOUNT_UUID).build())
+          .provisionedThroughput(ProvisionedThroughput.builder().readCapacityUnits(10L).writeCapacityUnits(10L).build())
+          .build())
       .build();
 
   @Test
   void test() {
 
     final DeletedAccounts deletedAccounts = new DeletedAccounts(dynamoDbExtension.getDynamoDbClient(),
-        dynamoDbExtension.getTableName());
+        dynamoDbExtension.getTableName(),
+        NEEDS_RECONCILIATION_INDEX_NAME);
 
     UUID firstUuid = UUID.randomUUID();
     UUID secondUuid = UUID.randomUUID();