Merge pull request #37727: [Drain] Reduce Fn Runner changes

Author: kennknowles

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

@@ -438,7 +438,15 @@ public <T> void output(TupleTag<T> tag, T value) {
     @Override
     public <T> void outputWithTimestamp(TupleTag<T> tag, T value, Instant timestamp) {
       outputReceiver.output(
-          tag, WindowedValues.of(value, timestamp, element.getWindows(), element.getPaneInfo()));
+          tag,
+          WindowedValues.of(
+              value,
+              timestamp,
+              element.getWindows(),
+              element.getPaneInfo(),
+              element.getRecordId(),
+              element.getRecordOffset(),
+              element.causedByDrain()));
     }
 
     @Override
@@ -452,7 +460,16 @@ public <T> void outputWindowedValue(
       if (watermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
         ((TimestampObservingWatermarkEstimator) watermarkEstimator).observeTimestamp(timestamp);
       }
-      outputReceiver.output(tag, WindowedValues.of(value, timestamp, windows, paneInfo));
+      outputReceiver.output(
+          tag,
+          WindowedValues.of(
+              value,
+              timestamp,
+              windows,
+              paneInfo,
+              element.getRecordId(),
+              element.getRecordOffset(),
+              element.causedByDrain()));
     }
 
     private void noteOutput() {

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

@@ -22,6 +22,7 @@
 import org.apache.beam.sdk.state.Timers;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
+import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.joda.time.Instant;
 
@@ -59,6 +60,7 @@ public abstract class ProcessValueContext extends Context {
     /** Return the actual value being processed. */
     public abstract InputT value();
 
+    public abstract CausedByDrain causedByDrain();
     /** Return the timestamp associated with the value. */
     public abstract Instant timestamp();
   }
@@ -75,6 +77,7 @@ public abstract class OnTriggerContext extends Context {
     /** Returns the {@link PaneInfo} for the trigger firing being processed. */
     public abstract PaneInfo paneInfo();
 
+    public abstract CausedByDrain causedByDrain();
     /** Output the given value in the current window. */
     public abstract void output(OutputT value);
   }

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

@@ -99,13 +99,19 @@ public ReduceFn<K, InputT, OutputT, W>.Context base(W window, StateStyle style)
   }
 
   public ReduceFn<K, InputT, OutputT, W>.ProcessValueContext forValue(
-      W window, InputT value, Instant timestamp, StateStyle style) {
-    return new ProcessValueContextImpl(stateAccessor(window, style), value, timestamp);
+      W window, InputT value, Instant timestamp, StateStyle style, CausedByDrain causedByDrain) {
+    return new ProcessValueContextImpl(
+        stateAccessor(window, style), value, timestamp, causedByDrain);
   }
 
   public ReduceFn<K, InputT, OutputT, W>.OnTriggerContext forTrigger(
-      W window, PaneInfo paneInfo, StateStyle style, OnTriggerCallbacks<OutputT> callbacks) {
-    return new OnTriggerContextImpl(stateAccessor(window, style), paneInfo, callbacks);
+      W window,
+      PaneInfo paneInfo,
+      StateStyle style,
+      OnTriggerCallbacks<OutputT> callbacks,
+      CausedByDrain causedByDrain) {
+    return new OnTriggerContextImpl(
+        stateAccessor(window, style), paneInfo, callbacks, causedByDrain);
   }
 
   public ReduceFn<K, InputT, OutputT, W>.OnMergeContext forMerge(
@@ -358,14 +364,19 @@ private class ProcessValueContextImpl
     private final Instant timestamp;
     private final StateAccessorImpl<K, W> state;
     private final TimersImpl timers;
+    private final CausedByDrain causedByDrain;
 
     private ProcessValueContextImpl(
-        StateAccessorImpl<K, W> state, InputT value, Instant timestamp) {
+        StateAccessorImpl<K, W> state,
+        InputT value,
+        Instant timestamp,
+        CausedByDrain causedByDrain) {
       reduceFn.super();
       this.state = state;
       this.value = value;
       this.timestamp = timestamp;
       this.timers = new TimersImpl(state.namespace());
+      this.causedByDrain = causedByDrain;
     }
 
     @Override
@@ -393,6 +404,11 @@ public InputT value() {
       return value;
     }
 
+    @Override
+    public CausedByDrain causedByDrain() {
+      return causedByDrain;
+    }
+
     @Override
     public Instant timestamp() {
       return timestamp;
@@ -409,14 +425,19 @@ private class OnTriggerContextImpl extends ReduceFn<K, InputT, OutputT, W>.OnTri
     private final PaneInfo paneInfo;
     private final OnTriggerCallbacks<OutputT> callbacks;
     private final TimersImpl timers;
+    private final CausedByDrain causedByDrain;
 
     private OnTriggerContextImpl(
-        StateAccessorImpl<K, W> state, PaneInfo paneInfo, OnTriggerCallbacks<OutputT> callbacks) {
+        StateAccessorImpl<K, W> state,
+        PaneInfo paneInfo,
+        OnTriggerCallbacks<OutputT> callbacks,
+        CausedByDrain causedByDrain) {
       reduceFn.super();
       this.state = state;
       this.paneInfo = paneInfo;
       this.callbacks = callbacks;
       this.timers = new TimersImpl(state.namespace());
+      this.causedByDrain = causedByDrain;
     }
 
     @Override
@@ -444,6 +465,11 @@ public PaneInfo paneInfo() {
       return paneInfo;
     }
 
+    @Override
+    public CausedByDrain causedByDrain() {
+      return causedByDrain;
+    }
+
     @Override
     public void output(OutputT value) {
       callbacks.output(value);

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

@@ -48,6 +48,7 @@
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
 import org.apache.beam.sdk.util.WindowTracing;
 import org.apache.beam.sdk.util.WindowedValueReceiver;
+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.WindowedValue;
@@ -374,7 +375,8 @@ public void processElements(Iterable<WindowedValue<InputT>> values) throws Excep
     for (W window : windowsToFire) {
       emit(
           contextFactory.base(window, StateStyle.DIRECT),
-          contextFactory.base(window, StateStyle.RENAMED));
+          contextFactory.base(window, StateStyle.RENAMED),
+          null);
     }
 
     // We're all done with merging and emitting elements so can compress the activeWindow state.
@@ -583,7 +585,11 @@ private void processElement(Map<W, W> windowToMergeResult, WindowedValue<InputT>
     for (W window : windows) {
       ReduceFn<K, InputT, OutputT, W>.ProcessValueContext directContext =
           contextFactory.forValue(
-              window, value.getValue(), value.getTimestamp(), StateStyle.DIRECT);
+              window,
+              value.getValue(),
+              value.getTimestamp(),
+              StateStyle.DIRECT,
+              value.causedByDrain());
       if (triggerRunner.isClosed(directContext.state())) {
         // This window has already been closed.
         droppedDueToClosedWindow.inc();
@@ -601,7 +607,11 @@ private void processElement(Map<W, W> windowToMergeResult, WindowedValue<InputT>
       activeWindows.ensureWindowIsActive(window);
       ReduceFn<K, InputT, OutputT, W>.ProcessValueContext renamedContext =
           contextFactory.forValue(
-              window, value.getValue(), value.getTimestamp(), StateStyle.RENAMED);
+              window,
+              value.getValue(),
+              value.getTimestamp(),
+              StateStyle.RENAMED,
+              value.causedByDrain());
 
       nonEmptyPanes.recordContent(renamedContext.state());
       scheduleGarbageCollectionTimer(directContext);
@@ -639,12 +649,15 @@ private class WindowActivation {
     // garbage collect the window.  We'll consider any timer at or after the
     // end-of-window time to be a signal to garbage collect.
     public final boolean isGarbageCollection;
+    public final CausedByDrain causedByDrain;
 
     WindowActivation(
         ReduceFn<K, InputT, OutputT, W>.Context directContext,
-        ReduceFn<K, InputT, OutputT, W>.Context renamedContext) {
+        ReduceFn<K, InputT, OutputT, W>.Context renamedContext,
+        CausedByDrain causedByDrain) {
       this.directContext = directContext;
       this.renamedContext = renamedContext;
+      this.causedByDrain = causedByDrain;
       W window = directContext.window();
 
       // The output watermark is before the end of the window if it is either unknown
@@ -701,12 +714,13 @@ public void onTimers(Iterable<TimerData> timers) throws Exception {
 
       WindowTracing.debug(
           "ReduceFnRunner: Received timer key:{}; window:{}; data:{} with "
-              + "inputWatermark:{}; outputWatermark:{}",
+              + "inputWatermark:{}; outputWatermark:{}; draining:{}",
           key,
           window,
           timer,
           timerInternals.currentInputWatermarkTime(),
-          timerInternals.currentOutputWatermarkTime());
+          timerInternals.currentOutputWatermarkTime(),
+          timer.causedByDrain());
 
       // Processing time timers for an expired window are ignored, just like elements
       // that show up too late. Window GC is management by an event time timer
@@ -722,11 +736,13 @@ public void onTimers(Iterable<TimerData> timers) throws Exception {
         continue;
       }
 
+      // pass draining bit
       ReduceFn<K, InputT, OutputT, W>.Context directContext =
           contextFactory.base(window, StateStyle.DIRECT);
       ReduceFn<K, InputT, OutputT, W>.Context renamedContext =
           contextFactory.base(window, StateStyle.RENAMED);
-      WindowActivation windowActivation = new WindowActivation(directContext, renamedContext);
+      WindowActivation windowActivation =
+          new WindowActivation(directContext, renamedContext, timer.causedByDrain());
       windowActivations.put(window, windowActivation);
 
       // Perform prefetching of state to determine if the trigger should fire.
@@ -757,11 +773,12 @@ public void onTimers(Iterable<TimerData> timers) throws Exception {
 
       if (windowActivation.isGarbageCollection) {
         WindowTracing.debug(
-            "ReduceFnRunner: Cleaning up for key:{}; window:{} with inputWatermark:{}; outputWatermark:{}",
+            "ReduceFnRunner: Cleaning up for key:{}; window:{} with inputWatermark:{}; outputWatermark:{}; draining:{}",
             key,
             directContext.window(),
             timerInternals.currentInputWatermarkTime(),
-            timerInternals.currentOutputWatermarkTime());
+            timerInternals.currentOutputWatermarkTime(),
+            windowActivation.causedByDrain);
 
         boolean windowIsActiveAndOpen = windowActivation.windowIsActiveAndOpen();
         if (windowIsActiveAndOpen) {
@@ -774,7 +791,8 @@ public void onTimers(Iterable<TimerData> timers) throws Exception {
                   directContext,
                   renamedContext,
                   true /* isFinished */,
-                  windowActivation.isEndOfWindow);
+                  windowActivation.isEndOfWindow,
+                  windowActivation.causedByDrain);
           checkState(newHold == null, "Hold placed at %s despite isFinished being true.", newHold);
         }
 
@@ -792,7 +810,7 @@ public void onTimers(Iterable<TimerData> timers) throws Exception {
         if (windowActivation.windowIsActiveAndOpen()
             && triggerRunner.shouldFire(
                 directContext.window(), directContext.timers(), directContext.state())) {
-          emit(directContext, renamedContext);
+          emit(directContext, renamedContext, windowActivation.causedByDrain);
         }
 
         if (windowActivation.isEndOfWindow) {
@@ -915,7 +933,8 @@ private void prefetchEmit(
   /** Emit if a trigger is ready to fire or timers require it, and cleanup state. */
   private void emit(
       ReduceFn<K, InputT, OutputT, W>.Context directContext,
-      ReduceFn<K, InputT, OutputT, W>.Context renamedContext)
+      ReduceFn<K, InputT, OutputT, W>.Context renamedContext,
+      CausedByDrain causedByDrain)
       throws Exception {
     checkState(
         triggerRunner.shouldFire(
@@ -931,7 +950,7 @@ 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*/);
+    onTrigger(directContext, renamedContext, isFinished, false /*isEndOfWindow*/, causedByDrain);
 
     // Now that we've triggered, the pane is empty.
     nonEmptyPanes.clearPane(renamedContext.state());
@@ -989,7 +1008,8 @@ private void prefetchOnTrigger(
       final ReduceFn<K, InputT, OutputT, W>.Context directContext,
       ReduceFn<K, InputT, OutputT, W>.Context renamedContext,
       final boolean isFinished,
-      boolean isEndOfWindow)
+      boolean isEndOfWindow,
+      CausedByDrain causedByDrain)
       throws Exception {
     // Extract the window hold, and as a side effect clear it.
     final WatermarkHold.OldAndNewHolds pair =
@@ -1061,10 +1081,12 @@ private void prefetchOnTrigger(
                     .setValue(KV.of(key, toOutput))
                     .setTimestamp(outputTimestamp)
                     .setWindows(windows)
+                    .setCausedByDrain(causedByDrain)
                     .setPaneInfo(paneInfo)
                     .setReceiver(outputter)
                     .output();
-              });
+              },
+              causedByDrain);
 
       reduceFn.onTrigger(renamedTriggerContext);
     }

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

@@ -599,6 +599,7 @@ public String getErrorContext() {
           timerInternals.currentProcessingTime().plus(result.getContinuation().resumeDelay());
       holdState.add(futureOutputWatermark);
       // Set a timer to continue processing this element.
+      // todo radoslws@ decide if draining should be set on timer
       timerInternals.setTimer(
           TimerInternals.TimerData.of(
               stateNamespace,