Model changes to allow OpenTelemetry context propagation

stankiewicz · stankiewicz · commit d8b9e070e5d3 · 2026-03-11T10:52:47.000+01:00
add nullable
diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -634,6 +634,8 @@ class BeamModulePlugin implements Plugin<Project> {
     def netty_version = "4.1.124.Final"
     // [bomupgrader] determined by: io.opentelemetry:opentelemetry-sdk, consistent with: google_cloud_platform_libraries_bom
     def opentelemetry_version = "1.51.0"
+    def opentelemetry_sdk_version = "1.56.0"
+    def opentelemetry_contrib_version = "1.52.0"
     def postgres_version = "42.2.16"
     // [bomupgrader] determined by: com.google.protobuf:protobuf-java, consistent with: google_cloud_platform_libraries_bom
     def protobuf_version = "4.33.2"
@@ -855,8 +857,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_bom                           : "io.opentelemetry:opentelemetry-bom-alpha:$opentelemetry_version-alpha", // alpha required by extensions
+        opentelemetry_api                           : "io.opentelemetry:opentelemetry-api:$opentelemetry_sdk_version",
+        opentelemetry_bom                           : "io.opentelemetry:opentelemetry-bom-alpha:$opentelemetry_sdk_version-alpha", // alpha required by extensions
+        opentelemetry_context                       : "io.opentelemetry:opentelemetry-context:$opentelemetry_sdk_version",
+        opentelemetry_gcp_auth                      : "io.opentelemetry.contrib:opentelemetry-gcp-auth-extension:$opentelemetry_contrib_version-alpha",
+        opentelemetry_sdk                           : "io.opentelemetry:opentelemetry-sdk:$opentelemetry_sdk_version",
+        opentelemetry_exporter_otlp                 : "io.opentelemetry:opentelemetry-exporter-otlp:$opentelemetry_sdk_version",
+        opentelemetry_extension_autoconfigure       : "io.opentelemetry:opentelemetry-sdk-extension-autoconfigure:$opentelemetry_sdk_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/runners/google-cloud-dataflow-java/build.gradle b/runners/google-cloud-dataflow-java/build.gradle
@@ -112,6 +112,7 @@ dependencies {
   permitUnusedDeclared library.java.google_http_client_gson // BEAM-11761
   implementation library.java.google_cloud_logging
   permitUnusedDeclared library.java.google_cloud_logging // 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;
@@ -1404,6 +1405,11 @@ public PaneInfo getPaneInfo() {
       return null;
     }
 
+    @Override
+    public @Nullable Context getOpenTelemetryContext() {
+      return null;
+    }
+
     @Override
     public @Nullable Long getRecordOffset() {
       return null;
diff --git a/runners/google-cloud-dataflow-java/worker/build.gradle b/runners/google-cloud-dataflow-java/worker/build.gradle
@@ -212,6 +212,7 @@ dependencies {
     implementation library.java.jackson_core
     implementation library.java.jackson_databind
     implementation library.java.joda_time
+    implementation library.java.opentelemetry_context
     implementation library.java.slf4j_api
     implementation library.java.vendored_grpc_1_69_0
     implementation library.java.error_prone_annotations
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/UngroupedWindmillReader.java
@@ -155,18 +155,28 @@ protected WindowedValue<T> decodeMessage(Windmill.Message message) throws IOExce
         @SuppressWarnings("unchecked")
         T result =
             (T) KV.of(decode(kvCoder.getKeyCoder(), key), decode(kvCoder.getValueCoder(), data));
+        // todo #37030 parse context from previous stage
         return WindowedValues.of(
-            result, timestampMillis, windows, paneInfo, null, null, drainingValueFromUpstream);
+            result,
+            timestampMillis,
+            windows,
+            paneInfo,
+            null,
+            null,
+            drainingValueFromUpstream,
+            null);
       } else {
         notifyElementRead(data.available() + metadata.available());
+        // todo #37030 parse context from previous stage
         return WindowedValues.of(
             decode(valueCoder, data),
             timestampMillis,
             windows,
             paneInfo,
             null,
             null,
-            drainingValueFromUpstream);
+            drainingValueFromUpstream,
+            null);
       }
     }
 
diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillKeyedWorkItem.java
@@ -159,7 +159,7 @@ public Iterable<TimerData> timersIterable() {
       InputStream inputStream = message.getData().newInput();
       ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER);
       return WindowedValues.of(
-          value, timestamp, windows, paneInfo, null, null, drainingValueFromUpstream);
+          value, timestamp, windows, paneInfo, null, null, drainingValueFromUpstream, null);
     } catch (RuntimeException | IOException e) {
       if (!skipUndecodableElements) {
         throw new RuntimeException(e);
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;
@@ -65,6 +66,11 @@ public CausedByDrain causedByDrain() {
     return CausedByDrain.NORMAL;
   }
 
+  @Override
+  public @Nullable Context getOpenTelemetryContext() {
+    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;
@@ -116,6 +117,11 @@ public PaneInfo getPaneInfo() {
       return null;
     }
 
+    @Override
+    public @Nullable Context getOpenTelemetryContext() {
+      return null;
+    }
+
     @Override
     public CausedByDrain causedByDrain() {
       return CausedByDrain.NORMAL;
diff --git a/sdks/java/container/license_scripts/dep_urls_java.yaml b/sdks/java/container/license_scripts/dep_urls_java.yaml
@@ -66,12 +66,12 @@ org.eclipse.jgit:
     license: "https://www.eclipse.org/org/documents/edl-v10.html"
     type: "Eclipse Distribution License - v1.0"
 opentelemetry-bom:
-  '1.51.0':
-    license: "https://raw.githubusercontent.com/open-telemetry/opentelemetry-java/v1.51.0/LICENSE"
+  '1.56.0':
+    license: "https://raw.githubusercontent.com/open-telemetry/opentelemetry-java/v1.56.0/LICENSE"
     type: "Apache License 2.0"
 opentelemetry-bom-alpha:
-  '1.51.0-alpha':
-    license: "https://raw.githubusercontent.com/open-telemetry/opentelemetry-java/v1.51.0/LICENSE"
+   '1.56.0-alpha':
+    license: "https://raw.githubusercontent.com/open-telemetry/opentelemetry-java/v1.56.0/LICENSE"
     type: "Apache License 2.0"
 zstd-jni:
   '1.5.2-5':
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
@@ -497,7 +497,7 @@ public <T> void output(TupleTag<T> tag, T output, Instant timestamp, BoundedWind
                       PaneInfo.NO_FIRING,
                       null,
                       null,
-                      CausedByDrain.NORMAL));
+                      CausedByDrain.NORMAL, null));
         }
       };
     }
@@ -641,7 +641,7 @@ public <T> void outputWithTimestamp(TupleTag<T> tag, T output, Instant timestamp
                   element.getPaneInfo(),
                   null,
                   null,
-                  CausedByDrain.NORMAL));
+                  CausedByDrain.NORMAL, null));
     }
 
     @Override
@@ -655,7 +655,7 @@ public <T> void outputWindowedValue(
         getMutableOutput(tag)
             .add(
                 ValueInSingleWindow.of(
-                    output, timestamp, w, paneInfo, null, null, CausedByDrain.NORMAL));
+                    output, timestamp, w, paneInfo, null, null, CausedByDrain.NORMAL, null));
       }
     }
   }
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OpenTelemetryContextPropagator.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/OpenTelemetryContextPropagator.java
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.values;
+
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.TextMapSetter;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+class OpenTelemetryContextPropagator {
+
+  private static final TextMapSetter<BeamFnApi.Elements.ElementMetadata.Builder> SETTER =
+      (carrier, key, value) -> {
+        if (carrier == null) {
+          return;
+        }
+        if ("traceparent".equals(key)) {
+          carrier.setTraceparent(value);
+        } else if ("tracestate".equals(key)) {
+          carrier.setTracestate(value);
+        }
+      };
+
+  private static final TextMapGetter<BeamFnApi.Elements.ElementMetadata> GETTER =
+      new TextMapGetter<BeamFnApi.Elements.ElementMetadata>() {
+        @Override
+        public Iterable<String> keys(BeamFnApi.Elements.ElementMetadata carrier) {
+          return Lists.newArrayList("traceparent", "tracestate");
+        }
+
+        @Override
+        public @Nullable String get(
+            BeamFnApi.Elements.@Nullable ElementMetadata carrier, String key) {
+          if (carrier == null) {
+            return null;
+          }
+          if ("traceparent".equals(key)) {
+            return carrier.getTraceparent();
+          } else if ("tracestate".equals(key)) {
+            return carrier.getTracestate();
+          }
+          return null;
+        }
+      };
+
+  static void write(Context from, BeamFnApi.Elements.ElementMetadata.Builder builder) {
+    W3CTraceContextPropagator.getInstance().inject(from, builder, SETTER);
+  }
+
+  static Context read(BeamFnApi.Elements.ElementMetadata from) {
+    return W3CTraceContextPropagator.getInstance().extract(Context.root(), from, GETTER);
+  }
+}
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(CausedByDrain causedByDrain);
 
