Refactor metadata propagation in ReduceFnRunner to support extensible PipelineMetadata

Author: stankiewicz

runners/core-java/src/main/java/org/apache/beam/runners/core/MetadataCombiner.java

@@ -0,0 +1,25 @@
+/*
+ * 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.runners.core;
+
+/** Interface for combining pipeline metadata. */
+interface MetadataCombiner<T> {
+  T createAccumulator();
+
+  T addInput(T accumulator, T input);
+}

runners/core-java/src/main/java/org/apache/beam/runners/core/PipelineMetadataCombiner.java

@@ -0,0 +1,65 @@
+/*
+ * 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.runners.core;
+
+import org.apache.beam.sdk.values.CausedByDrain;
+import org.apache.beam.sdk.values.PipelineMetadata;
+
+/** Combiner for PipelineMetadata. */
+class PipelineMetadataCombiner implements MetadataCombiner<PipelineMetadata> {
+  private static final PipelineMetadataCombiner INSTANCE = new PipelineMetadataCombiner();
+
+  public static PipelineMetadataCombiner of() {
+    return INSTANCE;
+  }
+
+  private final CausedByDrainCombiner causedByDrainCombiner = CausedByDrainCombiner.of();
+
+  @Override
+  public PipelineMetadata createAccumulator() {
+    return PipelineMetadata.create(causedByDrainCombiner.createAccumulator());
+  }
+
+  @Override
+  public PipelineMetadata addInput(PipelineMetadata accumulator, PipelineMetadata input) {
+    return PipelineMetadata.create(
+        causedByDrainCombiner.addInput(accumulator.causedByDrain(), input.causedByDrain()));
+  }
+
+  /** Combiner for CausedByDrain metadata. */
+  static class CausedByDrainCombiner implements MetadataCombiner<CausedByDrain> {
+    private static final CausedByDrainCombiner INSTANCE = new CausedByDrainCombiner();
+
+    public static CausedByDrainCombiner of() {
+      return INSTANCE;
+    }
+
+    @Override
+    public CausedByDrain createAccumulator() {
+      return CausedByDrain.NORMAL;
+    }
+
+    @Override
+    public CausedByDrain addInput(CausedByDrain current, CausedByDrain input) {
+      if (current == CausedByDrain.CAUSED_BY_DRAIN || input == CausedByDrain.CAUSED_BY_DRAIN) {
+        return CausedByDrain.CAUSED_BY_DRAIN;
+      }
+      return CausedByDrain.NORMAL;
+    }
+  }
+}

runners/core-java/src/main/java/org/apache/beam/runners/core/ReduceFnRunner.java

@@ -39,6 +39,7 @@
 import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.options.PipelineOptions;
 import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.ValueState;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
@@ -51,6 +52,7 @@
 import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PipelineMetadata;
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.sdk.values.WindowedValues;
 import org.apache.beam.sdk.values.WindowingStrategy;
@@ -107,6 +109,10 @@ public class ReduceFnRunner<K, InputT, OutputT, W extends BoundedWindow> {
    *   <li>It uses discarding or accumulation mode according to the {@link WindowingStrategy}.
    * </ul>
    */
+  static final StateTag<ValueState<PipelineMetadata>> METADATA_TAG =
+      StateTags.makeSystemTagInternal(
+          StateTags.value("pipelineMetadata", PipelineMetadata.Coder.of()));
+
   private final WindowingStrategy<Object, W> windowingStrategy;
 
   private final WindowedValueReceiver<KV<K, OutputT>> outputter;
@@ -376,7 +382,7 @@ public void processElements(Iterable<WindowedValue<InputT>> values) throws Excep
       emit(
           contextFactory.base(window, StateStyle.DIRECT),
           contextFactory.base(window, StateStyle.RENAMED),
-          CausedByDrain.NORMAL);
+          PipelineMetadata.createDefault());
     }
 
     // We're all done with merging and emitting elements so can compress the activeWindow state.
@@ -590,6 +596,17 @@ private void processElement(Map<W, W> windowToMergeResult, WindowedValue<InputT>
               value.getTimestamp(),
               StateStyle.DIRECT,
               value.causedByDrain());
