[Java] Dataflow runner v1 - Propagate drain mode (#36534)

Author: stankiewicz

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

@@ -389,6 +389,7 @@ private StreamingWorkerHarnessFactoryOutput createFanOutStreamingEngineWorkerHar
                 serializedWorkItemSize,
                 watermarks,
                 processingContext,
+                drainMode,
                 getWorkStreamLatencies) ->
                 computationStateCache
                     .get(processingContext.computationId())
@@ -401,6 +402,7 @@ private StreamingWorkerHarnessFactoryOutput createFanOutStreamingEngineWorkerHar
                               serializedWorkItemSize,
                               watermarks,
                               processingContext,
+                              drainMode,
                               getWorkStreamLatencies);
                         }),
             ChannelCachingRemoteStubFactory.create(options.getGcpCredential(), channelCache),

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

@@ -196,6 +196,10 @@ public boolean workIsFailed() {
     return work != null && work.isFailed();
   }
 
+  public boolean getDrainMode() {
+    return work != null ? work.getDrainMode() : false;
+  }
+
   public boolean offsetBasedDeduplicationSupported() {
     return activeReader != null
         && activeReader.getCurrentSource().offsetBasedDeduplicationSupported();
@@ -820,7 +824,10 @@ public <W extends BoundedWindow> TimerData getNextFiredTimer(Coder<W> windowCode
                 .transform(
                     timer ->
                         WindmillTimerInternals.windmillTimerToTimerData(
-                            WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, timer, windowCoder))
+                            WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX,
+                            timer,
+                            windowCoder,
+                            getDrainMode()))
                 .iterator();
       }
 
@@ -880,7 +887,10 @@ public <W extends BoundedWindow> TimerData getNextFiredUserTimer(Coder<W> window
                     .transform(
                         timer ->
                             WindmillTimerInternals.windmillTimerToTimerData(
-                                WindmillNamespacePrefix.USER_NAMESPACE_PREFIX, timer, windowCoder))
+                                WindmillNamespacePrefix.USER_NAMESPACE_PREFIX,
+                                timer,
+                                windowCoder,
+                                getDrainMode()))
                     .iterator());
       }
 

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

@@ -24,6 +24,7 @@
 import java.io.InputStream;
 import java.util.Collection;
 import java.util.Map;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
 import org.apache.beam.runners.dataflow.util.CloudObject;
 import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
@@ -117,8 +118,16 @@ protected WindowedValue<T> decodeMessage(Windmill.Message message) throws IOExce
       Collection<? extends BoundedWindow> windows =
           WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata());
       PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata());
+      /**
+       * https://s.apache.org/beam-drain-mode - propagate drain bit if aggregation/expiry induced by
+       * drain happened upstream
+       */
+      boolean drainingValueFromUpstream = false;
       if (WindowedValues.WindowedValueCoder.isMetadataSupported()) {
-        WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata());
+        BeamFnApi.Elements.ElementMetadata elementMetadata =
+            WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata());
+        drainingValueFromUpstream =
+            elementMetadata.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING;
       }
       if (valueCoder instanceof KvCoder) {
         KvCoder<?, ?> kvCoder = (KvCoder<?, ?>) valueCoder;
@@ -128,12 +137,18 @@ 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 #33176 propagate metadata to windowed value
-        return WindowedValues.of(result, timestampMillis, windows, paneInfo);
+        return WindowedValues.of(
+            result, timestampMillis, windows, paneInfo, null, null, drainingValueFromUpstream);
       } else {
         notifyElementRead(data.available() + metadata.available());
-        // todo #33176 propagate metadata to windowed value
-        return WindowedValues.of(decode(valueCoder, data), timestampMillis, windows, paneInfo);
+        return WindowedValues.of(
+            decode(valueCoder, data),
+            timestampMillis,
+            windows,
+            paneInfo,
+            null,
+            null,
+            drainingValueFromUpstream);
       }
     }
 

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

@@ -24,6 +24,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
+import org.apache.beam.model.fnexecution.v1.BeamFnApi;
 import org.apache.beam.runners.core.KeyedWorkItem;
 import org.apache.beam.runners.core.KeyedWorkItemCoder;
 import org.apache.beam.runners.core.TimerInternals.TimerData;
