diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java index ea98a8884507..807f91c9ea01 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java @@ -49,10 +49,13 @@ import com.google.cloud.pubsub.v1.stub.PublisherStubSettings; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.protobuf.CodedOutputStream; import com.google.pubsub.v1.PublishRequest; import com.google.pubsub.v1.PublishResponse; +import com.google.pubsub.v1.PubsubClientTelemetry; import com.google.pubsub.v1.PubsubMessage; import com.google.pubsub.v1.TopicName; import com.google.pubsub.v1.TopicNames; @@ -63,6 +66,7 @@ import java.io.IOException; import java.time.Duration; import java.util.ArrayList; +import java.util.Base64; import java.util.Collections; import java.util.HashMap; import java.util.Iterator; @@ -105,6 +109,8 @@ public class Publisher implements PublisherInterface { private static final Logger logger = Logger.getLogger(Publisher.class.getName()); private LoggingUtil loggingUtil = new LoggingUtil(); + @VisibleForTesting static final String TELEMETRY_HEADER_KEY = "x-goog-pubsub-client-telemetry"; + private static final String GZIP_COMPRESSION = "gzip"; private static final String OPEN_TELEMETRY_TRACER_NAME = "com.google.cloud.pubsub.v1"; @@ -540,6 +546,18 @@ private void publishAllWithoutInflightForKey(final String orderingKey) { } } + private Map> createTelemetryHeader(int attemptNumber) { + PubsubClientTelemetry telemetry = + PubsubClientTelemetry.newBuilder() + .setPublishOperation( + PubsubClientTelemetry.PublishOperation.newBuilder() + .setHedgedAttemptCount(attemptNumber) + .build()) + .build(); + String encodedHeader = Base64.getEncoder().encodeToString(telemetry.toByteArray()); + return ImmutableMap.of(TELEMETRY_HEADER_KEY, ImmutableList.of(encodedHeader)); + } + private ApiFuture publishCall(OutstandingBatch outstandingBatch) { return publishCall(outstandingBatch, 0, null); } @@ -561,7 +579,7 @@ private ApiFuture publishCall( outstandingBatch.getMessageWrappers().get(0)); context = context.withRetryableCodes(Collections.emptySet()); } - + context = context.withExtraHeaders(createTelemetryHeader(attemptNumber)); int numMessagesInBatch = outstandingBatch.size(); List pubsubMessagesList = new ArrayList(numMessagesInBatch); List messageWrappers = outstandingBatch.getMessageWrappers(); diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java index ddb3b5b42be6..9c2c4fed9807 100644 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java @@ -42,6 +42,7 @@ import com.google.pubsub.v1.ProjectTopicName; import com.google.pubsub.v1.PublishRequest; import com.google.pubsub.v1.PublishResponse; +import com.google.pubsub.v1.PubsubClientTelemetry; import com.google.pubsub.v1.PubsubMessage; import io.grpc.ManagedChannel; import io.grpc.Metadata; @@ -61,6 +62,7 @@ import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; import io.opentelemetry.sdk.trace.data.SpanData; import java.time.Duration; +import java.util.Base64; import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; @@ -1737,4 +1739,73 @@ private void shutdownTestPublisher(Publisher publisher) throws InterruptedExcept fakeExecutor.advanceTime(Duration.ofSeconds(10)); assertTrue(publisher.awaitTermination(1, TimeUnit.MINUTES)); } + + private PubsubClientTelemetry extractTelemetryHeader(Metadata headers) throws Exception { + Metadata.Key key = + Metadata.Key.of(Publisher.TELEMETRY_HEADER_KEY, Metadata.ASCII_STRING_MARSHALLER); + String headerValue = headers.get(key); + assertThat(headerValue).isNotNull(); + byte[] decodedBytes = Base64.getDecoder().decode(headerValue); + return PubsubClientTelemetry.parseFrom(decodedBytes); + } + + @Test + public void testTelemetryHeaderOnNormalPublish() throws Exception { + testPublisherServiceImpl.setAutoPublishResponse(true); + Publisher publisher = + getTestPublisherBuilder() + .setBatchingSettings( + Publisher.Builder.DEFAULT_BATCHING_SETTINGS.toBuilder() + .setElementCountThreshold(1L) + .build()) + .build(); + + ApiFuture future = sendTestMessage(publisher, "msg-normal"); + assertEquals("1", future.get(5, TimeUnit.SECONDS)); + + List capturedHeaders = testPublisherServiceImpl.getCapturedHeaders(); + assertThat(capturedHeaders).hasSize(1); + + PubsubClientTelemetry telemetry = extractTelemetryHeader(capturedHeaders.get(0)); + assertThat(telemetry.hasPublishOperation()).isTrue(); + assertThat(telemetry.getPublishOperation().getHedgedAttemptCount()).isEqualTo(0); + + shutdownTestPublisher(publisher); + } + + @Test + public void testTelemetryHeaderOnHedgedPublish() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 5); + + // Delay response so hedge fires + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); + + ApiFuture future = sendTestMessage(publisher, "msg-hedged"); + waitForRequests(testPublisherServiceImpl, 1); + + // Advance past hedge delay to trigger hedge attempt + fakeExecutor.advanceTime(Duration.ofMillis(120)); + waitForRequests(testPublisherServiceImpl, 2); + + // Finish request + fakeExecutor.advanceTime(Duration.ofMillis(100)); + assertEquals("1", future.get(5, TimeUnit.SECONDS)); + + List capturedHeaders = testPublisherServiceImpl.getCapturedHeaders(); + assertThat(capturedHeaders).hasSize(2); + + // Verify Attempt 0 (Original) + PubsubClientTelemetry initialTelemetry = extractTelemetryHeader(capturedHeaders.get(0)); + assertThat(initialTelemetry.getPublishOperation().getHedgedAttemptCount()).isEqualTo(0); + + // Verify Attempt 1 (Hedge) + PubsubClientTelemetry hedgedTelemetry = extractTelemetryHeader(capturedHeaders.get(1)); + assertThat(hedgedTelemetry.getPublishOperation().getHedgedAttemptCount()).isEqualTo(1); + + shutdownTestPublisher(publisher); + } }