From f7f870fe624b295bd0333e304cbfd1140f1d2413 Mon Sep 17 00:00:00 2001
From: Ehren Kret <ehren@signal.org>
Date: Thu, 12 Aug 2021 10:25:09 -0500
Subject: [PATCH] Execute send multi-recipient message loop in parallel

---
 .../controllers/MessageController.java        | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
index d94d8e043..fec4684a7 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java
@@ -17,6 +17,7 @@ import com.google.protobuf.ByteString;
 import io.dropwizard.auth.Auth;
 import io.dropwizard.util.DataSize;
 import io.lettuce.core.ScriptOutputType;
+import io.micrometer.core.instrument.Counter;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Tag;
 import io.micrometer.core.instrument.Tags;
@@ -361,7 +362,7 @@ public class MessageController {
     Map<UUID, Account> uuidToAccountMap = Arrays.stream(multiRecipientMessage.getRecipients())
         .map(Recipient::getUuid)
         .distinct()
-        .collect(Collectors.toMap(Function.identity(), uuid -> {
+        .collect(Collectors.toUnmodifiableMap(Function.identity(), uuid -> {
           Optional<Account> account = accountsManager.get(uuid);
           if (account.isEmpty()) {
             throw new WebApplicationException(Status.NOT_FOUND);
@@ -417,21 +418,21 @@ public class MessageController {
         UserAgentTagUtil.getPlatformTag(userAgent),
         Tag.of(EPHEMERAL_TAG_NAME, String.valueOf(online)),
         Tag.of(SENDER_TYPE_TAG_NAME, "unidentified"));
-    List<UUID> uuids404 = new ArrayList<>();
-    for (Recipient recipient : multiRecipientMessage.getRecipients()) {
-
+    List<UUID> uuids404 = Collections.synchronizedList(new ArrayList<>());
+    final Counter counter = Metrics.counter(SENT_MESSAGE_COUNTER_NAME, tags);
+    Arrays.stream(multiRecipientMessage.getRecipients()).parallel().forEach(recipient -> {
       Account destinationAccount = uuidToAccountMap.get(recipient.getUuid());
-      // we asserted this must be true in validateCompleteDeviceList
-      //noinspection OptionalGetWithoutIsPresent
-      Device destinationDevice = destinationAccount.getDevice(recipient.getDeviceId()).get();
-      Metrics.counter(SENT_MESSAGE_COUNTER_NAME, tags).increment();
+
+      // we asserted this must exist in validateCompleteDeviceList
+      Device destinationDevice = destinationAccount.getDevice(recipient.getDeviceId()).orElseThrow();
+      counter.increment();
       try {
         sendMessage(destinationAccount, destinationDevice, timestamp, online, recipient,
             multiRecipientMessage.getCommonPayload());
       } catch (NoSuchUserException e) {
         uuids404.add(destinationAccount.getUuid());
       }
-    }
+    });
     return Response.ok(new SendMultiRecipientMessageResponse(uuids404)).build();
   }