@@ -60,6 +61,8 @@ public class WindmillKeyedWorkItem<K, ElemT> implements KeyedWorkItem<K, ElemT>
 
   private final Windmill.WorkItem workItem;
   private final K key;
+  // used to inform that timer was caused by drain
+  private final boolean drainMode;
 
   private final transient Coder<? extends BoundedWindow> windowCoder;
   private final transient Coder<Collection<? extends BoundedWindow>> windowsCoder;
@@ -70,12 +73,14 @@ public WindmillKeyedWorkItem(
       Windmill.WorkItem workItem,
       Coder<? extends BoundedWindow> windowCoder,
       Coder<Collection<? extends BoundedWindow>> windowsCoder,
-      Coder<ElemT> valueCoder) {
+      Coder<ElemT> valueCoder,
+      boolean drainMode) {
     this.key = key;
     this.workItem = workItem;
     this.windowCoder = windowCoder;
     this.windowsCoder = windowsCoder;
     this.valueCoder = valueCoder;
+    this.drainMode = drainMode;
   }
 
   @Override
@@ -93,7 +98,10 @@ public Iterable<TimerData> timersIterable() {
         .transform(
             timer ->
                 WindmillTimerInternals.windmillTimerToTimerData(
-                    WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX, timer, windowCoder));
+                    WindmillNamespacePrefix.SYSTEM_NAMESPACE_PREFIX,
+                    timer,
+                    windowCoder,
+                    drainMode));
   }
 
   @Override
@@ -108,13 +116,21 @@ public Iterable<WindowedValue<ElemT>> elementsIterable() {
                 Collection<? extends BoundedWindow> windows =
                     WindmillSink.decodeMetadataWindows(windowsCoder, message.getMetadata());
                 PaneInfo paneInfo = WindmillSink.decodeMetadataPane(message.getMetadata());
+                /**
+                 * https://s.apache.org/beam-drain-mode - propagate drain bit if aggregation/expiry
+                 * induced by drain happened upstream
+                 */
+                boolean drainingValueFromUpstream = false;
                 if (WindowedValues.WindowedValueCoder.isMetadataSupported()) {
-                  WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata());
+                  BeamFnApi.Elements.ElementMetadata elementMetadata =
+                      WindmillSink.decodeAdditionalMetadata(windowsCoder, message.getMetadata());
+                  drainingValueFromUpstream =
+                      elementMetadata.getDrain() == BeamFnApi.Elements.DrainMode.Enum.DRAINING;
                 }
                 InputStream inputStream = message.getData().newInput();
                 ElemT value = valueCoder.decode(inputStream, Coder.Context.OUTER);
-                // todo #33176 specify additional metadata in the future
-                return WindowedValues.of(value, timestamp, windows, paneInfo);
+                return WindowedValues.of(
+                    value, timestamp, windows, paneInfo, null, null, drainingValueFromUpstream);
               } catch (IOException e) {
                 throw new RuntimeException(e);
               }

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

