OTEL in pubsub

Author: stankiewicz

sdks/java/io/google-cloud-platform/build.gradle

@@ -58,6 +58,8 @@ dependencies {
   implementation library.java.google_api_client
   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

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java

@@ -26,9 +26,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;
@@ -50,6 +58,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;
@@ -854,6 +863,8 @@ public abstract static class Read<T> extends PTransform<PBegin, PCollection<T>>
 
     abstract boolean getNeedsMessageId();
 
+    abstract boolean isEnableOpenTelemetryTracing();
+
     abstract boolean getNeedsOrderingKey();
 
     abstract BadRecordRouter getBadRecordRouter();
@@ -874,6 +885,7 @@ static <T> Builder<T> newBuilder(SerializableFunction<PubsubMessage, T> parseFn)
       builder.setBadRecordRouter(BadRecordRouter.THROWING_ROUTER);
       builder.setBadRecordErrorHandler(new DefaultErrorHandler<>());
       builder.setValidate(false);
+      builder.setEnableOpenTelemetryTracing(false);
       return builder;
     }
 
@@ -896,6 +908,8 @@ abstract static class Builder<T> {
 
       abstract Builder<T> setIdAttribute(String idAttribute);
 
+      abstract Builder<T> setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing);
+
       abstract Builder<T> setCoder(Coder<T> coder);
 
       abstract Builder<T> setParseFn(SerializableFunction<PubsubMessage, T> parseFn);
@@ -1079,6 +1093,66 @@ public Read<T> withIdAttribute(String idAttribute) {
       return toBuilder().setIdAttribute(idAttribute).build();
     }
 
+    public Read<T> withEnableOpenTelemetryTracing() {
+      return toBuilder().setEnableOpenTelemetryTracing(true).setNeedsAttributes(true).build();
+    }
+
+    static class OpenTelemetryHeaderConsumer extends DoFn<PubsubMessage, PubsubMessage> {
+      Context extractSpanContext(PubsubMessage message) {
+        TextMapGetter<PubsubMessage> extractMessageAttributes =
+            new TextMapGetter<PubsubMessage>() {
+              @Override
+              public String get(@Nullable PubsubMessage carrier, String key) {
+                if (carrier == null) {
+                  return null;
+                }
+                return carrier.getAttribute("googclient_" + key);
+              }
+
+              @Override
+              public Iterable<String> keys(PubsubMessage carrier) {
+                Map<String, String> attributeMap = carrier.getAttributeMap();
+                if (attributeMap == null) {
+                  return ImmutableList.of();
+                }
+                List<String> 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 Tracer tracer = null;
+
+      @ProcessElement
+      public void processElement(
+          @Element PubsubMessage message,
+          OutputReceiver<PubsubMessage> output) {
+        Context context = extractSpanContext(message);
+        Span psSub =
+            Preconditions.checkArgumentNotNull(tracer)
+                .spanBuilder("PubSubIO.Read")
+                .setParent(context)
+                .startSpan();
+        try (Scope s = psSub.makeCurrent()) {
+          output.output(message);
+        } finally {
+          psSub.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
@@ -1205,6 +1279,14 @@ public T apply(PubsubMessage input) {
             }
           };
       PCollection<T> read;
+
+      if (isEnableOpenTelemetryTracing()) {
+        preParse =
+            preParse.apply(
+                "Extract OpenTelemetry context from Header",
+                ParDo.of(new OpenTelemetryHeaderConsumer()));
+      }
+
       if (getDeadLetterTopicProvider() == null
           && (getBadRecordRouter() instanceof ThrowingBadRecordRouter)) {
         read = preParse.apply(MapElements.into(typeDescriptor).via(parseFnWrapped));
@@ -1377,6 +1459,8 @@ public abstract static class Write<T> extends PTransform<PCollection<T>, PDone>
 
     abstract @Nullable String getPubsubRootUrl();
 
+    abstract boolean isEnableOpenTelemetryTracing();
+
     abstract boolean getPublishWithOrderingKey();
 
     abstract BadRecordRouter getBadRecordRouter();
@@ -1396,6 +1480,7 @@ static <T> Builder<T> newBuilder(
       builder.setBadRecordErrorHandler(new DefaultErrorHandler<>());
       builder.setPublishWithOrderingKey(false);
       builder.setValidate(false);
+      builder.setEnableOpenTelemetryTracing(false);
       return builder;
     }
 
@@ -1420,6 +1505,8 @@ abstract Builder<T> setTopicFunction(
 
       abstract Builder<T> setTimestampAttribute(String timestampAttribute);
 
+      abstract Builder<T> setEnableOpenTelemetryTracing(boolean enableOpenTelemetryTracing);
+
       abstract Builder<T> setIdAttribute(String idAttribute);
 
       abstract Builder<T> setFormatFn(
@@ -1439,6 +1526,42 @@ abstract Builder<T> setBadRecordErrorHandler(
       abstract Write<T> build();
     }
 
+    static class OpenTelemetryHeaderPropagator extends DoFn<PubsubMessage, PubsubMessage> {
+      void injectSpanContext(Map<String, String> attr) {
+        TextMapSetter<Map<String, String>> inject =
+            new TextMapSetter<Map<String, String>>() {
+              @Override
+              public void set(@Nullable Map<String, String> attr, String key, String value) {
+                if (attr != null) {
+                  attr.put("googclient_" + key, value);
+                }
+              }
+            };
+        W3CTraceContextPropagator.getInstance().inject(Context.current(), attr, inject);
+      }
+
+      @ProcessElement
+      public void processElement(
+          ProcessContext c,
+          @Element PubsubMessage message,
+          OutputReceiver<PubsubMessage> output,
+          PipelineOptions po) {
+        Map<String, String> attributeMap = message.getAttributeMap();
+        Map<String, String> 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
+        ps = ps.withTopic(message.getTopic());
+        output.output(ps);
+      }
+    }
+
     /**
      * Publishes to the specified topic.
      *
@@ -1514,6 +1637,10 @@ public Write<T> withMaxBatchBytesSize(int maxBatchBytesSize) {
       return toBuilder().setMaxBatchBytesSize(maxBatchBytesSize).build();
     }
 
+    public Write<T> 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
@@ -1620,6 +1747,11 @@ public PDone expand(PCollection<T> 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(