Remove migration paths for lazy message deletion

2026-04-20 11:18:04 +01:00 · 2024-07-23 14:07:19 -07:00
parent 6eed458ceb
commit f12a6ff73f
7 changed files with 34 additions and 345 deletions
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -397,7 +397,6 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
    MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
        config.getDynamoDbTables().getMessages().getTableName(),
        config.getDynamoDbTables().getMessages().getExpiration(),
-        dynamicConfigurationManager,
        messageDeletionAsyncExecutor);
    RemoteConfigs remoteConfigs = new RemoteConfigs(dynamoDbClient,
        config.getDynamoDbTables().getRemoteConfig().getTableName());
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesDynamoDb.java
@@ -29,8 +29,6 @@ import java.util.function.Predicate;
 import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
-import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicMessagesConfiguration.DynamoKeyScheme;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
 import org.whispersystems.textsecuregcm.util.AttributeValues;
 import reactor.core.publisher.Flux;
@@ -69,27 +67,22 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
  private final Timer storeTimer = timer(name(getClass(), "store"));
  private final String DELETE_BY_ACCOUNT_TIMER_NAME = name(getClass(), "delete", "account");
  private final String DELETE_BY_DEVICE_TIMER_NAME = name(getClass(), "delete", "device");
-  private final String MESSAGES_STORED_BY_SCHEME_COUNTER_NAME = name(getClass(), "messagesStored");
-  private final String MESSAGES_LOADED_BY_SCHEME_COUNTER_NAME = name(getClass(), "messagesLoaded");
-  private final String MESSAGES_DELETED_BY_SCHEME_COUNTER_NAME = name(getClass(), "messagesDeleted");

  private final DynamoDbAsyncClient dbAsyncClient;
  private final String tableName;
  private final Duration timeToLive;
