diff --git a/sdks/java/io/google-cloud-platform/build.gradle b/sdks/java/io/google-cloud-platform/build.gradle index d72afac96b0d..43e16288348e 100644 --- a/sdks/java/io/google-cloud-platform/build.gradle +++ b/sdks/java/io/google-cloud-platform/build.gradle @@ -65,6 +65,7 @@ dependencies { implementation library.java.google_api_common implementation library.java.google_api_services_bigquery implementation library.java.opentelemetry_api + implementation library.java.opentelemetry_context implementation library.java.google_api_services_healthcare implementation library.java.google_api_services_pubsub implementation library.java.google_api_services_storage diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java index ee4e283cc2f2..fa2399abd91c 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java @@ -28,9 +28,17 @@ import com.google.protobuf.DynamicMessage; import com.google.protobuf.InvalidProtocolBufferException; import com.google.protobuf.Message; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.context.propagation.TextMapGetter; +import io.opentelemetry.context.propagation.TextMapSetter; import java.io.IOException; import java.io.Serializable; import java.nio.charset.StandardCharsets; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -52,6 +60,7 @@ import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath; import org.apache.beam.sdk.metrics.Lineage; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.SdkHarnessOptions; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; @@ -90,6 +99,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import org.slf4j.Logger; @@ -880,6 +890,8 @@ public abstract static class Read extends PTransform> abstract boolean getNeedsMessageId(); + abstract boolean isEnableOpenTelemetryTracing(); + abstract boolean getNeedsOrderingKey(); abstract BadRecordRouter getBadRecordRouter(); @@ -900,6 +912,7 @@ static Builder newBuilder(SerializableFunction parseFn) builder.setBadRecordRouter(BadRecordRouter.THROWING_ROUTER); builder.setBadRecordErrorHandler(new DefaultErrorHandler<>()); builder.setValidate(false); + builder.setEnableOpenTelemetryTracing(false); return builder; } @@ -922,6 +935,8 @@ abstract static class Builder { abstract Builder setIdAttribute(String idAttribute); + abstract Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing); + abstract Builder setCoder(Coder coder); abstract Builder setParseFn(SerializableFunction parseFn); @@ -1105,6 +1120,66 @@ public Read withIdAttribute(String idAttribute) { return toBuilder().setIdAttribute(idAttribute).build(); } + public Read withEnableOpenTelemetryTracing() { + return toBuilder().setEnableOpenTelemetryTracing(true).setNeedsAttributes(true).build(); + } + + static class OpenTelemetryHeaderConsumer extends DoFn { + + Context extractSpanContext(PubsubMessage message) { + TextMapGetter extractMessageAttributes = + new TextMapGetter() { + @Override + public @Nullable String get(@Nullable PubsubMessage carrier, String key) { + if (carrier == null) { + return null; + } + return carrier.getAttribute("googclient_" + key); + } + + @Override + public Iterable keys(PubsubMessage carrier) { + Map attributeMap = carrier.getAttributeMap(); + if (attributeMap == null) { + return ImmutableList.of(); + } + List keys = new java.util.ArrayList<>(); + for (String key : attributeMap.keySet()) { + if (key.startsWith("googclient_")) { + keys.add(key.substring("googclient_".length())); + } + } + return keys; + } + }; + return W3CTraceContextPropagator.getInstance() + .extract(Context.current(), message, extractMessageAttributes); + } + + @Setup + public void setup(PipelineOptions po) { + tracer = po.as(SdkHarnessOptions.class).getOpenTelemetry().getTracer("PubSubIO"); + } + + private transient @MonotonicNonNull Tracer tracer = null; + + @ProcessElement + public void processElement( + @Element PubsubMessage message, OutputReceiver output) { + @Nullable Context context = extractSpanContext(message); + Span span = + checkArgumentNotNull(tracer) + .spanBuilder("PubSubIO.Read") + .setParent(context) + .startSpan(); + try (Scope s = span.makeCurrent()) { + output.output(message); + } finally { + span.end(); + } + } + } + /** * Causes the source to return a PubsubMessage that includes Pubsub attributes, and uses the * given parsing function to transform the PubsubMessage into an output type. A Coder for the @@ -1234,6 +1309,12 @@ public T apply(PubsubMessage input) { }; ValueProvider deadLetterTopicProvider = getDeadLetterTopicProvider(); PCollection read; + if (isEnableOpenTelemetryTracing()) { + preParse = + preParse.apply( + "Extract OpenTelemetry context from Header", + ParDo.of(new OpenTelemetryHeaderConsumer())); + } if (deadLetterTopicProvider == null && (getBadRecordRouter() instanceof ThrowingBadRecordRouter)) { read = preParse.apply(MapElements.into(typeDescriptor).via(parseFnWrapped)); @@ -1413,6 +1494,8 @@ public abstract static class Write extends PTransform, PDone> abstract @Nullable String getPubsubRootUrl(); + abstract boolean isEnableOpenTelemetryTracing(); + abstract boolean getPublishWithOrderingKey(); abstract BadRecordRouter getBadRecordRouter(); @@ -1432,6 +1515,7 @@ static Builder newBuilder( builder.setBadRecordErrorHandler(new DefaultErrorHandler<>()); builder.setPublishWithOrderingKey(false); builder.setValidate(false); + builder.setEnableOpenTelemetryTracing(false); return builder; } @@ -1456,6 +1540,8 @@ abstract Builder setTopicFunction( abstract Builder setTimestampAttribute(String timestampAttribute); + abstract Builder setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing); + abstract Builder setIdAttribute(String idAttribute); abstract Builder setFormatFn( @@ -1475,6 +1561,41 @@ abstract Builder setBadRecordErrorHandler( abstract Write build(); } + static class OpenTelemetryHeaderPropagator extends DoFn { + void injectSpanContext(Map attr) { + TextMapSetter> inject = + new TextMapSetter>() { + @Override + public void set(@Nullable Map attr, String key, String value) { + if (attr != null) { + attr.put("googclient_" + key, value); + } + } + }; + W3CTraceContextPropagator.getInstance().inject(Context.current(), attr, inject); + } + + @ProcessElement + public void processElement( + @Element PubsubMessage message, OutputReceiver output) { + Map attributeMap = message.getAttributeMap(); + Map attr = + attributeMap == null ? new HashMap<>() : new HashMap<>(attributeMap); + injectSpanContext(attr); + + // copy the message, multiple fields + PubsubMessage ps = + new PubsubMessage( + message.getPayload(), attr, message.getMessageId(), message.getOrderingKey()); + + // topic is copied seperately, not via constructor + if (message.getTopic() != null) { + ps = ps.withTopic(message.getTopic()); + } + output.output(ps); + } + } + /** * Publishes to the specified topic. * @@ -1550,6 +1671,10 @@ public Write withMaxBatchBytesSize(int maxBatchBytesSize) { return toBuilder().setMaxBatchBytesSize(maxBatchBytesSize).build(); } + public Write withEnableOpenTelemetryTracing() { + return toBuilder().setEnableOpenTelemetryTracing(true).build(); + } + /** * Writes to Pub/Sub with each record's ordering key. A subscription with message ordering * enabled will receive messages published in the same region with the same ordering key in the @@ -1657,6 +1782,11 @@ public PDone expand(PCollection input) { } else { pubsubMessages.setCoder(PubsubMessageWithTopicCoder.of()); } + if (isEnableOpenTelemetryTracing()) { + pubsubMessages = + pubsubMessages.apply( + "Propagate OpenTelemetry Tracing", ParDo.of(new OpenTelemetryHeaderPropagator())); + } switch (input.isBounded()) { case BOUNDED: pubsubMessages.apply(