Hoist user-agent tag out of publisher in copyMedia

This commit is contained in:
Ravi Khadiwala
2026-07-17 12:42:27 -05:00
committed by ravi-signal
parent ea411e2f1e
commit b255c37978
5 changed files with 39 additions and 10 deletions
@@ -8,6 +8,7 @@ import com.google.protobuf.ByteString;
import com.google.protobuf.Empty;
import java.util.Optional;
import java.util.concurrent.Flow;
import io.micrometer.core.instrument.Tag;
import org.signal.chat.backup.BackupStreamClosed;
import org.signal.chat.backup.CopyMediaRequest;
import org.signal.chat.backup.CopyMediaResponse;
@@ -235,10 +236,10 @@ public class BackupsAnonymousGrpcService extends SimpleBackupsAnonymousGrpc.Back
.build())
.build())));
}
final Tag platformTag = UserAgentTagUtil.getPlatformTag(RequestAttributesUtil.getUserAgent().orElse(null));
return JdkFlowAdapter.publisherToFlowPublisher(backupManager.copyToBackup(copyQuota)
.doOnNext(result -> backupMetrics.updateCopyCounter(
result,
UserAgentTagUtil.getPlatformTag(RequestAttributesUtil.getUserAgent().orElse(null))))
.doOnNext(result -> backupMetrics.updateCopyCounter(result, platformTag))
.map(copyResult -> {
CopyMediaResponse.Builder builder = CopyMediaResponse
.newBuilder()
@@ -70,6 +70,7 @@ import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
import org.whispersystems.textsecuregcm.metrics.BackupMetrics;
import org.whispersystems.textsecuregcm.util.TestRandomUtil;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;
class BackupsAnonymousGrpcServiceTest extends
SimpleBaseGrpcTest<BackupsAnonymousGrpcService, BackupsAnonymousGrpc.BackupsAnonymousBlockingStub> {
@@ -146,10 +147,11 @@ class BackupsAnonymousGrpcServiceTest extends
@Test
void putMediaBatchSuccess() {
final byte[][] mediaIds = {TestRandomUtil.nextBytes(15), TestRandomUtil.nextBytes(15)};
when(backupManager.copyToBackup(any()))
.thenReturn(Flux.just(
when(backupManager.copyToBackup(any())).thenReturn(Flux.just(
new CopyResult(CopyResult.Outcome.SUCCESS, mediaIds[0], 1),
new CopyResult(CopyResult.Outcome.SUCCESS, mediaIds[1], 1)));
new CopyResult(CopyResult.Outcome.SUCCESS, mediaIds[1], 1))
// helps catch any issues with streaming concurrency (especially context propagation)
.publishOn(Schedulers.parallel()));
final CopyMediaRequest request = CopyMediaRequest.newBuilder()
.setSignedPresentation(signedPresentation(presentation))
@@ -71,6 +71,7 @@ import org.whispersystems.textsecuregcm.storage.KeysManager;
import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
import org.whispersystems.textsecuregcm.tests.util.DevicesHelper;
import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
import org.whispersystems.textsecuregcm.util.CompletableFutureTestUtil;
import org.whispersystems.textsecuregcm.util.TestClock;
import org.whispersystems.textsecuregcm.util.TestRandomUtil;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
@@ -422,17 +423,19 @@ class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcS
final UUID mismatchedPniFingerprintAccountIdentifier = UUID.randomUUID();
final IdentityKey mismatchedPniFingerpringAccountIdentityKey = new IdentityKey(ECKeyPair.generate().getPublicKey());
// Complete futures asynchronously to catch potential async/context-propagation issues
final Duration futureDelay = Duration.ofMillis(1);
when(mismatchedAciFingerprintAccount.getIdentityKey(IdentityType.ACI)).thenReturn(mismatchedAciFingerprintAccountIdentityKey);
when(accountsManager.getByServiceIdentifierAsync(new AciServiceIdentifier(mismatchedAciFingerprintAccountIdentifier)))
.thenReturn(CompletableFuture.completedFuture(Optional.of(mismatchedAciFingerprintAccount)));
.thenReturn(delayedAccount(mismatchedAciFingerprintAccount));
when(matchingAciFingerprintAccount.getIdentityKey(IdentityType.ACI)).thenReturn(matchingAciFingerprintAccountIdentityKey);
when(accountsManager.getByServiceIdentifierAsync(new AciServiceIdentifier(matchingAciFingerprintAccountIdentifier)))
.thenReturn(CompletableFuture.completedFuture(Optional.of(matchingAciFingerprintAccount)));
.thenReturn(delayedAccount(matchingAciFingerprintAccount));
when(mismatchedPniFingerprintAccount.getIdentityKey(IdentityType.PNI)).thenReturn(mismatchedPniFingerpringAccountIdentityKey);
when(accountsManager.getByServiceIdentifierAsync(new PniServiceIdentifier(mismatchedPniFingerprintAccountIdentifier)))
.thenReturn(CompletableFuture.completedFuture(Optional.of(mismatchedPniFingerprintAccount)));
.thenReturn(delayedAccount(mismatchedPniFingerprintAccount));
final Map<UUID, IdentityKey> expectedResponses = Map.of(
mismatchedAciFingerprintAccountIdentifier, mismatchedAciFingerprintAccountIdentityKey,
@@ -489,6 +492,10 @@ class KeysAnonymousGrpcServiceTest extends SimpleBaseGrpcTest<KeysAnonymousGrpcS
assertEquals(expectedResponses, responses);
}
private static CompletableFuture<Optional<Account>> delayedAccount(final Account account) {
return CompletableFutureTestUtil.almostCompletedFuture(Optional.of(account), Duration.ofMillis(1));
}
private static CheckIdentityKeyRequest buildCheckIdentityKeyRequest(final org.signal.chat.common.IdentityType identityType,
final UUID uuid, final IdentityKey identityKey) {
return CheckIdentityKeyRequest.newBuilder()
@@ -89,6 +89,7 @@ import org.whispersystems.textsecuregcm.util.TestClock;
import org.whispersystems.textsecuregcm.util.TestRandomUtil;
import org.whispersystems.textsecuregcm.util.UUIDUtil;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;
import reactor.test.publisher.TestPublisher;
class MessagesGrpcServiceTest extends SimpleBaseGrpcTest<MessagesGrpcService, MessagesGrpc.MessagesBlockingV2Stub> {
@@ -809,6 +810,19 @@ class MessagesGrpcServiceTest extends SimpleBaseGrpcTest<MessagesGrpcService, Me
blockingCall.write(request);
assertEquals(Status.INVALID_ARGUMENT.getCode(), reader.join().getStatus().getCode());
}
@Test
void getMessages() throws StatusException, InterruptedException {
final BlockingClientCall<GetMessagesRequest, GetMessagesResponse> blockingCall = authenticatedServiceStub().getMessages();
when(messageDispatcher.getMessages(anyBoolean(), any(), any(), any(), any())).thenReturn(Flux.just(
GetMessagesResponse.newBuilder().setEnvelope(MessageProtos.Envelope.getDefaultInstance()).build(),
GetMessagesResponse.newBuilder().setQueueEmpty(Empty.getDefaultInstance()).build())
// helps catch any issues with streaming concurrency (especially context propagation)
.publishOn(Schedulers.parallel()));
blockingCall.write(GetMessagesRequest.newBuilder().setOptions(GetMessagesRequest.GetMessageOptions.getDefaultInstance()).build());
assertTrue(blockingCall.read().hasEnvelope());
assertTrue(blockingCall.read().hasQueueEmpty());
}
}
private static Executable convertStatusException(final Executable serviceCall) {
@@ -8,6 +8,7 @@ package org.whispersystems.textsecuregcm.util;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.TimeUnit;
@@ -30,7 +31,11 @@ public class CompletableFutureTestUtil {
}
public static <T> CompletableFuture<T> almostCompletedFuture(T result) {
return new CompletableFuture<T>().completeOnTimeout(result, 5, TimeUnit.MILLISECONDS);
return almostCompletedFuture(result, Duration.ofMillis(5));
}
public static <T> CompletableFuture<T> almostCompletedFuture(T result, Duration timeout) {
return new CompletableFuture<T>().completeOnTimeout(result, timeout.toNanos(), TimeUnit.NANOSECONDS);
}
}