+
+      ValueState<PipelineMetadata> metadataState = directContext.state().access(METADATA_TAG);
+      PipelineMetadata currentMetadata = metadataState.read();
+      if (currentMetadata == null) {
+        currentMetadata = PipelineMetadata.createDefault();
+      }
+      PipelineMetadata inputMetadata = PipelineMetadata.create(value.causedByDrain());
+      PipelineMetadata newMetadata =
+          PipelineMetadataCombiner.of().addInput(currentMetadata, inputMetadata);
+      metadataState.write(newMetadata);
+
       if (triggerRunner.isClosed(directContext.state())) {
         // This window has already been closed.
         droppedDueToClosedWindow.inc();
@@ -792,7 +809,7 @@ public void onTimers(Iterable<TimerData> timers) throws Exception {
                   renamedContext,
                   true /* isFinished */,
                   windowActivation.isEndOfWindow,
-                  windowActivation.causedByDrain);
+                  PipelineMetadata.create(windowActivation.causedByDrain));
           checkState(newHold == null, "Hold placed at %s despite isFinished being true.", newHold);
         }
 
@@ -810,7 +827,10 @@ public void onTimers(Iterable<TimerData> timers) throws Exception {
         if (windowActivation.windowIsActiveAndOpen()
             && triggerRunner.shouldFire(
                 directContext.window(), directContext.timers(), directContext.state())) {
-          emit(directContext, renamedContext, windowActivation.causedByDrain);
+          emit(
+              directContext,
+              renamedContext,
+              PipelineMetadata.create(windowActivation.causedByDrain));
         }
 
         if (windowActivation.isEndOfWindow) {
@@ -874,6 +894,7 @@ private void clearAllState(
       triggerRunner.clearState(
           directContext.window(), directContext.timers(), directContext.state());
       paneInfoTracker.clear(directContext.state());
+      directContext.state().access(METADATA_TAG).clear();
     } else {
       // If !windowIsActiveAndOpen then !activeWindows.isActive (1) or triggerRunner.isClosed (2).
       // For (1), if !activeWindows.isActive then the window must be merging and has been
@@ -934,8 +955,9 @@ private void prefetchEmit(
   private void emit(
       ReduceFn<K, InputT, OutputT, W>.Context directContext,
       ReduceFn<K, InputT, OutputT, W>.Context renamedContext,
-      CausedByDrain causedByDrain)
+      PipelineMetadata metadata)
       throws Exception {
+
     checkState(
         triggerRunner.shouldFire(
             directContext.window(), directContext.timers(), directContext.state()));
@@ -950,13 +972,14 @@ private void emit(
     // Run onTrigger to produce the actual pane contents.
     // As a side effect it will clear all element holds, but not necessarily any
     // end-of-window or garbage collection holds.
-    onTrigger(directContext, renamedContext, isFinished, false /*isEndOfWindow*/, causedByDrain);
+    onTrigger(directContext, renamedContext, isFinished, false /*isEndOfWindow*/, metadata);
 
     // Now that we've triggered, the pane is empty.
     nonEmptyPanes.clearPane(renamedContext.state());
 
     // Cleanup buffered data if appropriate
     if (shouldDiscard) {
+      directContext.state().access(METADATA_TAG).clear();
       // Cleanup flavor C: The user does not want any buffered data to persist between panes.
       reduceFn.clearState(renamedContext);
     }
@@ -1009,8 +1032,19 @@ private void prefetchOnTrigger(
       ReduceFn<K, InputT, OutputT, W>.Context renamedContext,
       final boolean isFinished,
       boolean isEndOfWindow,
-      CausedByDrain causedByDrain)
+      PipelineMetadata metadata)
       throws Exception {
+    ValueState<PipelineMetadata> metadataState = directContext.state().access(METADATA_TAG);
+    PipelineMetadata aggregatedMetadata = metadataState.read();
+    if (aggregatedMetadata == null) {
+      aggregatedMetadata = PipelineMetadata.createDefault();
+    }
+    PipelineMetadata fullyAggregatedMetadata =
+        PipelineMetadataCombiner.of().addInput(aggregatedMetadata, metadata);
+    final CausedByDrain aggregatedCausedByDrain = fullyAggregatedMetadata.causedByDrain();
+    if (isFinished) {
+      metadataState.clear();
+    }
     // Extract the window hold, and as a side effect clear it.
     final WatermarkHold.OldAndNewHolds pair =
         watermarkHold.extractAndRelease(renamedContext, isFinished).read();
@@ -1081,12 +1115,12 @@ private void prefetchOnTrigger(
                     .setValue(KV.of(key, toOutput))
                     .setTimestamp(outputTimestamp)
                     .setWindows(windows)
-                    .setCausedByDrain(causedByDrain)
+                    .setCausedByDrain(aggregatedCausedByDrain)
                     .setPaneInfo(paneInfo)
                     .setReceiver(outputter)
                     .output();
               },
-              causedByDrain);
+              aggregatedCausedByDrain);
 
       reduceFn.onTrigger(renamedTriggerContext);
     }

runners/core-java/src/test/java/org/apache/beam/runners/core/ReduceFnTester.java

@@ -318,7 +318,7 @@ public boolean hasNoActiveWindows() {
   public final void assertHasOnlyGlobalAndFinishedSetsFor(W... expectedWindows) {
     assertHasOnlyGlobalAndAllowedTags(
         ImmutableSet.copyOf(expectedWindows),
-        ImmutableSet.of(TriggerStateMachineRunner.FINISHED_BITS_TAG));
+        ImmutableSet.of(TriggerStateMachineRunner.FINISHED_BITS_TAG, ReduceFnRunner.METADATA_TAG));
   }
 
   @SafeVarargs
@@ -331,7 +331,8 @@ public final void assertHasOnlyGlobalAndStateFor(W... expectedWindows) {
             PaneInfoTracker.PANE_INFO_TAG,
             WatermarkHold.watermarkHoldTagForTimestampCombiner(
                 objectStrategy.getTimestampCombiner()),
-            WatermarkHold.EXTRA_HOLD_TAG));
+            WatermarkHold.EXTRA_HOLD_TAG,
+            ReduceFnRunner.METADATA_TAG));
   }
 
   @SafeVarargs

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

