Model changes to allow OpenTelemtry context propagation

stankiewicz · stankiewicz · commit 65e3fddd7227 · 2025-12-02T14:37:55.000+01:00
diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -855,8 +855,13 @@ class BeamModulePlugin implements Plugin<Project> {
         netty_tcnative_boringssl_static             : "io.netty:netty-tcnative-boringssl-static:2.0.52.Final",
         netty_transport                             : "io.netty:netty-transport:$netty_version",
         netty_transport_native_epoll                : "io.netty:netty-transport-native-epoll:$netty_version",
-        opentelemetry_api                           : "io.opentelemetry:opentelemetry-api", // google_cloud_platform_libraries_bom sets version
+        opentelemetry_api                           : "io.opentelemetry:opentelemetry-api:$opentelemetry_version",
         opentelemetry_bom                           : "io.opentelemetry:opentelemetry-bom-alpha:$opentelemetry_version-alpha", // alpha required by extensions
+        opentelemetry_context                       : "io.opentelemetry:opentelemetry-context:$opentelemetry_version",
+        opentelemetry_gcp_auth                      : "io.opentelemetry.contrib:opentelemetry-gcp-auth-extension:$opentelemetry_version-alpha",
+        opentelemetry_sdk                           : "io.opentelemetry:opentelemetry-sdk:$opentelemetry_version",
+        opentelemetry_exporter_otlp                 : "io.opentelemetry:opentelemetry-exporter-otlp:$opentelemetry_version",
+        opentelemetry_extension_autoconfigure       : "io.opentelemetry:opentelemetry-sdk-extension-autoconfigure:$opentelemetry_version",
         postgres                                    : "org.postgresql:postgresql:$postgres_version",
         protobuf_java                               : "com.google.protobuf:protobuf-java:$protobuf_version",
         protobuf_java_util                          : "com.google.protobuf:protobuf-java-util:$protobuf_version",
diff --git a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto
@@ -761,6 +761,8 @@ message Elements {
   // extensible and backward compatible
   message ElementMetadata {
     optional DrainMode.Enum drain = 1;
+    optional string traceparent = 2;
+    optional string tracestate = 3;
   }
 
   // Represent the encoded user timer for a given instruction, transform and
diff --git a/runners/google-cloud-dataflow-java/build.gradle b/runners/google-cloud-dataflow-java/build.gradle
@@ -112,6 +112,7 @@ dependencies {
   implementation library.java.google_http_client
   implementation library.java.google_http_client_gson
   permitUnusedDeclared library.java.google_http_client_gson // BEAM-11761
+  implementation library.java.opentelemetry_context
   implementation library.java.hamcrest
   implementation library.java.jackson_annotations
   implementation library.java.jackson_core
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java
@@ -20,6 +20,7 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
+import io.opentelemetry.context.Context;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -1403,6 +1404,11 @@ public PaneInfo getPaneInfo() {
       return null;
     }
 
+    @Override
+    public @Nullable Context getContext() {
+      return null;
+    }
+
     @Override
     public @Nullable Long getRecordOffset() {
       return null;
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/ValueInEmptyWindows.java
@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.dataflow.worker.util;
 
+import io.opentelemetry.context.Context;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Objects;
@@ -64,6 +65,11 @@ public boolean causedByDrain() {
     return false;
   }
 
+  @Override
+  public @Nullable Context getContext() {
+    return null;
+  }
+
   @Override
   public Iterable<WindowedValue<T>> explodeWindows() {
     return Collections.emptyList();
diff --git a/runners/spark/spark_runner.gradle b/runners/spark/spark_runner.gradle
@@ -167,6 +167,7 @@ dependencies {
   implementation library.java.jackson_annotations
   implementation library.java.slf4j_api
   implementation library.java.joda_time
+  implementation library.java.opentelemetry_context
   implementation library.java.commons_lang3
   implementation library.java.args4j
   implementation project(path: ":model:fn-execution", configuration: "shadow")
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java
@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.spark.util;
 
+import io.opentelemetry.context.Context;
 import java.io.Serializable;
 import java.util.Collection;
 import java.util.Collections;
@@ -115,6 +116,11 @@ public PaneInfo getPaneInfo() {
       return null;
     }
 
+    @Override
+    public @Nullable Context getContext() {
+      return null;
+    }
+
     @Override
     public boolean causedByDrain() {
       return false;
diff --git a/sdks/java/core/build.gradle b/sdks/java/core/build.gradle
@@ -98,6 +98,7 @@ dependencies {
   shadow library.java.jackson_databind
   shadow platform(library.java.opentelemetry_bom)
   shadow library.java.opentelemetry_api
+  shadow library.java.opentelemetry_context
   shadow library.java.slf4j_api
   shadow library.java.snappy_java
   shadow library.java.joda_time
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java
@@ -490,7 +490,7 @@ public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWind
           getMutableOutput(tag)
               .add(
                   ValueInSingleWindow.of(
-                      output, timestamp, window, PaneInfo.NO_FIRING, null, null));
+                      output, timestamp, window, PaneInfo.NO_FIRING, null, null, null));
         }
       };
     }
@@ -623,7 +623,7 @@ public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp
       getMutableOutput(tag)
           .add(
               ValueInSingleWindow.of(
-                  output, timestamp, element.getWindow(), element.getPaneInfo(), null, null));
+                  output, timestamp, element.getWindow(), element.getPaneInfo(), null, null, null));
     }
 
     @Override