-  private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfig;
  private final ExecutorService messageDeletionExecutor;
  private final Scheduler messageDeletionScheduler;

  private static final Logger logger = LoggerFactory.getLogger(MessagesDynamoDb.class);

  public MessagesDynamoDb(DynamoDbClient dynamoDb, DynamoDbAsyncClient dynamoDbAsyncClient, String tableName,
-      Duration timeToLive, DynamicConfigurationManager<DynamicConfiguration> dynamicConfig, ExecutorService messageDeletionExecutor) {
+      Duration timeToLive, ExecutorService messageDeletionExecutor) {
    super(dynamoDb);

    this.dbAsyncClient = dynamoDbAsyncClient;
    this.tableName = tableName;
    this.timeToLive = timeToLive;
-    this.dynamicConfig = dynamicConfig;

    this.messageDeletionExecutor = messageDeletionExecutor;
    this.messageDeletionScheduler = Schedulers.fromExecutor(messageDeletionExecutor);
@@ -102,20 +95,18 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {

  private void storeBatch(final List<MessageProtos.Envelope> messages, final UUID destinationAccountUuid,
      final Device destinationDevice) {
-    final byte destinationDeviceId = destinationDevice.getId();
    if (messages.size() > DYNAMO_DB_MAX_BATCH_SIZE) {
      throw new IllegalArgumentException("Maximum batch size of " + DYNAMO_DB_MAX_BATCH_SIZE + " exceeded with " + messages.size() + " messages");
    }

-    final DynamoKeyScheme scheme = dynamicConfig.getConfiguration().getMessagesConfiguration().writeKeyScheme();
-    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, destinationDevice, scheme);
+    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, destinationDevice);
    List<WriteRequest> writeItems = new ArrayList<>();
    for (MessageProtos.Envelope message : messages) {
      final UUID messageUuid = UUID.fromString(message.getServerGuid());

      final ImmutableMap.Builder<String, AttributeValue> item = ImmutableMap.<String, AttributeValue>builder()
          .put(KEY_PARTITION, partitionKey)
-          .put(KEY_SORT, convertSortKey(destinationDevice.getId(), message.getServerTimestamp(), messageUuid, scheme))
+          .put(KEY_SORT, convertSortKey(destinationDevice.getId(), message.getServerTimestamp(), messageUuid))
          .put(LOCAL_INDEX_MESSAGE_UUID_KEY_SORT, convertLocalIndexMessageUuidSortKey(messageUuid))
          .put(KEY_TTL, AttributeValues.fromLong(getTtlForMessage(message)))
          .put(KEY_ENVELOPE_BYTES, AttributeValue.builder().b(SdkBytes.fromByteArray(message.toByteArray())).build());
@@ -126,77 +117,31 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
    }

    executeTableWriteItemsUntilComplete(Map.of(tableName, writeItems));
-    Metrics.counter(MESSAGES_STORED_BY_SCHEME_COUNTER_NAME, Tags.of("scheme", scheme.name())).increment(writeItems.size());
  }

  public CompletableFuture<Boolean> mayHaveMessages(final UUID accountIdentifier, final Device device) {
-    return Flux.fromIterable(dynamicConfig.getConfiguration().getMessagesConfiguration().dynamoKeySchemes())
-        .flatMap(scheme -> mayHaveMessages(accountIdentifier, device, scheme))
-        .any(mayHaveMessages -> mayHaveMessages)
-        .toFuture();
-  }
-
-  private Mono<Boolean> mayHaveMessages(final UUID accountIdentifier, final Device device, final DynamoKeyScheme scheme) {
-    final AttributeValue partitionKey = convertPartitionKey(accountIdentifier, device, scheme);
+    final AttributeValue partitionKey = convertPartitionKey(accountIdentifier, device);

    QueryRequest.Builder queryRequestBuilder = QueryRequest.builder()
        .tableName(tableName)
        .consistentRead(false)
-        .limit(1);
+        .limit(1)
+        .keyConditionExpression("#part = :part")
+        .expressionAttributeNames(Map.of("#part", KEY_PARTITION))
+        .expressionAttributeValues(Map.of(":part", partitionKey));

-    queryRequestBuilder = switch (scheme) {
-      case TRADITIONAL -> queryRequestBuilder
-          .keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
-          .expressionAttributeNames(Map.of(
-              "#part", KEY_PARTITION,
-              "#sort", KEY_SORT))
-          .expressionAttributeValues(Map.of(
-              ":part", partitionKey,
-              ":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(device.getId(), scheme)));
-      case LAZY_DELETION -> queryRequestBuilder
-          .keyConditionExpression("#part = :part")
-          .expressionAttributeNames(Map.of("#part", KEY_PARTITION))
-          .expressionAttributeValues(Map.of(":part", partitionKey));
-    };
-
-    return Mono.fromFuture(dbAsyncClient.query(queryRequestBuilder.build())
-        .thenApply(queryResponse -> queryResponse.count() > 0));
+    return dbAsyncClient.query(queryRequestBuilder.build())
+        .thenApply(queryResponse -> queryResponse.count() > 0);
  }

  public Publisher<MessageProtos.Envelope> load(final UUID destinationAccountUuid, final Device device, final Integer limit) {
-    return Flux.concat(
-        dynamicConfig.getConfiguration().getMessagesConfiguration().dynamoKeySchemes()
-            .stream()
-            .map(scheme -> load(destinationAccountUuid, device, limit, scheme))
-            .toList())
-        .map(messageAndScheme -> {
-              Metrics.counter(MESSAGES_LOADED_BY_SCHEME_COUNTER_NAME, Tags.of("scheme", messageAndScheme.getT2().name())).increment();
-              return messageAndScheme.getT1();
-            });
-  }
-
-  private Publisher<Tuple2<MessageProtos.Envelope, DynamoKeyScheme>> load(final UUID destinationAccountUuid, final Device device, final Integer limit, final DynamoKeyScheme scheme) {
-    final byte destinationDeviceId = device.getId();
-
-    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, device, scheme);
+    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, device);
    QueryRequest.Builder queryRequestBuilder = QueryRequest.builder()
        .tableName(tableName)
-        .consistentRead(true);
-
-    queryRequestBuilder = switch (scheme) {
-      case TRADITIONAL -> queryRequestBuilder
-          .keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
-          .expressionAttributeNames(Map.of(
-              "#part", KEY_PARTITION,
-              "#sort", KEY_SORT))
-          .expressionAttributeValues(Map.of(
-              ":part", partitionKey,
-              ":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId, scheme)));
-      case LAZY_DELETION -> queryRequestBuilder
-          .keyConditionExpression("#part = :part")
-          .expressionAttributeNames(Map.of("#part", KEY_PARTITION))
-          .expressionAttributeValues(Map.of(":part", partitionKey));
-    };
+        .consistentRead(true)
+        .keyConditionExpression("#part = :part")
+        .expressionAttributeNames(Map.of("#part", KEY_PARTITION))
+        .expressionAttributeValues(Map.of(":part", partitionKey));

    if (limit != null) {
      // some callers don’t take advantage of reactive streams, so we want to support limiting the fetch size. Otherwise,
@@ -215,25 +160,12 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
            return null;
          }
        })
-        .filter(Predicate.not(Objects::isNull))
-        .map(m -> Tuples.of(m, scheme));
+        .filter(Predicate.not(Objects::isNull));
  }

  public CompletableFuture<Optional<MessageProtos.Envelope>> deleteMessageByDestinationAndGuid(
      final UUID destinationAccountUuid, final Device destinationDevice, final UUID messageUuid) {
-    return dynamicConfig.getConfiguration().getMessagesConfiguration().dynamoKeySchemes()
-        .stream()
-        .map(scheme -> deleteMessageByDestinationAndGuid(destinationAccountUuid, destinationDevice, messageUuid, scheme))
-        // this combines the futures by producing a future that returns an arbitrary nonempty
-        // result if there is one, which should be OK because only one of the keying schemes
-        // should produce a nonempty result for any given message uuid
-        .reduce((f, g) -> f.thenCombine(g, (a, b) -> a.or(() -> b)))
-        .get();                 // there is always at least one scheme
-  }
-
-  private CompletableFuture<Optional<MessageProtos.Envelope>> deleteMessageByDestinationAndGuid(
-      final UUID destinationAccountUuid, final Device destinationDevice, final UUID messageUuid, DynamoKeyScheme scheme) {
-    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, destinationDevice, scheme);
+    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, destinationDevice);
    final QueryRequest queryRequest = QueryRequest.builder()
        .tableName(tableName)
        .indexName(LOCAL_INDEX_MESSAGE_UUID_NAME)