@@ -0,0 +1,85 @@
+/*
+ * 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 com.google.auto.value.AutoValue;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
+import org.apache.beam.sdk.coders.AtomicCoder;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
+
+/**
+ * Encapsulates metadata that propagates with elements in the pipeline.
+ *
+ * <p>This metadata is sent along with {@link WindowedValue} elements. It currently includes fields
+ * like {@link CausedByDrain}, and is designed to be extensible to support future metadata fields
+ * such as OpenTelemetry context or CDC (Change Data Capture) kind.
+ *
+ * <p>The purpose of this class is to group targeted metadata fields together. This makes it easier
+ * to define combination strategies (e.g., when accumulating state in {@code ReduceFnRunner}) when
+ * multiple elements are merged or grouped, without having to extend method signatures or state
+ * handling for every new metadata field.
+ */
+@AutoValue
+public abstract class PipelineMetadata {
+  public abstract CausedByDrain causedByDrain();
+
+  public static PipelineMetadata create(CausedByDrain causedByDrain) {
+    return new AutoValue_PipelineMetadata(causedByDrain);
+  }
+
+  public static PipelineMetadata createDefault() {
+    return create(CausedByDrain.NORMAL);
+  }
+
+  public static class Coder extends AtomicCoder<PipelineMetadata> {
+    private static final Coder INSTANCE = new Coder();
+
+    public static Coder of() {
+      return INSTANCE;
+    }
+
+    @Override
+    public void encode(PipelineMetadata value, OutputStream outStream) throws IOException {
+      BeamFnApi.Elements.ElementMetadata.Builder builder =
+          BeamFnApi.Elements.ElementMetadata.newBuilder();
+      builder.setDrain(
+          value.causedByDrain() == CausedByDrain.CAUSED_BY_DRAIN
+              ? BeamFnApi.Elements.DrainMode.Enum.DRAINING
+              : BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING);
+
+      ByteArrayCoder.of().encode(builder.build().toByteArray(), outStream);
+    }
+
+    @Override
+    public PipelineMetadata decode(InputStream inStream) throws IOException {
+      byte[] bytes = ByteArrayCoder.of().decode(inStream);
+      BeamFnApi.Elements.ElementMetadata proto =
+          BeamFnApi.Elements.ElementMetadata.parseFrom(bytes);
+
+      CausedByDrain causedByDrain =
+          proto.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING
+              ? CausedByDrain.CAUSED_BY_DRAIN
+              : CausedByDrain.NORMAL;
+
+      return PipelineMetadata.create(causedByDrain);
+    }
+  }
+}

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

@@ -864,6 +864,10 @@ public static void setMetadataSupported() {
       metadataSupported = true;
     }
 
+    public static void setMetadataNotSupported() {
+      metadataSupported = false;
+    }
+
     public static boolean isMetadataSupported() {
       return metadataSupported;
     }