Skip to content

[OpenTelemetry] Fix OpenTelemetry span kinds and add write spans in PubsubIO and KafkaIO #40483

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stankiewicz merged 2 commits into from
Oct 13, 2026
Merged

[OpenTelemetry] Fix OpenTelemetry span kinds and add write spans in PubsubIO and KafkaIO #40483

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8071ce2
Fix OpenTelemetry span kinds and add write spans in PubsubIO and KafkaIO
stankiewicz Oct 9, 2026
a90d3ce
address comments
stankiewicz Oct 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 33 additions & 14 deletions ...va/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
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;
Expand Down Expand Up @@ -1170,6 +1171,7 @@ public void processElement(
Span span =
checkArgumentNotNull(tracer)
.spanBuilder("PubSubIO.Read")
.setSpanKind(SpanKind.CONSUMER)
.setParent(context)
.startSpan();
try (Scope s = span.makeCurrent()) {
Expand Down Expand Up @@ -1573,24 +1575,41 @@ public void set(@Nullable Map<String, String> attr, String key, String value) {
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);
}
}

Expand Down
121 changes: 121 additions & 0 deletions ...o/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,19 @@
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;
Expand Down Expand Up @@ -62,11 +72,13 @@
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;
Expand All @@ -82,6 +94,7 @@
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;
Expand All @@ -98,6 +111,7 @@
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. */
Expand Down Expand Up @@ -1083,4 +1097,111 @@ public void testWithoutValidation() throws IOException {

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();
}
}
29 changes: 27 additions & 2 deletions sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
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;
Expand Down Expand Up @@ -142,6 +143,7 @@
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;
Expand Down Expand Up @@ -2251,7 +2253,7 @@ public void populateDisplayData(DisplayData.Builder builder) {

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) {
Expand Down Expand Up @@ -2304,6 +2306,7 @@ public void processElement(
Span span =
Preconditions.checkArgumentNotNull(tracer)
.spanBuilder("KafkaIO.Read")
.setSpanKind(SpanKind.CONSUMER)
.setParent(context)
.startSpan();
try (Scope ignored = span.makeCurrent()) {
Expand All @@ -2316,6 +2319,18 @@ public void processElement(

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());
Expand All @@ -2340,7 +2355,17 @@ ProducerRecord<K, V> injectTraceContext(ProducerRecord<K, V> message) {
@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();
}
}
}

Expand Down
Loading
Loading