Estimate message byte limit exceeded error count

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java

@@ -35,6 +35,7 @@ import org.whispersystems.textsecuregcm.configuration.GcpAttachmentsConfiguratio
 import org.whispersystems.textsecuregcm.configuration.GenericZkConfig;
 import org.whispersystems.textsecuregcm.configuration.HCaptchaConfiguration;
 import org.whispersystems.textsecuregcm.configuration.MaxDeviceConfiguration;
+import org.whispersystems.textsecuregcm.configuration.MessageByteLimitCardinalityEstimatorConfiguration;
 import org.whispersystems.textsecuregcm.configuration.MessageCacheConfiguration;
 import org.whispersystems.textsecuregcm.configuration.OneTimeDonationConfiguration;
 import org.whispersystems.textsecuregcm.configuration.PaymentsServiceConfiguration;
@@ -288,6 +289,11 @@ public class WhisperServerConfiguration extends Configuration {
   @JsonProperty
   private ClientReleaseConfiguration clientRelease = new ClientReleaseConfiguration(Duration.ofHours(4));
 
+  @Valid
+  @NotNull
+  @JsonProperty
+  private MessageByteLimitCardinalityEstimatorConfiguration messageByteLimitCardinalityEstimator = new MessageByteLimitCardinalityEstimatorConfiguration(Duration.ofDays(1));
+
   public AdminEventLoggingConfiguration getAdminEventLoggingConfiguration() {
     return adminEventLoggingConfiguration;
   }
@@ -478,4 +484,8 @@ public class WhisperServerConfiguration extends Configuration {
   public ClientReleaseConfiguration getClientReleaseConfiguration() {
     return clientRelease;
   }
+
+  public MessageByteLimitCardinalityEstimatorConfiguration getMessageByteLimitCardinalityEstimator() {
+    return messageByteLimitCardinalityEstimator;
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -123,6 +123,7 @@ import org.whispersystems.textsecuregcm.filters.RequestStatisticsFilter;
 import org.whispersystems.textsecuregcm.filters.TimestampResponseFilter;
 import org.whispersystems.textsecuregcm.grpc.KeysGrpcService;
 import org.whispersystems.textsecuregcm.grpc.KeysAnonymousGrpcService;
+import org.whispersystems.textsecuregcm.limits.CardinalityEstimator;
 import org.whispersystems.textsecuregcm.limits.PushChallengeManager;
 import org.whispersystems.textsecuregcm.limits.RateLimitChallengeManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
@@ -571,6 +572,11 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     final TurnTokenGenerator turnTokenGenerator = new TurnTokenGenerator(dynamicConfigurationManager,
         config.getTurnSecretConfiguration().secret().value());
 
+    final CardinalityEstimator messageByteLimitCardinalityEstimator = new CardinalityEstimator(
+        rateLimitersCluster,
+        "message_byte_limit",
+        config.getMessageByteLimitCardinalityEstimator().period());
+
     RecaptchaClient recaptchaClient = new RecaptchaClient(
         config.getRecaptchaConfiguration().projectPath(),
         useSecondaryCredentialsJson
@@ -755,9 +761,10 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         new DirectoryV2Controller(directoryV2CredentialsGenerator),
         new DonationController(clock, zkReceiptOperations, redeemedReceiptsManager, accountsManager, config.getBadges(),
             ReceiptCredentialPresentation::new),
-        new MessageController(rateLimiters, messageSender, receiptSender, accountsManager, deletedAccounts,
-            messagesManager, pushNotificationManager, reportMessageManager, multiRecipientMessageExecutor,
-            messageDeliveryScheduler, reportSpamTokenProvider, clientReleaseManager, dynamicConfigurationManager),
+        new MessageController(rateLimiters, messageByteLimitCardinalityEstimator, messageSender, receiptSender,
+            accountsManager, deletedAccounts, messagesManager, pushNotificationManager, reportMessageManager,
+            multiRecipientMessageExecutor, messageDeliveryScheduler, reportSpamTokenProvider, clientReleaseManager,
+            dynamicConfigurationManager),
         new PaymentsController(currencyManager, paymentsCredentialsGenerator),
         new ProfileController(clock, rateLimiters, accountsManager, profilesManager, dynamicConfigurationManager,
             profileBadgeConverter, config.getBadges(), cdnS3Client, profileCdnPolicyGenerator, profileCdnPolicySigner,

service/src/main/java/org/whispersystems/textsecuregcm/configuration/MessageByteLimitCardinalityEstimatorConfiguration.java

@@ -0,0 +1,11 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.configuration;
+
+import javax.validation.constraints.NotNull;
+import java.time.Duration;
+
+public record MessageByteLimitCardinalityEstimatorConfiguration(@NotNull Duration period) {}

service/src/main/java/org/whispersystems/textsecuregcm/controllers/MessageController.java

@@ -85,6 +85,7 @@ import org.whispersystems.textsecuregcm.entities.SpamReport;
 import org.whispersystems.textsecuregcm.entities.StaleDevices;
 import org.whispersystems.textsecuregcm.identity.AciServiceIdentifier;
 import org.whispersystems.textsecuregcm.identity.ServiceIdentifier;
+import org.whispersystems.textsecuregcm.limits.CardinalityEstimator;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.MessageMetrics;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
@@ -118,6 +119,7 @@ public class MessageController {
   private static final Logger logger = LoggerFactory.getLogger(MessageController.class);
 
   private final RateLimiters rateLimiters;
+  private final CardinalityEstimator messageByteLimitEstimator;
   private final MessageSender messageSender;
   private final ReceiptSender receiptSender;
   private final AccountsManager accountsManager;
@@ -153,6 +155,7 @@ public class MessageController {
 
   public MessageController(
       RateLimiters rateLimiters,
+      CardinalityEstimator messageByteLimitEstimator,
       MessageSender messageSender,
       ReceiptSender receiptSender,
       AccountsManager accountsManager,
@@ -166,6 +169,7 @@ public class MessageController {
       final ClientReleaseManager clientReleaseManager,
       final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager) {
     this.rateLimiters = rateLimiters;
+    this.messageByteLimitEstimator = messageByteLimitEstimator;
     this.messageSender = messageSender;
     this.receiptSender = receiptSender;
     this.accountsManager = accountsManager;
@@ -237,6 +241,7 @@ public class MessageController {
       rateLimiters.getInboundMessageBytes().validate(destinationIdentifier.uuid(), totalContentLength);
     } catch (final RateLimitExceededException e) {
       if (dynamicConfigurationManager.getConfiguration().getInboundMessageByteLimitConfiguration().enforceInboundLimit()) {
+        messageByteLimitEstimator.add(destinationIdentifier.uuid().toString());
         throw e;
       }
     }

service/src/main/java/org/whispersystems/textsecuregcm/limits/CardinalityEstimator.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2023 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.textsecuregcm.limits;
+
+import com.google.common.annotations.VisibleForTesting;
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.Tags;
+import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
+import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisCluster;
+import org.whispersystems.textsecuregcm.util.Util;
+
+/**
+ * Estimate the number of unique items seen over a configurable period and update a metric
+ */
+public class CardinalityEstimator {
+
+  private volatile double uniqueElementCount;
+  private final FaultTolerantRedisCluster redisCluster;
+  private final String hllName;
+  private final Duration period;
+
+  public CardinalityEstimator(final FaultTolerantRedisCluster redisCluster, final String name, final Duration period) {
+    this.redisCluster = redisCluster;
+    this.hllName = "cardinality_estimator::" + name;
+    this.period = period;
+    Metrics.gauge(
+        MetricsUtil.name(getClass(), "unique"),
+        Tags.of("name", name),
+        this,
+        obj -> obj.uniqueElementCount);
+  }
+
+  public void add(String element) {
+    addAsync(element).toCompletableFuture().join();
+  }
+
+  public CompletionStage<Void> addAsync(String element) {
+    return redisCluster.withCluster(connection -> connection.async()
+        .pfadd(hllName, element)
+        .thenCompose(modCount -> {
+          if (modCount == 0) {
+            return CompletableFuture.completedFuture(false);
+          }
+
+          // The hll changed - update our local view of the cardinality, and
+          // initialize the TTL if required
+          return connection.async()
+              .pfcount(hllName)
+              .thenCompose(count -> {
+                uniqueElementCount = count;
+                // check if this is a new hll with no TTL set
+                return connection.async().ttl(hllName).thenApply(ttl -> ttl == -1);
+              });
+        })
+        .thenCompose(isNewHll -> {
+          if (!isNewHll) {
+            return CompletableFuture.completedFuture(null);
+          }
+
+          // If this is a new hll, we need to set the TTL. This could be
+          // a single atomic op in redis 7.x with EXPIRE NX
+          return connection.async().expire(hllName, period).thenRun(Util.NOOP);
+        }));
+  }
+
+  @VisibleForTesting
+  long estimate() {
+    return (long) this.uniqueElementCount;
+  }
+}