+  OutputBuilder<T> setOpenTelemetryContext(@Nullable Context openTelemetryContext);
+
   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;
@@ -68,6 +69,8 @@ public T getValue() {
 
   public abstract CausedByDrain getCausedByDrain();
 
+  public abstract @Nullable Context getOpenTelemetryContext();
+
   // todo #33176 specify additional metadata in the future
   public static <T> ValueInSingleWindow<T> of(
       T value,
@@ -76,14 +79,15 @@ public static <T> ValueInSingleWindow<T> of(
       PaneInfo paneInfo,
       @Nullable String currentRecordId,
       @Nullable Long currentRecordOffset,
-      CausedByDrain causedByDrain) {
+      CausedByDrain causedByDrain,
+      @Nullable Context openTelemetryContext) {
     return new AutoValue_ValueInSingleWindow<>(
-        value, timestamp, window, paneInfo, currentRecordId, currentRecordOffset, causedByDrain);
+        value, timestamp, window, paneInfo, currentRecordId, currentRecordOffset, causedByDrain, openTelemetryContext);
   }
 
   public static <T> ValueInSingleWindow<T> of(
       T value, Instant timestamp, BoundedWindow window, PaneInfo paneInfo) {
-    return of(value, timestamp, window, paneInfo, null, null, CausedByDrain.NORMAL);
+    return of(value, timestamp, window, paneInfo, null, null, CausedByDrain.NORMAL, null);
   }
 
   /** A coder for {@link ValueInSingleWindow}. */
@@ -108,11 +112,11 @@ 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, Coder.Context.NESTED);
     }
 
     @Override
-    public void encode(ValueInSingleWindow<T> windowedElem, OutputStream outStream, Context context)
+    public void encode(ValueInSingleWindow<T> windowedElem, OutputStream outStream, Coder.Context context)
         throws IOException {
       InstantCoder.of().encode(windowedElem.getTimestamp(), outStream);
       windowCoder.encode(windowedElem.getWindow(), outStream);
@@ -123,6 +127,12 @@ 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 openTelemetryContext =
+            windowedElem.getOpenTelemetryContext();
+        if (openTelemetryContext != null) {
+
+          OpenTelemetryContextPropagator.write(openTelemetryContext, builder);
+        }
         BeamFnApi.Elements.ElementMetadata metadata = builder.build();
         ByteArrayCoder.of().encode(metadata.toByteArray(), outStream);
       }
@@ -132,29 +142,31 @@ 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, Coder.Context.NESTED);
     }
 
     @Override
     @SuppressWarnings("IgnoredPureGetter")
-    public ValueInSingleWindow<T> decode(InputStream inStream, Context context) throws IOException {
+    public ValueInSingleWindow<T> decode(InputStream inStream, Coder.Context context)
+        throws IOException {
       Instant timestamp = InstantCoder.of().decode(inStream);
       BoundedWindow window = windowCoder.decode(inStream);
       PaneInfo paneInfo = PaneInfo.PaneInfoCoder.INSTANCE.decode(inStream);
       CausedByDrain causedByDrain = CausedByDrain.NORMAL;
+      io.opentelemetry.context. @Nullable Context openTelemetryContext = null;
       if (WindowedValues.WindowedValueCoder.isMetadataSupported() && paneInfo.isElementMetadata()) {
         BeamFnApi.Elements.ElementMetadata elementMetadata =
             BeamFnApi.Elements.ElementMetadata.parseFrom(ByteArrayCoder.of().decode(inStream));
         causedByDrain =
             elementMetadata.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING
                 ? CausedByDrain.CAUSED_BY_DRAIN
                 : CausedByDrain.NORMAL;
+        openTelemetryContext = OpenTelemetryContextPropagator.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, causedByDrain);
+      return new AutoValue_ValueInSingleWindow<>(value, timestamp, window, paneInfo, null, causedByDrain, 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
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
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/GcpApiSurfaceTest.java
