Skip to content

Commit 226ad87

Browse files
stankiewiczstankiewicz
authored
enable otel context propagation - runner v1 sink, source changes, doFnRunner changes for per element propagation (#39152)
1 parent 6741628 commit 226ad87

9 files changed

Lines changed: 128 additions & 19 deletions

File tree

‎runners/core-java/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ dependencies {
4949
implementation library.java.slf4j_api
5050
implementation library.java.jackson_core
5151
implementation library.java.jackson_databind
52+
implementation library.java.opentelemetry_context
5253
implementation library.java.hamcrest
5354
testImplementation project(path: ":sdks:java:core", configuration: "shadowTest")
5455
testImplementation library.java.junit

‎runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
2222
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
2323

24+
import io.opentelemetry.context.Context;
25+
import io.opentelemetry.context.Scope;
2426
import java.util.Collection;
2527
import java.util.Collections;
2628
import java.util.HashMap;
@@ -184,6 +186,17 @@ public void startBundle() {
184186

185187
@Override
186188
public void processElement(WindowedValue<InputT> compressedElem) {
189+
Context openTelemetryContext = compressedElem.getOpenTelemetryContext();
190+
if (openTelemetryContext == null) {
191+
processElementInternal(compressedElem);
192+
} else {
193+
try (Scope ignore = openTelemetryContext.makeCurrent()) {
194+
processElementInternal(compressedElem);
195+
}
196+
}
197+
}
198+
199+
private void processElementInternal(WindowedValue<InputT> compressedElem) {
187200
if (observesWindow) {
188201
for (WindowedValue<InputT> elem : compressedElem.explodeWindows()) {
189202
invokeProcessElement(elem);

‎runners/google-cloud-dataflow-java/worker/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@ dependencies {
229229
implementation library.java.jackson_databind
230230
implementation library.java.joda_time
231231
implementation library.java.opentelemetry_context
232+
implementation library.java.opentelemetry_api
232233
implementation library.java.slf4j_api
233234
implementation library.java.vendored_grpc_1_69_0
234235
implementation library.java.error_prone_annotations

‎runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
2222

2323
import com.google.auto.service.AutoService;
24+
import io.opentelemetry.context.Context;
2425
import java.io.IOException;
2526
import java.io.InputStream;
2627
import java.util.Collection;
@@ -133,6 +134,7 @@ protected WindowedValue<T> decodeMessage(Windmill.Message message) throws IOExce
133134
*/
134135
CausedByDrain drainingValueFromUpstream = CausedByDrain.NORMAL;
135136
ValueKind valueKind = ValueKind.INSERT;
137+
Context openTelemetryContext = null;
136138
if (WindowedValues.WindowedValueCoder.isMetadataSupported()) {
137139
BeamFnApi.Elements.ElementMetadata elementMetadata =
138140
WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata());
@@ -141,6 +143,7 @@ protected WindowedValue<T> decodeMessage(Windmill.Message message) throws IOExce
141143
? CausedByDrain.CAUSED_BY_DRAIN
142144
: CausedByDrain.NORMAL;
143145
valueKind = WindmillValueKindHelper.fromProto(elementMetadata.getValueKind());
146+
openTelemetryContext = WindmillOpenTelemetryContextPropagator.read(elementMetadata);
144147
}
145148
if (valueCoder instanceof KvCoder) {
146149
KvCoder<?, ?> kvCoder = (KvCoder<?, ?>) valueCoder;
@@ -159,7 +162,7 @@ protected WindowedValue<T> decodeMessage(Windmill.Message message) throws IOExce
159162
null,
160163
null,
161164
drainingValueFromUpstream,
162-
null,
165+
openTelemetryContext,
163166
valueKind);
164167
} else {
165168
notifyElementRead(data.available() + metadata.available());
@@ -172,7 +175,7 @@ protected WindowedValue<T> decodeMessage(Windmill.Message message) throws IOExce
172175
null,
173176
null,
174177
drainingValueFromUpstream,
175-
null,
178+
openTelemetryContext,
176179
valueKind);
177180
}
178181
}

‎runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
2121

22+
import io.opentelemetry.context.Context;
2223
import java.io.IOException;
2324
import java.io.InputStream;
2425
import java.io.OutputStream;
@@ -150,6 +151,7 @@ public Iterable<TimerData> timersIterable() {
150151
*/
151152
CausedByDrain drainingValueFromUpstream = CausedByDrain.NORMAL;
152153
ValueKind valueKind = ValueKind.INSERT;
154+
Context openTelemetryContext = null;
153155
if (WindowedValues.WindowedValueCoder.isMetadataSupported()) {
154156
BeamFnApi.Elements.ElementMetadata elementMetadata =
155157
WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata());
@@ -158,6 +160,7 @@ public Iterable<TimerData> timersIterable() {
158160
? CausedByDrain.CAUSED_BY_DRAIN
159161
: CausedByDrain.NORMAL;
160162
valueKind = WindmillValueKindHelper.fromProto(elementMetadata.getValueKind());
163+
openTelemetryContext = WindmillOpenTelemetryContextPropagator.read(elementMetadata);
161164
}
162165
InputStream inputStream = message.getData().newInput();
163166
ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER);
@@ -169,7 +172,7 @@ public Iterable<TimerData> timersIterable() {
169172
null,
170173
null,
171174
drainingValueFromUpstream,
172-
null,
175+
openTelemetryContext,
173176
valueKind);
174177
} catch (RuntimeException | IOException e) {
175178
if (!skipUndecodableElements) {

‎runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillOpenTelemetryContextPropagator.java‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.runners.dataflow.worker;
19+
20+
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
21+
import io.opentelemetry.context.Context;
22+
import io.opentelemetry.context.propagation.TextMapGetter;
23+
import io.opentelemetry.context.propagation.TextMapSetter;
24+
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
25+
import org.apache.beam.sdk.annotations.Internal;
26+
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
27+
import org.checkerframework.checker.nullness.qual.Nullable;
28+
29+
@Internal
30+
public class WindmillOpenTelemetryContextPropagator {
31+
32+
public static final String TRACEPARENT = "traceparent";
33+
public static final String TRACESTATE = "tracestate";
34+
private static final TextMapSetter<BeamFnApi.Elements.ElementMetadata.Builder> SETTER =
35+
(carrier, key, value) -> {
36+
if (carrier == null) {
37+
return;
38+
}
39+
if (TRACEPARENT.equals(key)) {
40+
carrier.setTraceparent(value);
41+
} else if (TRACESTATE.equals(key)) {
42+
carrier.setTracestate(value);
43+
}
44+
};
45+
46+
private static final TextMapGetter<BeamFnApi.Elements.ElementMetadata> GETTER =
47+
new TextMapGetter<BeamFnApi.Elements.ElementMetadata>() {
48+
@Override
49+
public Iterable<String> keys(BeamFnApi.Elements.ElementMetadata carrier) {
50+
return Lists.newArrayList(TRACEPARENT, TRACESTATE);
51+
}
52+
53+
@Override
54+
public @Nullable String get(
55+
BeamFnApi.Elements.@Nullable ElementMetadata carrier, String key) {
56+
if (carrier == null) {
57+
return null;
58+
}
59+
if (TRACEPARENT.equalsIgnoreCase(key)) {
60+
return carrier.getTraceparent();
61+
} else if (TRACESTATE.equalsIgnoreCase(key)) {
62+
return carrier.getTracestate();
63+
}
64+
return null;
65+
}
66+
};
67+
68+
public static void set(Context from, BeamFnApi.Elements.ElementMetadata.Builder builder) {
69+
W3CTraceContextPropagator.getInstance().inject(from, builder, SETTER);
70+
}
71+
72+
public static Context read(BeamFnApi.Elements.ElementMetadata from) {
73+
return W3CTraceContextPropagator.getInstance().extract(Context.root(), from, GETTER);
74+
}
75+
}

‎runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java‎

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
2222

2323
import com.google.auto.service.AutoService;
24+
import io.opentelemetry.context.Context;
2425
import java.io.IOException;
2526
import java.io.InputStream;
2627
import java.nio.charset.StandardCharsets;
@@ -220,17 +221,27 @@ public long add(WindowedValue<T> data) throws IOException {
220221
ByteString key, value;
221222
ByteString id = ByteString.EMPTY;
222223
// todo #33176 specify additional metadata in the future
223-
BeamFnApi.Elements.ElementMetadata additionalMetadata =
224-
BeamFnApi.Elements.ElementMetadata.newBuilder()
225-
.setDrain(
226-
data.causedByDrain() == CausedByDrain.CAUSED_BY_DRAIN
227-
? BeamFnApi.Elements.DrainMode.Enum.DRAINING
228-
: BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING)
229-
.setValueKind(WindmillValueKindHelper.toProto(data.getValueKind()))
230-
.build();
224+
BeamFnApi.Elements.ElementMetadata.Builder additionalMetadataBuilder =
225+
BeamFnApi.Elements.ElementMetadata.newBuilder();
226+
additionalMetadataBuilder
227+
.setDrain(
228+
data.causedByDrain() == CausedByDrain.CAUSED_BY_DRAIN
229+
? BeamFnApi.Elements.DrainMode.Enum.DRAINING
230+
: BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING)
231+
.setValueKind(WindmillValueKindHelper.toProto(data.getValueKind()));
232+
Context openTelemetryContext = data.getOpenTelemetryContext();
233+
if (openTelemetryContext != null) {
234+
// TODO replace with OpenTelemetryContextPropagator
235+
WindmillOpenTelemetryContextPropagator.set(openTelemetryContext, additionalMetadataBuilder);
236+
}
237+
231238
ByteString metadata =
232239
encodeMetadata(
233-
stream, windowsCoder, data.getWindows(), data.getPaneInfo(), additionalMetadata);
240+
stream,
241+
windowsCoder,
242+
data.getWindows(),
243+
data.getPaneInfo(),
244+
additionalMetadataBuilder.build());
234245
if (valueCoder instanceof KvCoder) {
235246
KvCoder kvCoder = (KvCoder) valueCoder;
236247
KV kv = checkNotNull((KV) data.getValue());

‎sdks/java/core/src/main/java/org/apache/beam/sdk/values/OpenTelemetryContextPropagator.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,12 @@
2222
import io.opentelemetry.context.propagation.TextMapGetter;
2323
import io.opentelemetry.context.propagation.TextMapSetter;
2424
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
25+
import org.apache.beam.sdk.annotations.Internal;
2526
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
2627
import org.checkerframework.checker.nullness.qual.Nullable;
2728

28-
class OpenTelemetryContextPropagator {
29+
@Internal
30+
public class OpenTelemetryContextPropagator {
2931

3032
private static final TextMapSetter<BeamFnApi.Elements.ElementMetadata.Builder> SETTER =
3133
(carrier, key, value) -> {
@@ -61,11 +63,11 @@ public Iterable<String> keys(BeamFnApi.Elements.ElementMetadata carrier) {
6163
}
6264
};
6365

64-
static void set(Context from, BeamFnApi.Elements.ElementMetadata.Builder builder) {
66+
public static void set(Context from, BeamFnApi.Elements.ElementMetadata.Builder builder) {
6567
W3CTraceContextPropagator.getInstance().inject(from, builder, SETTER);
6668
}
6769

68-
static Context read(BeamFnApi.Elements.ElementMetadata from) {
70+
public static Context read(BeamFnApi.Elements.ElementMetadata from) {
6971
return W3CTraceContextPropagator.getInstance().extract(Context.root(), from, GETTER);
7072
}
7173
}

‎sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,6 @@ public static class Builder<T> implements OutputBuilder<T> {
107107
private @Nullable String recordId;
108108
private @Nullable Long recordOffset;
109109
private CausedByDrain causedByDrain = CausedByDrain.NORMAL;
110-
private @Nullable Context openTelemetryContext;
111110
private ValueKind valueKind = ValueKind.INSERT;
112111

113112
@Override
@@ -160,8 +159,7 @@ public Builder<T> setCausedByDrain(CausedByDrain causedByDrain) {
160159
}
161160

162161
@Override
163-
public Builder<T> setOpenTelemetryContext(@Nullable Context openTelemetryContext) {
164-
this.openTelemetryContext = openTelemetryContext;
162+
public Builder<T> setOpenTelemetryContext(@Nullable Context ignored) {
165163
return this;
166164
}
167165

@@ -200,7 +198,9 @@ public Instant getTimestamp() {
200198

201199
@Override
202200
public @Nullable Context getOpenTelemetryContext() {
203-
return openTelemetryContext;
201+
// builder may have different context set at the beginning of parDo
202+
// when building WindowedValue we should take current context from storage.
203+
return Context.current();
204204
}
205205

206206
@Override

0 commit comments

Comments
 (0)