You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ra...@apache.org on 2026/10/13 15:32:51 UTC
(beam) branch master updated: [OpenTelemetry] Fix OpenTelemetry span kinds and add write spans in PubsubIO and KafkaIO (#40483)
This is an automated email from the ASF dual-hosted git repository.
stankiewicz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new bfa17fde29d [OpenTelemetry] Fix OpenTelemetry span kinds and add write spans in PubsubIO and KafkaIO (#40483)
bfa17fde29d is described below
commit bfa17fde29d36a15402b75b56e89895b34a0c31f
Author: Radosław Stankiewicz <ra...@google.com>
AuthorDate: Tue Oct 13 17:32:43 2026 +0200
[OpenTelemetry] Fix OpenTelemetry span kinds and add write spans in PubsubIO and KafkaIO (#40483)
* Fix OpenTelemetry span kinds and add write spans in PubsubIO and KafkaIO
---
.../apache/beam/sdk/io/gcp/pubsub/PubsubIO.java | 47 +++++---
.../beam/sdk/io/gcp/pubsub/PubsubIOTest.java | 121 ++++++++++++++++++++
.../java/org/apache/beam/sdk/io/kafka/KafkaIO.java | 29 ++++-
.../org/apache/beam/sdk/io/kafka/KafkaIOTest.java | 125 +++++++++++++++++++++
4 files changed, 306 insertions(+), 16 deletions(-)
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 3c63b3386e9..0eb2b44bb73 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
@@ -29,6 +29,7 @@ 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.SpanKind;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
import io.opentelemetry.context.Context;
@@ -1170,6 +1171,7 @@ public class PubsubIO {
Span span =
checkArgumentNotNull(tracer)
.spanBuilder("PubSubIO.Read")
+ .setSpanKind(SpanKind.CONSUMER)
.setParent(context)
.startSpan();
try (Scope s = span.makeCurrent()) {
@@ -1573,24 +1575,41 @@ public class PubsubIO {
W3CTraceContextPropagator.getInstance().inject(Context.current(), attr, inject);
}
+ @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<PubsubMessage> output) {
- 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
- if (message.getTopic() != null) {
- ps = ps.withTopic(message.getTopic());
+ Span span =
+ checkArgumentNotNull(tracer)
+ .spanBuilder("PubSubIO.Write")
+ .setSpanKind(SpanKind.PRODUCER)
+ .setParent(Context.current())
+ .startSpan();
+ try (Scope s = span.makeCurrent()) {
+ 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 separately, not via constructor
+ if (message.getTopic() != null) {
+ ps = ps.withTopic(message.getTopic());
+ }
+ output.output(ps);
+ } finally {
+ span.end();
}
- output.output(ps);
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
index 4aefac0379a..af20679ea7a 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
@@ -32,9 +32,19 @@ import com.google.api.client.util.Clock;
import com.google.protobuf.ByteString;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.InvalidProtocolBufferException;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
import java.io.IOException;
import java.io.Serializable;
import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
@@ -62,11 +72,13 @@ import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO.Read;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubTestClient.PubsubTestClientFactory;
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.StaticValueProvider;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.transforms.display.DisplayData;
@@ -82,6 +94,7 @@ import org.apache.beam.sdk.values.TypeDescriptors;
import org.apache.beam.sdk.values.ValueInSingleWindow;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+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.commons.lang3.RandomStringUtils;
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -98,6 +111,7 @@ import org.junit.runner.Description;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.junit.runners.model.Statement;
+import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
/** Tests for PubsubIO Read and Write transforms. */
@@ -1083,4 +1097,111 @@ public class PubsubIOTest {
write.validate(options);
}
+
+ @Test
+ public void testOpenTelemetryHeaderConsumer() {
+ String traceId = "0af7651916cd43dd8448eb211c80319c";
+ String parentSpanId = "b7ad6b7169203331";
+ String childSpanId = "1111111111111111";
+ String traceparent = "00-" + traceId + "-" + parentSpanId + "-01";
+
+ OpenTelemetry mockOpenTelemetry = Mockito.mock(OpenTelemetry.class);
+ Tracer mockTracer = Mockito.mock(Tracer.class);
+ SpanBuilder mockSpanBuilder = Mockito.mock(SpanBuilder.class, Mockito.RETURNS_SELF);
+ Span mockSpan = Mockito.mock(Span.class);
+ SpanContext childSpanContext =
+ SpanContext.create(traceId, childSpanId, TraceFlags.getSampled(), TraceState.getDefault());
+ Mockito.when(mockSpan.getSpanContext()).thenReturn(childSpanContext);
+ Mockito.when(mockSpan.storeInContext(Mockito.any())).thenCallRealMethod();
+ Mockito.when(mockSpan.makeCurrent()).thenCallRealMethod();
+ Mockito.when(mockOpenTelemetry.getTracer("PubSubIO")).thenReturn(mockTracer);
+ Mockito.when(mockTracer.spanBuilder("PubSubIO.Read")).thenReturn(mockSpanBuilder);
+ Mockito.when(mockSpanBuilder.startSpan()).thenReturn(mockSpan);
+
+ PipelineOptions options = TestPipeline.testingPipelineOptions();
+ options.as(SdkHarnessOptions.class).setOpenTelemetry(mockOpenTelemetry);
+
+ PubsubIO.Read.OpenTelemetryHeaderConsumer consumer =
+ new PubsubIO.Read.OpenTelemetryHeaderConsumer();
+ consumer.setup(options);
+
+ PubsubMessage message =
+ new PubsubMessage(
+ "payload".getBytes(StandardCharsets.UTF_8),
+ ImmutableMap.of("googclient_traceparent", traceparent));
+ List<PubsubMessage> outputs = new ArrayList<>();
+ @SuppressWarnings("unchecked")
+ DoFn.OutputReceiver<PubsubMessage> receiver = Mockito.mock(DoFn.OutputReceiver.class);
+ Mockito.doAnswer(
+ invocation -> {
+ assertEquals(childSpanId, Span.current().getSpanContext().getSpanId());
+ outputs.add(invocation.getArgument(0));
+ return null;
+ })
+ .when(receiver)
+ .output(Mockito.any());
+ consumer.processElement(message, receiver);
+
+ assertEquals(1, outputs.size());
+ Mockito.verify(mockTracer).spanBuilder("PubSubIO.Read");
+ Mockito.verify(mockSpanBuilder).setSpanKind(SpanKind.CONSUMER);
+ ArgumentCaptor<Context> parentContextCaptor = ArgumentCaptor.forClass(Context.class);
+ Mockito.verify(mockSpanBuilder).setParent(parentContextCaptor.capture());
+ assertEquals(
+ parentSpanId,
+ Span.fromContext(parentContextCaptor.getValue()).getSpanContext().getSpanId());
+ Mockito.verify(mockSpan).end();
+ }
+
+ @Test
+ public void testOpenTelemetryHeaderPropagator() {
+ String traceId = "0af7651916cd43dd8448eb211c80319c";
+ String producerSpanId = "2222222222222222";
+
+ OpenTelemetry mockOpenTelemetry = Mockito.mock(OpenTelemetry.class);
+ Tracer mockTracer = Mockito.mock(Tracer.class);
+ SpanBuilder mockSpanBuilder = Mockito.mock(SpanBuilder.class, Mockito.RETURNS_SELF);
+ Span mockSpan = Mockito.mock(Span.class);
+ SpanContext producerSpanContext =
+ SpanContext.create(
+ traceId, producerSpanId, TraceFlags.getSampled(), TraceState.getDefault());
+ Mockito.when(mockSpan.getSpanContext()).thenReturn(producerSpanContext);
+ Mockito.when(mockSpan.storeInContext(Mockito.any())).thenCallRealMethod();
+ Mockito.when(mockSpan.makeCurrent()).thenCallRealMethod();
+ Mockito.when(mockOpenTelemetry.getTracer("PubSubIO")).thenReturn(mockTracer);
+ Mockito.when(mockTracer.spanBuilder("PubSubIO.Write")).thenReturn(mockSpanBuilder);
+ Mockito.when(mockSpanBuilder.startSpan()).thenReturn(mockSpan);
+
+ PipelineOptions options = TestPipeline.testingPipelineOptions();
+ options.as(SdkHarnessOptions.class).setOpenTelemetry(mockOpenTelemetry);
+
+ PubsubIO.Write.OpenTelemetryHeaderPropagator propagator =
+ new PubsubIO.Write.OpenTelemetryHeaderPropagator();
+ propagator.setup(options);
+
+ PubsubMessage message =
+ new PubsubMessage("payload".getBytes(StandardCharsets.UTF_8), ImmutableMap.of("k", "v"));
+ List<PubsubMessage> outputs = new ArrayList<>();
+ @SuppressWarnings("unchecked")
+ DoFn.OutputReceiver<PubsubMessage> receiver = Mockito.mock(DoFn.OutputReceiver.class);
+ Mockito.doAnswer(
+ invocation -> {
+ assertEquals(producerSpanId, Span.current().getSpanContext().getSpanId());
+ outputs.add(invocation.getArgument(0));
+ return null;
+ })
+ .when(receiver)
+ .output(Mockito.any());
+ propagator.processElement(message, receiver);
+
+ assertEquals(1, outputs.size());
+ PubsubMessage propagated = outputs.get(0);
+ assertEquals("v", propagated.getAttribute("k"));
+ assertEquals(
+ "00-" + traceId + "-" + producerSpanId + "-01",
+ propagated.getAttribute("googclient_traceparent"));
+ Mockito.verify(mockTracer).spanBuilder("PubSubIO.Write");
+ Mockito.verify(mockSpanBuilder).setSpanKind(SpanKind.PRODUCER);
+ Mockito.verify(mockSpan).end();
+ }
}
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index 99f8a72e8d1..a41db9fbeaf 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -27,6 +27,7 @@ import com.google.auto.value.AutoValue;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import io.confluent.kafka.serializers.KafkaAvroDeserializer;
import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
import io.opentelemetry.context.Context;
@@ -142,6 +143,7 @@ import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.apache.kafka.common.utils.AppInfoParser;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.checkerframework.dataflow.qual.Pure;
import org.joda.time.Duration;
@@ -2251,7 +2253,7 @@ public class KafkaIO {
static class OpenTelemetryHeaderConsumer<K, V>
extends DoFn<KafkaRecord<K, V>, KafkaRecord<K, V>> {
- @Nullable Tracer tracer = null;
+ private transient @MonotonicNonNull Tracer tracer = null;
@Setup
public void setup(PipelineOptions options) {
@@ -2304,6 +2306,7 @@ public class KafkaIO {
Span span =
Preconditions.checkArgumentNotNull(tracer)
.spanBuilder("KafkaIO.Read")
+ .setSpanKind(SpanKind.CONSUMER)
.setParent(context)
.startSpan();
try (Scope ignored = span.makeCurrent()) {
@@ -2316,6 +2319,18 @@ public class KafkaIO {
static class OpenTelemetryHeaderPropagator<K, V>
extends DoFn<ProducerRecord<K, V>, ProducerRecord<K, V>> {
+ transient @MonotonicNonNull Tracer tracer = null;
+
+ @Setup
+ public void setup(PipelineOptions options) {
+ // inject tracer via options
+ io.opentelemetry.api.OpenTelemetry openTelemetry =
+ options.as(SdkHarnessOptions.class).getOpenTelemetry();
+ if (openTelemetry != null) {
+ tracer = openTelemetry.getTracer("KafkaIO");
+ }
+ }
+
ProducerRecord<K, V> injectTraceContext(ProducerRecord<K, V> message) {
org.apache.kafka.common.header.internals.RecordHeaders headers =
new org.apache.kafka.common.header.internals.RecordHeaders(message.headers());
@@ -2340,7 +2355,17 @@ public class KafkaIO {
@ProcessElement
public void processElement(
@Element ProducerRecord<K, V> element, OutputReceiver<ProducerRecord<K, V>> receiver) {
- receiver.output(injectTraceContext(element));
+ Span span =
+ Preconditions.checkArgumentNotNull(tracer)
+ .spanBuilder("KafkaIO.Write")
+ .setSpanKind(SpanKind.PRODUCER)
+ .setParent(Context.current())
+ .startSpan();
+ try (Scope ignored = span.makeCurrent()) {
+ receiver.output(injectTraceContext(element));
+ } finally {
+ span.end();
+ }
}
}
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index 1e320681414..8e56f08f7a1 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -42,6 +42,15 @@ import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient;
import io.confluent.kafka.schemaregistry.testutil.MockSchemaRegistry;
import io.confluent.kafka.serializers.AbstractKafkaAvroSerDeConfig;
import io.confluent.kafka.serializers.KafkaAvroSerializer;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
import java.io.IOException;
import java.lang.reflect.Method;
import java.nio.ByteBuffer;
@@ -88,7 +97,9 @@ import org.apache.beam.sdk.metrics.MetricResult;
import org.apache.beam.sdk.metrics.MetricsFilter;
import org.apache.beam.sdk.metrics.SinkMetrics;
import org.apache.beam.sdk.metrics.SourceMetrics;
+import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.SdkHarnessOptions;
import org.apache.beam.sdk.options.StreamingOptions;
import org.apache.beam.sdk.schemas.NoSuchSchemaException;
import org.apache.beam.sdk.schemas.Schema;
@@ -142,6 +153,7 @@ import org.apache.kafka.common.errors.SerializationException;
import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.record.TimestampType;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.Deserializer;
@@ -164,6 +176,8 @@ import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -2544,6 +2558,117 @@ public class KafkaIOTest {
assertEquals(Schema.Field.nullable("value", Schema.FieldType.BYTES), schema.getField(1));
}
+ @Test
+ public void testOpenTelemetryHeaderConsumer() {
+ String traceId = "0af7651916cd43dd8448eb211c80319c";
+ String parentSpanId = "b7ad6b7169203331";
+ String childSpanId = "1111111111111111";
+ String traceparent = "00-" + traceId + "-" + parentSpanId + "-01";
+
+ OpenTelemetry mockOpenTelemetry = Mockito.mock(OpenTelemetry.class);
+ Tracer mockTracer = Mockito.mock(Tracer.class);
+ SpanBuilder mockSpanBuilder = Mockito.mock(SpanBuilder.class, Mockito.RETURNS_SELF);
+ Span mockSpan = Mockito.mock(Span.class);
+ SpanContext childSpanContext =
+ SpanContext.create(traceId, childSpanId, TraceFlags.getSampled(), TraceState.getDefault());
+ Mockito.when(mockSpan.getSpanContext()).thenReturn(childSpanContext);
+ Mockito.when(mockSpan.storeInContext(Mockito.any())).thenCallRealMethod();
+ Mockito.when(mockSpan.makeCurrent()).thenCallRealMethod();
+ Mockito.when(mockOpenTelemetry.getTracer("KafkaIO")).thenReturn(mockTracer);
+ Mockito.when(mockTracer.spanBuilder("KafkaIO.Read")).thenReturn(mockSpanBuilder);
+ Mockito.when(mockSpanBuilder.startSpan()).thenReturn(mockSpan);
+
+ PipelineOptions options = PipelineOptionsFactory.create();
+ options.as(SdkHarnessOptions.class).setOpenTelemetry(mockOpenTelemetry);
+
+ KafkaIO.OpenTelemetryHeaderConsumer<String, String> consumer =
+ new KafkaIO.OpenTelemetryHeaderConsumer<>();
+ consumer.setup(options);
+
+ RecordHeaders headers = new RecordHeaders();
+ headers.add("traceparent", traceparent.getBytes(StandardCharsets.UTF_8));
+ KafkaRecord<String, String> record =
+ new KafkaRecord<>(
+ "topic", 0, 0L, 0L, KafkaTimestampType.CREATE_TIME, headers, "key", "value");
+
+ List<KafkaRecord<String, String>> outputs = new ArrayList<>();
+ @SuppressWarnings("unchecked")
+ DoFn.OutputReceiver<KafkaRecord<String, String>> receiver =
+ Mockito.mock(DoFn.OutputReceiver.class);
+ Mockito.doAnswer(
+ invocation -> {
+ assertEquals(childSpanId, Span.current().getSpanContext().getSpanId());
+ outputs.add(invocation.getArgument(0));
+ return null;
+ })
+ .when(receiver)
+ .output(Mockito.any());
+ consumer.processElement(record, receiver);
+
+ assertEquals(1, outputs.size());
+ Mockito.verify(mockTracer).spanBuilder("KafkaIO.Read");
+ Mockito.verify(mockSpanBuilder).setSpanKind(SpanKind.CONSUMER);
+ ArgumentCaptor<Context> parentContextCaptor = ArgumentCaptor.forClass(Context.class);
+ Mockito.verify(mockSpanBuilder).setParent(parentContextCaptor.capture());
+ assertEquals(
+ parentSpanId,
+ Span.fromContext(parentContextCaptor.getValue()).getSpanContext().getSpanId());
+ Mockito.verify(mockSpan).end();
+ }
+
+ @Test
+ public void testOpenTelemetryHeaderPropagator() {
+ String traceId = "0af7651916cd43dd8448eb211c80319c";
+ String producerSpanId = "2222222222222222";
+
+ OpenTelemetry mockOpenTelemetry = Mockito.mock(OpenTelemetry.class);
+ Tracer mockTracer = Mockito.mock(Tracer.class);
+ SpanBuilder mockSpanBuilder = Mockito.mock(SpanBuilder.class, Mockito.RETURNS_SELF);
+ Span mockSpan = Mockito.mock(Span.class);
+ SpanContext producerSpanContext =
+ SpanContext.create(
+ traceId, producerSpanId, TraceFlags.getSampled(), TraceState.getDefault());
+ Mockito.when(mockSpan.getSpanContext()).thenReturn(producerSpanContext);
+ Mockito.when(mockSpan.storeInContext(Mockito.any())).thenCallRealMethod();
+ Mockito.when(mockSpan.makeCurrent()).thenCallRealMethod();
+ Mockito.when(mockOpenTelemetry.getTracer("KafkaIO")).thenReturn(mockTracer);
+ Mockito.when(mockTracer.spanBuilder("KafkaIO.Write")).thenReturn(mockSpanBuilder);
+ Mockito.when(mockSpanBuilder.startSpan()).thenReturn(mockSpan);
+
+ PipelineOptions options = PipelineOptionsFactory.create();
+ options.as(SdkHarnessOptions.class).setOpenTelemetry(mockOpenTelemetry);
+
+ KafkaIO.OpenTelemetryHeaderPropagator<String, String> propagator =
+ new KafkaIO.OpenTelemetryHeaderPropagator<>();
+ propagator.setup(options);
+
+ ProducerRecord<String, String> record = new ProducerRecord<>("topic", "key", "value");
+ List<ProducerRecord<String, String>> outputs = new ArrayList<>();
+ @SuppressWarnings("unchecked")
+ DoFn.OutputReceiver<ProducerRecord<String, String>> receiver =
+ Mockito.mock(DoFn.OutputReceiver.class);
+ Mockito.doAnswer(
+ invocation -> {
+ assertEquals(producerSpanId, Span.current().getSpanContext().getSpanId());
+ outputs.add(invocation.getArgument(0));
+ return null;
+ })
+ .when(receiver)
+ .output(Mockito.any());
+ propagator.processElement(record, receiver);
+
+ assertEquals(1, outputs.size());
+ ProducerRecord<String, String> propagated = outputs.get(0);
+ Header traceparentHeader = propagated.headers().lastHeader("traceparent");
+ assertNotNull(traceparentHeader);
+ assertEquals(
+ "00-" + traceId + "-" + producerSpanId + "-01",
+ new String(traceparentHeader.value(), StandardCharsets.UTF_8));
+ Mockito.verify(mockTracer).spanBuilder("KafkaIO.Write");
+ Mockito.verify(mockSpanBuilder).setSpanKind(SpanKind.PRODUCER);
+ Mockito.verify(mockSpan).end();
+ }
+
private static void verifyProducerRecords(
MockProducer<Integer, Long> mockProducer,
String topic,