@@ -260,7 +192,6 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
        .mapNotNull(deleteItemResponse -> {
          try {
            if (deleteItemResponse.attributes() != null && deleteItemResponse.attributes().containsKey(KEY_PARTITION)) {
-              Metrics.counter(MESSAGES_DELETED_BY_SCHEME_COUNTER_NAME, Tags.of("scheme", scheme.name())).increment();
              return convertItemToEnvelope(deleteItemResponse.attributes());
            }
          } catch (final InvalidProtocolBufferException e) {
@@ -276,20 +207,8 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {

  public CompletableFuture<Optional<MessageProtos.Envelope>> deleteMessage(final UUID destinationAccountUuid,
      final Device destinationDevice, final UUID messageUuid, final long serverTimestamp) {
-    return dynamicConfig.getConfiguration().getMessagesConfiguration().dynamoKeySchemes()
-        .stream()
-        .map(scheme -> deleteMessage(destinationAccountUuid, destinationDevice, messageUuid, serverTimestamp, scheme))
-        // this combines the futures by producing a future that returns an arbitrary nonempty
-        // result if there is one, which should be OK because only one of the keying schemes
-        // should produce a nonempty result for any given message uuid
-        .reduce((f, g) -> f.thenCombine(g, (a, b) -> a.or(() -> b)))
-        .orElseThrow();         // there is always at least one scheme
-  }
-
-  private CompletableFuture<Optional<MessageProtos.Envelope>> deleteMessage(final UUID destinationAccountUuid,
-      final Device destinationDevice, final UUID messageUuid, final long serverTimestamp, final DynamoKeyScheme scheme) {
-    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, destinationDevice, scheme);
-    final AttributeValue sortKey = convertSortKey(destinationDevice.getId(), serverTimestamp, messageUuid, scheme);
+    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, destinationDevice);
+    final AttributeValue sortKey = convertSortKey(destinationDevice.getId(), serverTimestamp, messageUuid);
    DeleteItemRequest.Builder deleteItemRequest = DeleteItemRequest.builder()
        .tableName(tableName)
        .key(Map.of(KEY_PARTITION, partitionKey, KEY_SORT, sortKey))
@@ -299,7 +218,6 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
        .thenApplyAsync(deleteItemResponse -> {
          if (deleteItemResponse.attributes() != null && deleteItemResponse.attributes().containsKey(KEY_PARTITION)) {
            try {
-              Metrics.counter(MESSAGES_DELETED_BY_SCHEME_COUNTER_NAME, Tags.of("scheme", scheme.name())).increment();
              return Optional.of(convertItemToEnvelope(deleteItemResponse.attributes()));
            } catch (final InvalidProtocolBufferException e) {
              logger.error("Failed to parse envelope", e);
@@ -310,69 +228,6 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
        }, messageDeletionExecutor);
  }

-  // Deletes all messages stored for the supplied account that were stored under the traditional (uuid+device id) keying scheme.
-  // Messages stored under the lazy-message-deletion keying scheme will not be affected.
-  public CompletableFuture<Void> deleteAllMessagesForAccount(final UUID destinationAccountUuid) {
-    final Timer.Sample sample = Timer.start();
-
-    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, null, DynamoKeyScheme.TRADITIONAL);
-
-    return Flux.from(dbAsyncClient.queryPaginator(QueryRequest.builder()
-            .tableName(tableName)
-            .projectionExpression(KEY_SORT)
-            .consistentRead(true)
-            .keyConditionExpression("#part = :part")
-            .expressionAttributeNames(Map.of("#part", KEY_PARTITION))
-            .expressionAttributeValues(Map.of(":part", partitionKey))
-            .build())
-        .items())
-        .flatMap(item -> Mono.fromFuture(() -> dbAsyncClient.deleteItem(DeleteItemRequest.builder()
-            .tableName(tableName)
-            .key(Map.of(
-                KEY_PARTITION, partitionKey,
-                KEY_SORT, item.get(KEY_SORT)))
-            .build())),
-            DYNAMO_DB_MAX_BATCH_SIZE)
-        .then()
-        .doOnSuccess(ignored -> sample.stop(timer(DELETE_BY_ACCOUNT_TIMER_NAME, "outcome", "success")))
-        .doOnError(ignored -> sample.stop(timer(DELETE_BY_ACCOUNT_TIMER_NAME, "outcome", "error")))
-        .toFuture();
-  }
-
-  // Deletes all messages stored for the supplied account and device that were stored under the
-  // traditional (uuid+device id) keying scheme. Messages stored under the lazy-message-deletion
-  // keying scheme will not be affected.
-  public CompletableFuture<Void> deleteAllMessagesForDevice(final UUID destinationAccountUuid,
-      final byte destinationDeviceId) {
-    final Timer.Sample sample = Timer.start();
-    final AttributeValue partitionKey = convertPartitionKey(destinationAccountUuid, null, DynamoKeyScheme.TRADITIONAL);
-
-    return Flux.from(dbAsyncClient.queryPaginator(QueryRequest.builder()
-                .tableName(tableName)
-                .keyConditionExpression("#part = :part AND begins_with ( #sort , :sortprefix )")
-                .expressionAttributeNames(Map.of(
-                    "#part", KEY_PARTITION,
-                    "#sort", KEY_SORT))
-                .expressionAttributeValues(Map.of(
-                    ":part", partitionKey,
-                    ":sortprefix", convertDestinationDeviceIdToSortKeyPrefix(destinationDeviceId, DynamoKeyScheme.TRADITIONAL)))
-                .projectionExpression(KEY_SORT)
-                .consistentRead(true)
-                .build())
-            .items())
-        .flatMap(item -> Mono.fromFuture(() -> dbAsyncClient.deleteItem(DeleteItemRequest.builder()
-            .tableName(tableName)
-            .key(Map.of(
-                KEY_PARTITION, partitionKey,
-                KEY_SORT, item.get(KEY_SORT)))
-            .build())),
-            DYNAMO_DB_MAX_BATCH_SIZE)
-        .then()
-        .doOnSuccess(ignored -> sample.stop(timer(DELETE_BY_DEVICE_TIMER_NAME, "outcome", "success")))
-        .doOnError(ignored -> sample.stop(timer(DELETE_BY_DEVICE_TIMER_NAME, "outcome", "error")))
-        .toFuture();
-  }
-
  @VisibleForTesting
  static MessageProtos.Envelope convertItemToEnvelope(final Map<String, AttributeValue> item)
      throws InvalidProtocolBufferException {
@@ -384,40 +239,22 @@ public class MessagesDynamoDb extends AbstractDynamoDbStore {
    return message.getServerTimestamp() / 1000 + timeToLive.getSeconds();
  }

-  private static AttributeValue convertPartitionKey(final UUID destinationAccountUuid, final Device destinationDevice, final DynamoKeyScheme scheme) {
-    return switch (scheme) {
-      case TRADITIONAL -> AttributeValues.fromUUID(destinationAccountUuid);
-      case LAZY_DELETION -> {
-          final ByteBuffer byteBuffer = ByteBuffer.allocate(24);
-          byteBuffer.putLong(destinationAccountUuid.getMostSignificantBits());
-          byteBuffer.putLong(destinationAccountUuid.getLeastSignificantBits());
-          byteBuffer.putLong(destinationDevice.getCreated() & ~0x7f + destinationDevice.getId());
-          yield AttributeValues.fromByteBuffer(byteBuffer.flip());
-      }
-    };
+  private static AttributeValue convertPartitionKey(final UUID destinationAccountUuid, final Device destinationDevice) {
+    final ByteBuffer byteBuffer = ByteBuffer.allocate(24);
+    byteBuffer.putLong(destinationAccountUuid.getMostSignificantBits());
+    byteBuffer.putLong(destinationAccountUuid.getLeastSignificantBits());
+    byteBuffer.putLong((destinationDevice.getCreated() & ~0x7f) + destinationDevice.getId());
+    return AttributeValues.fromByteBuffer(byteBuffer.flip());
  }

-  private static AttributeValue convertSortKey(final byte destinationDeviceId, final long serverTimestamp,
-      final UUID messageUuid, final DynamoKeyScheme scheme) {
-
-    final ByteBuffer byteBuffer = ByteBuffer.allocate(32);
-    if (scheme == DynamoKeyScheme.TRADITIONAL) {
-        // for compatibility - destinationDeviceId was previously `long`
-        byteBuffer.putLong(destinationDeviceId);
-    }
+  private static AttributeValue convertSortKey(final byte destinationDeviceId, final long serverTimestamp, final UUID messageUuid) {
+    final ByteBuffer byteBuffer = ByteBuffer.allocate(24);
    byteBuffer.putLong(serverTimestamp);
    byteBuffer.putLong(messageUuid.getMostSignificantBits());
    byteBuffer.putLong(messageUuid.getLeastSignificantBits());
    return AttributeValues.fromByteBuffer(byteBuffer.flip());
  }

-  private static AttributeValue convertDestinationDeviceIdToSortKeyPrefix(final byte destinationDeviceId, final DynamoKeyScheme scheme) {
-    return switch (scheme) {
-      case TRADITIONAL -> AttributeValues.fromByteBuffer(ByteBuffer.allocate(8).putLong(destinationDeviceId).flip());
-      case LAZY_DELETION -> AttributeValues.b(new byte[0]);
-    };
-  }
-
  private static AttributeValue convertLocalIndexMessageUuidSortKey(final UUID messageUuid) {
    return AttributeValues.fromUUID(messageUuid);
  }
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
@@ -107,15 +107,11 @@ public class MessagesManager {
  }

  public CompletableFuture<Void> clear(UUID destinationUuid) {
-    return CompletableFuture.allOf(
-        messagesCache.clear(destinationUuid),
-        messagesDynamoDb.deleteAllMessagesForAccount(destinationUuid));
+    return messagesCache.clear(destinationUuid);
  }

  public CompletableFuture<Void> clear(UUID destinationUuid, byte deviceId) {
-    return CompletableFuture.allOf(
-        messagesCache.clear(destinationUuid, deviceId),
-        messagesDynamoDb.deleteAllMessagesForDevice(destinationUuid, deviceId));
+    return messagesCache.clear(destinationUuid, deviceId);
  }

  public CompletableFuture<Optional<Envelope>> delete(UUID destinationUuid, Device destinationDevice, UUID guid,
--- a/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java
@@ -169,7 +169,6 @@ record CommandDependencies(
    MessagesDynamoDb messagesDynamoDb = new MessagesDynamoDb(dynamoDbClient, dynamoDbAsyncClient,
        configuration.getDynamoDbTables().getMessages().getTableName(),
        configuration.getDynamoDbTables().getMessages().getExpiration(),
-        dynamicConfigurationManager,
        messageDeletionExecutor);
    FaultTolerantRedisCluster messagesCluster = configuration.getMessageCacheConfiguration()
        .getRedisClusterConfiguration().build("messages", redisClientResourcesBuilder);