Check story rate limits in parallel

This commit is contained in:
Jon Chambers
2023-12-01 16:01:14 -05:00
committed by Jon Chambers
parent e9708b9259
commit 417d99a17e
2 changed files with 55 additions and 3 deletions

View File

@@ -112,6 +112,7 @@ import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
import org.whispersystems.textsecuregcm.storage.MessagesManager;
import org.whispersystems.textsecuregcm.storage.ReportMessageManager;
import org.whispersystems.textsecuregcm.util.DestinationDeviceValidator;
import org.whispersystems.textsecuregcm.util.ExceptionUtils;
import org.whispersystems.textsecuregcm.util.Util;
import org.whispersystems.textsecuregcm.websocket.WebSocketConnection;
import org.whispersystems.websocket.Stories;
@@ -150,6 +151,8 @@ public class MessageController {
private static final int MAX_FETCH_ACCOUNT_CONCURRENCY = 8;
private static final CompletableFuture<?>[] EMPTY_FUTURE_ARRAY = new CompletableFuture<?>[0];
private static final String REJECT_OVERSIZE_MESSAGE_COUNTER = name(MessageController.class, "rejectOversizeMessage");
private static final String SENT_MESSAGE_COUNTER_NAME = name(MessageController.class, "sentMessages");
private static final String CONTENT_SIZE_DISTRIBUTION_NAME = name(MessageController.class, "messageContentSize");
@@ -447,8 +450,22 @@ public class MessageController {
if (recipients.isEmpty()) {
return Response.ok(new SendMultiRecipientMessageResponse(List.of())).build();
}
for (MultiRecipientDeliveryData recipient : recipients.values()) {
rateLimiters.getStoriesLimiter().validate(recipient.account().getUuid());
try {
CompletableFuture.allOf(recipients.values()
.stream()
.map(recipient -> recipient.account().getUuid())
.map(accountIdentifier ->
rateLimiters.getStoriesLimiter().validateAsync(accountIdentifier).toCompletableFuture())
.toList()
.toArray(EMPTY_FUTURE_ARRAY))
.join();
} catch (final Exception e) {
if (ExceptionUtils.unwrap(e) instanceof RateLimitExceededException rateLimitExceededException) {
throw rateLimitExceededException;
} else {
throw ExceptionUtils.wrap(e);
}
}
}