Add metric for requested attachment upload size

This commit is contained in:
Ravi Khadiwala 2026-06-23 13:59:26 -05:00 committed by ravi-signal
parent 28aefe0ebe
commit 8b617b64f8
2 changed files with 36 additions and 3 deletions

View File

@ -5,16 +5,21 @@
package org.whispersystems.textsecuregcm.controllers;
import com.google.common.net.HttpHeaders;
import io.dropwizard.auth.Auth;
import io.micrometer.core.instrument.DistributionSummary;
import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.Tags;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.headers.Header;
import io.swagger.v3.oas.annotations.responses.ApiResponse;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.validation.Valid;
import jakarta.validation.constraints.Positive;
import jakarta.ws.rs.ClientErrorException;
import jakarta.ws.rs.GET;
import jakarta.ws.rs.HeaderParam;
import jakarta.ws.rs.Path;
import jakarta.ws.rs.Produces;
import jakarta.ws.rs.QueryParam;
@ -24,6 +29,7 @@ import java.security.SecureRandom;
import java.util.Map;
import java.util.Optional;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import org.whispersystems.textsecuregcm.attachments.AttachmentGenerator;
import org.whispersystems.textsecuregcm.attachments.AttachmentUtil;
import org.whispersystems.textsecuregcm.attachments.GcsAttachmentGenerator;
@ -33,6 +39,8 @@ import org.whispersystems.textsecuregcm.entities.AttachmentDescriptorV3;
import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
import org.whispersystems.textsecuregcm.limits.RateLimiter;
import org.whispersystems.textsecuregcm.limits.RateLimiters;
import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
/**
@ -40,7 +48,7 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
* (message attachments) to a remote storage location. The location may be selected by the server at runtime.
*/
@Path("/v4/attachments")
@Tag(name = "Attachments")
@io.swagger.v3.oas.annotations.tags.Tag(name = "Attachments")
public class AttachmentControllerV4 {
private final ExperimentEnrollmentManager experimentEnrollmentManager;
@ -50,6 +58,9 @@ public class AttachmentControllerV4 {
private final Map<Integer, AttachmentGenerator> attachmentGenerators;
private static final String ATTACHMENT_SIZE_NAME =
MetricsUtil.name(AttachmentControllerV4.class, "attachmentSize");
@Nonnull
private final SecureRandom secureRandom;
@ -91,7 +102,8 @@ public class AttachmentControllerV4 {
public AttachmentDescriptorV3 getAttachmentUploadForm(
@Auth AuthenticatedDevice auth,
@Parameter(description = "The size of the attachment to upload in bytes")
@QueryParam("uploadLength") final @Valid Optional<@Positive Long> maybeUploadLength)
@QueryParam("uploadLength") final @Valid Optional<@Positive Long> maybeUploadLength,
@HeaderParam(HttpHeaders.USER_AGENT) @Nullable final String userAgent)
throws RateLimitExceededException {
final long uploadLength = maybeUploadLength.orElse(maxUploadLength);
@ -111,6 +123,12 @@ public class AttachmentControllerV4 {
}
}
DistributionSummary.builder(ATTACHMENT_SIZE_NAME)
.tags(Tags.of(UserAgentTagUtil.getPlatformTag(userAgent),
Tag.of("uploadLengthSupplied", Boolean.toString(maybeUploadLength.isPresent()))))
.register(Metrics.globalRegistry)
.record(uploadLength);
final String key = AttachmentUtil.generateAttachmentKey(secureRandom);
final boolean useCdn3 = this.experimentEnrollmentManager.isEnrolled(auth.accountIdentifier(), AttachmentUtil.CDN3_EXPERIMENT_NAME);
int cdn = useCdn3 ? 3 : 2;

View File

@ -11,6 +11,10 @@ import java.time.Instant;
import java.util.HexFormat;
import java.util.Map;
import java.util.stream.IntStream;
import io.micrometer.core.instrument.DistributionSummary;
import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.Tags;
import org.signal.chat.attachments.GetStickerUploadFormRequest;
import org.signal.chat.attachments.GetStickerUploadFormResponse;
import org.signal.chat.attachments.GetUploadFormRequest;
@ -25,11 +29,14 @@ import org.whispersystems.textsecuregcm.attachments.GcsAttachmentGenerator;
import org.whispersystems.textsecuregcm.attachments.TusAttachmentGenerator;
import org.whispersystems.textsecuregcm.auth.grpc.AuthenticatedDevice;
import org.whispersystems.textsecuregcm.auth.grpc.AuthenticationUtil;
import org.whispersystems.textsecuregcm.controllers.AttachmentControllerV4;
import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
import org.whispersystems.textsecuregcm.controllers.StickerController;
import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
import org.whispersystems.textsecuregcm.limits.RateLimiter;
import org.whispersystems.textsecuregcm.limits.RateLimiters;
import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
import org.whispersystems.textsecuregcm.s3.PostPolicyGenerator;
public class AttachmentsGrpcService extends SimpleAttachmentsGrpc.AttachmentsImplBase {
@ -49,6 +56,9 @@ public class AttachmentsGrpcService extends SimpleAttachmentsGrpc.AttachmentsImp
.setAlgorithm(PostPolicyGenerator.ALGORITHM)
.build();
private static final String ATTACHMENT_SIZE_NAME =
MetricsUtil.name(AttachmentsGrpcService.class, "attachmentSize");
public AttachmentsGrpcService(
final ExperimentEnrollmentManager experimentEnrollmentManager,
final RateLimiters rateLimiters,
@ -89,6 +99,11 @@ public class AttachmentsGrpcService extends SimpleAttachmentsGrpc.AttachmentsImp
throw e;
}
DistributionSummary.builder(ATTACHMENT_SIZE_NAME)
.tags(Tags.of(UserAgentTagUtil.getPlatformTag(RequestAttributesUtil.getUserAgent().orElse(null))))
.register(Metrics.globalRegistry)
.record(request.getUploadLength());
final String key = AttachmentUtil.generateAttachmentKey(secureRandom);
final boolean useCdn3 = this.experimentEnrollmentManager.isEnrolled(auth.accountIdentifier(),
AttachmentUtil.CDN3_EXPERIMENT_NAME);