@@ -635,7 +635,7 @@ public <T> void outputWindowedValue(
         PaneInfo paneInfo) {
       for (BoundedWindow w : windows) {
         getMutableOutput(tag)
-            .add(ValueInSingleWindow.of(output, timestamp, w, paneInfo, null, null));
+            .add(ValueInSingleWindow.of(output, timestamp, w, paneInfo, null, null, null));
       }
     }
   }
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OutputBuilder.java
@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.values;
 
+import io.opentelemetry.context.Context;
 import java.util.Collection;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
@@ -50,5 +51,7 @@ public interface OutputBuilder<T> extends WindowedValue<T> {
 
   OutputBuilder<T> setCausedByDrain(boolean causedByDrain);
 
+  OutputBuilder<T> setContext(@Nullable Context context);
+
   void output();
 }
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java
@@ -18,6 +18,7 @@
 package org.apache.beam.sdk.values;
 
 import com.google.auto.value.AutoValue;
+import io.opentelemetry.context.Context;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -66,21 +67,24 @@ public T getValue() {
 
   public abstract @Nullable Long getCurrentRecordOffset();
 
+  public abstract @Nullable Context getContext();
+
   // todo #33176 specify additional metadata in the future
   public static <T> ValueInSingleWindow<T> of(
       T value,
       Instant timestamp,
       BoundedWindow window,
       PaneInfo paneInfo,
       @Nullable String currentRecordId,
-      @Nullable Long currentRecordOffset) {
+      @Nullable Long currentRecordOffset,
+      @Nullable Context context) {
     return new AutoValue_ValueInSingleWindow<>(
-        value, timestamp, window, paneInfo, currentRecordId, currentRecordOffset);
+        value, timestamp, window, paneInfo, currentRecordId, currentRecordOffset, context);
   }
 
   public static <T> ValueInSingleWindow<T> of(
       T value, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) {
-    return of(value, timestamp, window, paneInfo, null, null);
+    return of(value, timestamp, window, paneInfo, null, null, null);
   }
 
   /** A coder for {@link ValueInSingleWindow}. */
