Add MRM views experiment to MessagesCache.getMessagesToPersist()

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

@@ -15,22 +15,14 @@ import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Timer;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.TimeUnit;
-import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.entities.MessageProtos;
-import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
 import org.whispersystems.textsecuregcm.util.Util;
-import reactor.core.publisher.Flux;
-import reactor.util.function.Tuple2;
-import reactor.util.function.Tuples;
 import software.amazon.awssdk.services.dynamodb.model.ItemCollectionSizeLimitExceededException;
 
 public class MessagePersister implements Managed {
@@ -38,8 +30,6 @@ public class MessagePersister implements Managed {
   private final MessagesCache messagesCache;
   private final MessagesManager messagesManager;
   private final AccountsManager accountsManager;
-  private final ClientPresenceManager clientPresenceManager;
-  private final KeysManager keysManager;
 
   private final Duration persistDelay;
 
@@ -72,17 +62,14 @@ public class MessagePersister implements Managed {
   private static final Logger logger = LoggerFactory.getLogger(MessagePersister.class);
 
   public MessagePersister(final MessagesCache messagesCache, final MessagesManager messagesManager,
-      final AccountsManager accountsManager, final ClientPresenceManager clientPresenceManager,
-      final KeysManager keysManager,
-      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
-      final Duration persistDelay,
+      final AccountsManager accountsManager,
+      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager, final Duration persistDelay,
       final int dedicatedProcessWorkerThreadCount
   ) {
+
     this.messagesCache = messagesCache;
     this.messagesManager = messagesManager;
     this.accountsManager = accountsManager;
-    this.clientPresenceManager = clientPresenceManager;
-    this.keysManager = keysManager;
     this.persistDelay = persistDelay;
     this.workerThreads = new Thread[dedicatedProcessWorkerThreadCount];
 

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

@@ -11,7 +11,6 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.InvalidProtocolBufferException;
 import io.dropwizard.lifecycle.Managed;
-import io.lettuce.core.ScoredValue;
 import io.lettuce.core.ZAddArgs;
 import io.lettuce.core.cluster.SlotHash;
 import io.lettuce.core.cluster.models.partitions.RedisClusterNode;
@@ -493,8 +492,7 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
         .subscribe();
   }
 
-  private Mono<Pair<List<byte[]>, Long>> getNextMessagePage(final UUID destinationUuid,
-      final byte destinationDevice,
+  private Mono<Pair<List<byte[]>, Long>> getNextMessagePage(final UUID destinationUuid, final byte destinationDevice,
       long messageId) {
 
     return getItemsScript.execute(destinationUuid, destinationDevice, PAGE_SIZE, messageId)
@@ -520,22 +518,40 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
   @VisibleForTesting
   List<MessageProtos.Envelope> getMessagesToPersist(final UUID accountUuid, final byte destinationDevice,
       final int limit) {
-    return getMessagesTimer.record(() -> {
-      final List<ScoredValue<byte[]>> scoredMessages = redisCluster.withBinaryCluster(
-          connection -> connection.sync()
-              .zrangeWithScores(getMessageQueueKey(accountUuid, destinationDevice), 0, limit));
-      final List<MessageProtos.Envelope> envelopes = new ArrayList<>(scoredMessages.size());
 
-      for (final ScoredValue<byte[]> scoredMessage : scoredMessages) {
-        try {
-          envelopes.add(MessageProtos.Envelope.parseFrom(scoredMessage.getValue()));
-        } catch (InvalidProtocolBufferException e) {
-          logger.warn("Failed to parse envelope", e);
-        }
-      }
+    final Timer.Sample sample = Timer.start();
 
-      return envelopes;
-    });
+    final List<byte[]> messages = redisCluster.withBinaryCluster(connection ->
+        connection.sync().zrange(getMessageQueueKey(accountUuid, destinationDevice), 0, limit));
+
+    return Flux.fromIterable(messages)
+        .mapNotNull(message -> {
+          try {
+            return MessageProtos.Envelope.parseFrom(message);
+          } catch (InvalidProtocolBufferException e) {
+            logger.warn("Failed to parse envelope", e);
+            return null;
+          }
+        })
+        .concatMap(message -> {
+          final Mono<MessageProtos.Envelope> messageMono;
+          if (message.hasSharedMrmKey()) {
+            final Mono<?> experimentMono = maybeRunMrmViewExperiment(message, accountUuid, destinationDevice);
+
+            // mrm views phase 1: messageMono for sharedMrmKey is always Mono.just(), because messages always have content
+            // To avoid races, wait for the experiment to run, but ignore any errors
+            messageMono = experimentMono
+                .onErrorComplete()
+                .then(Mono.just(message.toBuilder().clearSharedMrmKey().build()));
+          } else {
+            messageMono = Mono.just(message);
+          }
+
+          return messageMono;
+        })
+        .collectList()
+        .doOnTerminate(() -> sample.stop(getMessagesTimer))
+        .block(Duration.ofSeconds(5));
   }
 
   public CompletableFuture<Void> clear(final UUID destinationUuid) {

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

@@ -61,10 +61,7 @@ public class MessagePersisterServiceCommand extends ServerCommand<WhisperServerC
     }
 
     final MessagePersister messagePersister = new MessagePersister(deps.messagesCache(), deps.messagesManager(),
-        deps.accountsManager(),
-        deps.clientPresenceManager(),
-        deps.keysManager(),
-        deps.dynamicConfigurationManager(),
+        deps.accountsManager(), deps.dynamicConfigurationManager(),
         Duration.ofMinutes(configuration.getMessageCacheConfiguration().getPersistDelayMinutes()),
         namespace.getInt(WORKER_COUNT));