@@ -301,7 +301,10 @@ static Timer timerDataToWindmillTimer(
   }
 
   public static TimerData windmillTimerToTimerData(
-      WindmillNamespacePrefix prefix, Timer timer, Coder<? extends BoundedWindow> windowCoder) {
+      WindmillNamespacePrefix prefix,
+      Timer timer,
+      Coder<? extends BoundedWindow> windowCoder,
+      boolean draining) {
 
     // The tag is a path-structure string but cheaper to parse than a proper URI. It follows
     // this pattern, where no component but the ID can contain a slash
@@ -395,6 +398,8 @@ public static TimerData windmillTimerToTimerData(
         timestamp,
         outputTimestamp,
         timerTypeToTimeDomain(timer.getType()));
+    // todo add draining (https://github.com/apache/beam/issues/36884)
+
   }
 
   private static boolean useNewTimerTagEncoding(TimerData timerData) {

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

@@ -119,7 +119,8 @@ public NativeReaderIterator<WindowedValue<KeyedWorkItem<K, T>>> iterator() throw
     final K key = keyCoder.decode(context.getSerializedKey().newInput(), Coder.Context.OUTER);
     final WorkItem workItem = context.getWorkItem();
     KeyedWorkItem<K, T> keyedWorkItem =
-        new WindmillKeyedWorkItem<>(key, workItem, windowCoder, windowsCoder, valueCoder);
+        new WindmillKeyedWorkItem<>(
+            key, workItem, windowCoder, windowsCoder, valueCoder, context.getDrainMode());
     final boolean isEmptyWorkItem =
         (Iterables.isEmpty(keyedWorkItem.timersIterable())
             && Iterables.isEmpty(keyedWorkItem.elementsIterable()));

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java

@@ -78,19 +78,22 @@ public final class Work implements RefreshableWork {
   private volatile TimedState currentState;
   private volatile boolean isFailed;
   private volatile String processingThreadName = "";
+  private final boolean drainMode;
 
   private Work(
       WorkItem workItem,
       long serializedWorkItemSize,
       Watermarks watermarks,
       ProcessingContext processingContext,
+      boolean drainMode,
       Supplier<Instant> clock) {
     this.shardedKey = ShardedKey.create(workItem.getKey(), workItem.getShardingKey());
     this.workItem = workItem;
     this.serializedWorkItemSize = serializedWorkItemSize;
     this.processingContext = processingContext;
     this.watermarks = watermarks;
     this.clock = clock;
+    this.drainMode = drainMode;
     this.startTime = clock.get();
     Preconditions.checkState(EMPTY_ENUM_MAP.isEmpty());
     // Create by passing EMPTY_ENUM_MAP to avoid recreating
@@ -110,8 +113,10 @@ public static Work create(
       long serializedWorkItemSize,
       Watermarks watermarks,
       ProcessingContext processingContext,
+      boolean drainMode,
       Supplier<Instant> clock) {
-    return new Work(workItem, serializedWorkItemSize, watermarks, processingContext, clock);
+    return new Work(
+        workItem, serializedWorkItemSize, watermarks, processingContext, drainMode, clock);
   }
 
   public static ProcessingContext createProcessingContext(
@@ -207,6 +212,10 @@ public State getState() {
     return currentState.state();
   }
 
+  public boolean getDrainMode() {
+    return drainMode;
+  }
+
   public void setState(State state) {
     Instant now = clock.get();
     totalDurationPerState.compute(

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/SingleSourceWorkerHarness.java

@@ -155,6 +155,7 @@ private void streamingEngineDispatchLoop(
               (computationId,
                   inputDataWatermark,
                   synchronizedProcessingTime,
+                  drainMode,
                   workItem,
                   serializedWorkItemSize,
                   getWorkStreamLatencies) ->
@@ -178,6 +179,7 @@ private void streamingEngineDispatchLoop(
                                     getDataClient,
                                     workCommitter::commit,
                                     heartbeatSender),
+                                drainMode,
                                 getWorkStreamLatencies);
                           }));
       try {
@@ -239,6 +241,7 @@ private void applianceDispatchLoop(Supplier<Windmill.GetWorkResponse> getWorkFn)
               watermarks.setOutputDataWatermark(workItem.getOutputDataWatermark()).build(),
               Work.createProcessingContext(
                   computationId, getDataClient, workCommitter::commit, heartbeatSender),
+              computationWork.getDrainMode(),
               /* getWorkStreamLatencies= */ ImmutableList.of());
         }
       }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GetWorkResponseChunkAssembler.java

@@ -124,14 +124,17 @@ private static ComputationMetadata fromProto(
           metadataProto.getComputationId(),
           WindmillTimeUtils.windmillToHarnessWatermark(metadataProto.getInputDataWatermark()),
           WindmillTimeUtils.windmillToHarnessWatermark(
-              metadataProto.getDependentRealtimeInputWatermark()));
+              metadataProto.getDependentRealtimeInputWatermark()),
+          metadataProto.getDrainMode());
     }
 
     abstract String computationId();
 
     abstract @Nullable Instant inputDataWatermark();
 
     abstract @Nullable Instant synchronizedProcessingTime();
+
+    abstract boolean drainMode();
   }
 
   @AutoValue

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java

@@ -281,6 +281,7 @@ private void consumeAssembledWorkItem(AssembledWorkItem assembledWorkItem) {
         assembledWorkItem.bufferedSize(),
         createWatermarks(workItem, metadata),
         createProcessingContext(metadata.computationId()),
+        metadata.drainMode(),
         assembledWorkItem.latencyAttributions());
     budgetTracker.recordBudgetReceived(assembledWorkItem.bufferedSize());
     GetWorkBudget extension = budgetTracker.computeBudgetExtension();