From b255c379787bb949b76b7f2e9cc2b2ab8fcfa261 Mon Sep 17 00:00:00 2001
From: Ravi Khadiwala <ravi@signal.org>
Date: Thu, 16 Jul 2026 11:14:25 -0500
Subject: [PATCH] Hoist user-agent tag out of publisher in copyMedia

---
 .../grpc/BackupsAnonymousGrpcService.java          |  7 ++++---
 .../grpc/BackupsAnonymousGrpcServiceTest.java      |  8 +++++---
 .../grpc/KeysAnonymousGrpcServiceTest.java         | 13 ++++++++++---
 .../grpc/MessagesGrpcServiceTest.java              | 14 ++++++++++++++
 .../util/CompletableFutureTestUtil.java            |  7 ++++++-
 5 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/BackupsAnonymousGrpcService.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/BackupsAnonymousGrpcService.java
index 473f9d34d..7578c8a2e 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/BackupsAnonymousGrpcService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/BackupsAnonymousGrpcService.java
@@ -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()
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/BackupsAnonymousGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/BackupsAnonymousGrpcServiceTest.java
index 8abbdfcdf..6272b2265 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/BackupsAnonymousGrpcServiceTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/BackupsAnonymousGrpcServiceTest.java
@@ -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))
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysAnonymousGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysAnonymousGrpcServiceTest.java
index fa4e233af..7e42bf441 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysAnonymousGrpcServiceTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysAnonymousGrpcServiceTest.java
@@ -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()
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/MessagesGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/MessagesGrpcServiceTest.java
index 3592fca10..75652bced 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/MessagesGrpcServiceTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/MessagesGrpcServiceTest.java
@@ -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) {
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/util/CompletableFutureTestUtil.java b/service/src/test/java/org/whispersystems/textsecuregcm/util/CompletableFutureTestUtil.java
index 4ee4f0069..59501134e 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/util/CompletableFutureTestUtil.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/util/CompletableFutureTestUtil.java
@@ -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);
   }
 
 }