[OpenTelemetry] Add tracing support to GCP PubsubIO#39150
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances observability in Apache Beam by integrating OpenTelemetry tracing into the Google Cloud Pub/Sub IO connector. By allowing trace context propagation through Pub/Sub message attributes, it enables seamless distributed tracing for data pipelines. Additionally, the change includes a necessary fix to the validation logic for error handling configurations. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces OpenTelemetry tracing support to PubsubIO's read and write transforms by extracting and propagating trace contexts via message attributes. However, several critical issues were identified in the implementation: a logic inversion bug was introduced in the dead-letter topic validation; the TextMapGetter implementation in OpenTelemetryHeaderConsumer has potential null-pointer issues, incorrect return values, and inconsistent key mapping; the newly created tracing span is not correctly made current; and OpenTelemetryHeaderPropagator is susceptible to NullPointerExceptions when message attributes are null.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if (getDeadLetterTopicProvider() != null | ||
| && !(getBadRecordRouter() instanceof ThrowingBadRecordRouter)) { | ||
| && (getBadRecordRouter() instanceof ThrowingBadRecordRouter)) { | ||
| throw new IllegalArgumentException( | ||
| "PubSubIO cannot be configured with both a dead letter topic and a bad record router"); | ||
| } |
There was a problem hiding this comment.
This change introduces a critical logic inversion bug. The original code threw an exception if a dead letter topic was configured AND a custom bad record router was used (i.e., not the default ThrowingBadRecordRouter). By changing !(...) to (...), the code now throws an exception when a dead letter topic is configured with the DEFAULT ThrowingBadRecordRouter, which completely breaks the standard dead letter topic configuration for all users.
| if (getDeadLetterTopicProvider() != null | |
| && !(getBadRecordRouter() instanceof ThrowingBadRecordRouter)) { | |
| && (getBadRecordRouter() instanceof ThrowingBadRecordRouter)) { | |
| throw new IllegalArgumentException( | |
| "PubSubIO cannot be configured with both a dead letter topic and a bad record router"); | |
| } | |
| if (getDeadLetterTopicProvider() != null | |
| && !(getBadRecordRouter() instanceof ThrowingBadRecordRouter)) { | |
| throw new IllegalArgumentException( | |
| "PubSubIO cannot be configured with both a dead letter topic and a bad record router"); | |
| } |
| Context extractSpanContext(PubsubMessage message) { | ||
| TextMapGetter<PubsubMessage> extractMessageAttributes = | ||
| new TextMapGetter<PubsubMessage>() { | ||
| @Override | ||
| public String get(@Nullable PubsubMessage carrier, String key) { | ||
| String attribute = | ||
| Preconditions.checkArgumentNotNull(carrier).getAttribute("googclient_" + key); | ||
| return attribute == null ? "" : attribute; | ||
| } | ||
|
|
||
| @Override | ||
| public Iterable<String> keys(PubsubMessage carrier) { | ||
| Map<String, String> attributeMap = carrier.getAttributeMap(); | ||
| return attributeMap == null ? ImmutableList.of() : attributeMap.keySet(); | ||
| } | ||
| }; | ||
| return W3CTraceContextPropagator.getInstance() | ||
| .extract(Context.current(), message, extractMessageAttributes); | ||
| } |
There was a problem hiding this comment.
There are three issues with the TextMapGetter implementation:
- If
carrieris null,getwill throw aNullPointerExceptiondue toPreconditions.checkArgumentNotNull(carrier). According to the OpenTelemetry specification,getshould returnnullif the carrier is null or the key is not found. - Returning
""(empty string) instead ofnullwhen the attribute is missing can cause the propagator to fail or log warnings because it tries to parse an empty string as a traceparent. keysreturns the raw keys from the attribute map (e.g.,"googclient_traceparent"). However,getprepends"googclient_"to the key. If a propagator iterates over the keys returned bykeysand callsget, it will look up"googclient_googclient_traceparent".keysshould filter and strip the"googclient_"prefix to remain consistent withget.
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) {
if (carrier == null) {
return ImmutableList.of();
}
Map<String, String> attributeMap = carrier.getAttributeMap();
if (attributeMap == null) {
return ImmutableList.of();
}
java.util.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);
}| try (Scope s = context.makeCurrent()) { | ||
| output.output(message); | ||
| } finally { |
There was a problem hiding this comment.
The scope is being created using context.makeCurrent() instead of psSub.makeCurrent(). This means the newly created span psSub is never made current, and any downstream operations or nested spans created during output.output(message) will not be children of psSub ("Process Element"). Instead, they will be direct children of the extracted parent context.
| try (Scope s = context.makeCurrent()) { | |
| output.output(message); | |
| } finally { | |
| try (Scope s = psSub.makeCurrent()) { | |
| output.output(message); | |
| } finally { |
| Map<String, String> attributeMap = message.getAttributeMap(); | ||
| Map<String, String> attr = new HashMap<>(Objects.requireNonNull(attributeMap)); | ||
| injectSpanContext(attr); |
There was a problem hiding this comment.
message.getAttributeMap() can return null if the PubsubMessage was constructed without attributes. Calling Objects.requireNonNull(attributeMap) will throw a NullPointerException at runtime. We should handle the null case gracefully by initializing an empty HashMap.
| Map<String, String> attributeMap = message.getAttributeMap(); | |
| Map<String, String> attr = new HashMap<>(Objects.requireNonNull(attributeMap)); | |
| injectSpanContext(attr); | |
| Map<String, String> attributeMap = message.getAttributeMap(); | |
| Map<String, String> attr = | |
| attributeMap == null ? new HashMap<>() : new HashMap<>(attributeMap); | |
| injectSpanContext(attr); |
| 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) { | ||
| Preconditions.checkArgumentNotNull(attr).put("googclient_" + key, value); | ||
| } | ||
| }; |
There was a problem hiding this comment.
Using Preconditions.checkArgumentNotNull(attr) inside the TextMapSetter is risky because the setter's contract allows attr to be @Nullable. It is safer to perform a null check and skip injection if attr is null.
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);
}
}
};
stankiewicz
force-pushed
the
otel_pubsub
branch
2 times, most recently
from
8859ba1 to
d902a71
Compare
fdfd
|
Assigning reviewers: R: @Abacn for label java. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
1 participant
This PR introduces OpenTelemetry (OTel) support for Google Cloud Pub/Sub IO, enabling distributed tracing and better observability for Pub/Sub operations within Apache Beam pipelines.
Configuration: Added withEnableOpenTelemetryTracing() configuration methods to both PubsubIO.Read and PubsubIO.Write builders to allow users to toggle OTel tracing.
Context Propagation (Write): Introduced the OpenTelemetryHeaderPropagator DoFn. When OTel is enabled, this function injects the current OpenTelemetry trace context into outbound PubsubMessage attributes (using the googclient_ prefix) via the W3C Trace Context propagator.
Context Extraction (Read): Introduced the OpenTelemetryHeaderConsumer DoFn. When OTel is enabled, this function extracts the trace context from incoming PubsubMessage attributes and sets it as the parent span context for the "Process Element" span, linking the distributed trace.
Validation Fix: Updated the validation logic in PubsubIO.Read.expand() to properly throw an IllegalArgumentException if both a dead letter topic and a custom bad record router are configured.
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.