@@ -105,11 +109,14 @@ public static <T> Coder<T> of(
     @Override
     public void encode(ValueInSingleWindow<T> windowedElem, OutputStream outStream)
         throws IOException {
-      encode(windowedElem, outStream, Context.NESTED);
+      encode(windowedElem, outStream, org.apache.beam.sdk.coders.Coder.Context.NESTED);
     }
 
     @Override
-    public void encode(ValueInSingleWindow<T> windowedElem, OutputStream outStream, Context context)
+    public void encode(
+        ValueInSingleWindow<T> windowedElem,
+        OutputStream outStream,
+        org.apache.beam.sdk.coders.Coder.Context context)
         throws IOException {
       InstantCoder.of().encode(windowedElem.getTimestamp(), outStream);
       windowCoder.encode(windowedElem.getWindow(), outStream);
@@ -120,6 +127,10 @@ public void encode(ValueInSingleWindow<T> windowedElem, OutputStream outStream,
         BeamFnApi.Elements.ElementMetadata.Builder builder =
             BeamFnApi.Elements.ElementMetadata.newBuilder();
         // todo #33176 specify additional metadata in the future
+        io.opentelemetry.context.Context context1 = windowedElem.getContext();
+        if (context1 != null) {
+          WindowedValues.OpenTelemetryContextSerializer.write(context1, builder);
+        }
         BeamFnApi.Elements.ElementMetadata metadata = builder.build();
         ByteArrayCoder.of().encode(metadata.toByteArray(), outStream);
       }
@@ -129,22 +140,27 @@ public void encode(ValueInSingleWindow<T> windowedElem, OutputStream outStream,
 
     @Override
     public ValueInSingleWindow<T> decode(InputStream inStream) throws IOException {
-      return decode(inStream, Context.NESTED);
+      return decode(inStream, org.apache.beam.sdk.coders.Coder.Context.NESTED);
     }
 
     @Override
     @SuppressWarnings("IgnoredPureGetter")
-    public ValueInSingleWindow<T> decode(InputStream inStream, Context context) throws IOException {
+    public ValueInSingleWindow<T> decode(
+        InputStream inStream, org.apache.beam.sdk.coders.Coder.Context context) throws IOException {
       Instant timestamp = InstantCoder.of().decode(inStream);
       BoundedWindow window = windowCoder.decode(inStream);
       PaneInfo paneInfo = PaneInfo.PaneInfoCoder.INSTANCE.decode(inStream);
+      io.opentelemetry.context.Context openTelemetryContext = null;
       if (WindowedValues.WindowedValueCoder.isMetadataSupported() && paneInfo.isElementMetadata()) {
-        BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream));
+        BeamFnApi.Elements.ElementMetadata elementMetadata =
+            BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream));
+        openTelemetryContext = WindowedValues.OpenTelemetryContextSerializer.read(elementMetadata);
       }
 
       T value = valueCoder.decode(inStream, context);
       // todo #33176 specify additional metadata in the future
-      return new AutoValue_ValueInSingleWindow<>(value, timestamp, window, paneInfo, null, null);
+      return new AutoValue_ValueInSingleWindow<>(
+          value, timestamp, window, paneInfo, null, null, openTelemetryContext);
     }
 
     @Override
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValue.java
@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.values;
 
+import io.opentelemetry.context.Context;
 import java.util.Collection;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
@@ -49,6 +50,9 @@ public interface WindowedValue<T> {
   @Nullable
   String getRecordId();
 
+  @Nullable
+  Context getContext();
+
   @Nullable
   Long getRecordOffset();
 
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/WindowedValueTest.java

Original file line number	Diff line number	Diff line change
`@@ -761,6 +761,8 @@ message Elements {`
`761`	`761`	`// extensible and backward compatible`
`762`	`762`	`message ElementMetadata {`
`763`	`763`	`optional DrainMode.Enum drain = 1;`
	`764`	`+ optional string traceparent = 2;`
	`765`	`+ optional string tracestate = 3;`
`764`	`766`	`}`
`765`	`767`
`766`	`768`	`// Represent the encoded user timer for a given instruction, transform and`
Original file line number	Diff line number	Diff line change
`@@ -490,7 +490,7 @@ public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWind`
`490`	`490`	`getMutableOutput(tag)`
`491`	`491`	`.add(`
`492`	`492`	`ValueInSingleWindow.of(`
`493`		`- output, timestamp, window, PaneInfo.NO_FIRING, null, null));`
	`493`	`+ output, timestamp, window, PaneInfo.NO_FIRING, null, null, null));`
`494`	`494`	`}`
`495`	`495`	`};`
`496`	`496`	`}`
`@@ -623,7 +623,7 @@ public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp`
`623`	`623`	`getMutableOutput(tag)`
`624`	`624`	`.add(`
`625`	`625`	`ValueInSingleWindow.of(`
`626`		`- output, timestamp, element.getWindow(), element.getPaneInfo(), null, null));`
	`626`	`+ output, timestamp, element.getWindow(), element.getPaneInfo(), null, null, null));`
`627`	`627`	`}`
`628`	`628`
`629`	`629`	`@Override`
`@@ -635,7 +635,7 @@ public <T> void outputWindowedValue(`
`635`	`635`	`PaneInfo paneInfo) {`
`636`	`636`	`for (BoundedWindow w : windows) {`
`637`	`637`	`getMutableOutput(tag)`
`638`		`- .add(ValueInSingleWindow.of(output, timestamp, w, paneInfo, null, null));`
	`638`	`+ .add(ValueInSingleWindow.of(output, timestamp, w, paneInfo, null, null, null));`
`639`	`639`	`}`
`640`	`640`	`}`
`641`	`641`	`}`