Add back story ratelimiter with counter but do not enforce

This commit is contained in:
Katherine Yen 2023-10-13 14:56:59 -07:00 committed by Jon Chambers
parent 33b4f17945
commit b2ff016cc1
3 changed files with 13 additions and 4 deletions

View File

@ -136,6 +136,8 @@ public class MessageController {
private static final String CONTENT_SIZE_DISTRIBUTION_NAME = name(MessageController.class, "messageContentSize"); private static final String CONTENT_SIZE_DISTRIBUTION_NAME = name(MessageController.class, "messageContentSize");
private static final String OUTGOING_MESSAGE_LIST_SIZE_BYTES_DISTRIBUTION_NAME = name(MessageController.class, "outgoingMessageListSizeBytes"); private static final String OUTGOING_MESSAGE_LIST_SIZE_BYTES_DISTRIBUTION_NAME = name(MessageController.class, "outgoingMessageListSizeBytes");
private static final String RATE_LIMITED_MESSAGE_COUNTER_NAME = name(MessageController.class, "rateLimitedMessage"); private static final String RATE_LIMITED_MESSAGE_COUNTER_NAME = name(MessageController.class, "rateLimitedMessage");
private static final String RATE_LIMITED_STORIES_COUNTER_NAME = name(MessageController.class, "rateLimitedStory");
private static final String REJECT_INVALID_ENVELOPE_TYPE = name(MessageController.class, "rejectInvalidEnvelopeType"); private static final String REJECT_INVALID_ENVELOPE_TYPE = name(MessageController.class, "rejectInvalidEnvelopeType");
private static final String EPHEMERAL_TAG_NAME = "ephemeral"; private static final String EPHEMERAL_TAG_NAME = "ephemeral";
@ -276,7 +278,7 @@ public class MessageController {
} }
if (isStory) { if (isStory) {
checkStoryRateLimit(destination.get()); checkStoryRateLimit(destination.get(), userAgent);
} }
final Set<Long> excludedDeviceIds; final Set<Long> excludedDeviceIds;
@ -413,7 +415,7 @@ public class MessageController {
accountsByServiceIdentifier.forEach((serviceIdentifier, account) -> { accountsByServiceIdentifier.forEach((serviceIdentifier, account) -> {
if (isStory) { if (isStory) {
checkStoryRateLimit(account); checkStoryRateLimit(account, userAgent);
} }
Set<Long> deviceIds = accountToDeviceIdAndRegistrationIdMap Set<Long> deviceIds = accountToDeviceIdAndRegistrationIdMap
@ -732,10 +734,11 @@ public class MessageController {
} }
} }
private void checkStoryRateLimit(Account destination) { private void checkStoryRateLimit(Account destination, String userAgent) {
try { try {
rateLimiters.getMessagesLimiter().validate(destination.getUuid()); rateLimiters.getStoriesLimiter().validate(destination.getUuid());
} catch (final RateLimitExceededException e) { } catch (final RateLimitExceededException e) {
Metrics.counter(RATE_LIMITED_STORIES_COUNTER_NAME, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent))).increment();
} }
} }

View File

@ -28,6 +28,7 @@ public class RateLimiters extends BaseRateLimiters<RateLimiters.For> {
ATTACHMENT("attachmentCreate", false, new RateLimiterConfig(50, Duration.ofMillis(1200))), ATTACHMENT("attachmentCreate", false, new RateLimiterConfig(50, Duration.ofMillis(1200))),
PRE_KEYS("prekeys", false, new RateLimiterConfig(6, Duration.ofMinutes(10))), PRE_KEYS("prekeys", false, new RateLimiterConfig(6, Duration.ofMinutes(10))),
MESSAGES("messages", false, new RateLimiterConfig(60, Duration.ofSeconds(1))), MESSAGES("messages", false, new RateLimiterConfig(60, Duration.ofSeconds(1))),
STORIES("stories", false, new RateLimiterConfig(5_000, Duration.ofSeconds(8))),
ALLOCATE_DEVICE("allocateDevice", false, new RateLimiterConfig(2, Duration.ofMinutes(2))), ALLOCATE_DEVICE("allocateDevice", false, new RateLimiterConfig(2, Duration.ofMinutes(2))),
VERIFY_DEVICE("verifyDevice", false, new RateLimiterConfig(6, Duration.ofMinutes(10))), VERIFY_DEVICE("verifyDevice", false, new RateLimiterConfig(6, Duration.ofMinutes(10))),
TURN("turnAllocate", false, new RateLimiterConfig(60, Duration.ofSeconds(1))), TURN("turnAllocate", false, new RateLimiterConfig(60, Duration.ofSeconds(1))),
@ -213,4 +214,8 @@ public class RateLimiters extends BaseRateLimiters<RateLimiters.For> {
public RateLimiter getInboundMessageBytes() { public RateLimiter getInboundMessageBytes() {
return forDescriptor(For.INBOUND_MESSAGE_BYTES); return forDescriptor(For.INBOUND_MESSAGE_BYTES);
} }
public RateLimiter getStoriesLimiter() {
return forDescriptor(For.STORIES);
}
} }

View File

@ -220,6 +220,7 @@ class MessageControllerTest {
when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration); when(dynamicConfigurationManager.getConfiguration()).thenReturn(dynamicConfiguration);
when(rateLimiters.getMessagesLimiter()).thenReturn(rateLimiter); when(rateLimiters.getMessagesLimiter()).thenReturn(rateLimiter);
when(rateLimiters.getStoriesLimiter()).thenReturn(rateLimiter);
when(rateLimiters.getInboundMessageBytes()).thenReturn(rateLimiter); when(rateLimiters.getInboundMessageBytes()).thenReturn(rateLimiter);
} }