[WIP] Drain in element metadata v1

model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto

@@ -740,10 +740,18 @@ message Elements {
     bool is_last = 4;
   }
 
+  message DrainMode {
+    enum Enum {
+      UNSPECIFIED = 0;
+      NOT_DRAINING = 1;
+      DRAINING = 2;
+    }
+  }
+
   // Element metadata passed as part of WindowedValue to make WindowedValue
   // extensible and backward compatible
   message ElementMetadata {
-    // empty message - add drain, kind, tracing metadata in the future
+    optional DrainMode.Enum drain = 1;
   }
 
   // Represent the encoded user timer for a given instruction, transform and

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

@@ -46,7 +46,8 @@ public interface DoFnRunner<InputT extends @Nullable Object, OutputT extends @Nu
       BoundedWindow window,
       Instant timestamp,
       Instant outputTimestamp,
-      TimeDomain timeDomain);
+      TimeDomain timeDomain,
+      @Nullable Boolean draining);
 
   /**
    * Calls a {@link DoFn DoFn's} {@link DoFn.FinishBundle @FinishBundle} method and performs

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

@@ -30,6 +30,7 @@
 import org.apache.beam.sdk.values.WindowedValues;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Instant;
 
 /**
@@ -89,8 +90,10 @@ public <KeyT> void onTimer(
       BoundedWindow window,
       Instant timestamp,
       Instant outputTimestamp,
-      TimeDomain timeDomain) {
-    doFnRunner.onTimer(timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain);
+      TimeDomain timeDomain,
+      @Nullable Boolean draining) {
+    doFnRunner.onTimer(
+        timerId, timerFamilyId, key, window, timestamp, outputTimestamp, timeDomain, draining);
   }
 
   @Override

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

@@ -386,6 +386,11 @@ public PaneInfo pane() {
       return element.getPaneInfo();
     }
 
+    @Override
+    public Boolean draining() {
+      return element.isDraining();
+    }
+
     @Override
     public String currentRecordId() {
       return element.getRecordId();
@@ -439,7 +444,13 @@ public void outputWindowedValue(
       outputReceiver.output(
           mainOutputTag,
           WindowedValues.of(
-              value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset));
+              value,
+              timestamp,
+              windows,
+              paneInfo,
+              currentRecordId,
+              currentRecordOffset,
+              element.isDraining()));
     }
 
     @Override
@@ -450,7 +461,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.isDraining()));
     }
 
     @Override
@@ -479,7 +498,13 @@ public <T> void outputWindowedValue(
       outputReceiver.output(
           tag,
           WindowedValues.of(
-              value, timestamp, windows, paneInfo, currentRecordId, currentRecordOffset));
+              value,
+              timestamp,
+              windows,
+              paneInfo,
+              currentRecordId,
+              currentRecordOffset,
+              element.isDraining()));
     }
 
     private void noteOutput() {

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

@@ -59,6 +59,7 @@ public abstract class ProcessValueContext extends Context {
     /** Return the actual value being processed. */
     public abstract InputT value();
 
+    public abstract Boolean draining();
     /** Return the timestamp associated with the value. */
     public abstract Instant timestamp();
   }
@@ -75,6 +76,7 @@ public abstract class OnTriggerContext extends Context {
     /** Returns the {@link PaneInfo} for the trigger firing being processed. */
     public abstract PaneInfo paneInfo();
 
+    public abstract Boolean draining();
     /** 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

@@ -94,17 +94,26 @@ private StateAccessorImpl<K, W> stateAccessor(W window, StateStyle style) {
   }
 
   public ReduceFn<K, InputT, OutputT, W>.Context base(W window, StateStyle style) {
-    return new ContextImpl(stateAccessor(window, style));
+    return new ContextImpl(stateAccessor(window, style), null);
+  }
+
+  public ReduceFn<K, InputT, OutputT, W>.Context base(
+      W window, StateStyle style, Boolean draining) {
+    return new ContextImpl(stateAccessor(window, style), draining);
   }
 
   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, Boolean draining) {
+    return new ProcessValueContextImpl(stateAccessor(window, style), value, timestamp, draining);
   }
 
   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,
+      Boolean draining) {
+    return new OnTriggerContextImpl(stateAccessor(window, style), paneInfo, callbacks, draining);
   }
 
   public ReduceFn<K, InputT, OutputT, W>.OnMergeContext forMerge(
@@ -135,16 +144,19 @@ public TimersImpl(StateNamespace namespace) {
 
     @Override
     public void setTimer(Instant timestamp, TimeDomain timeDomain) {
+      // todo radoslaws@ should we pass draining bit to timerdata
       timerInternals.setTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain));
     }
 
     @Override
     public void setTimer(Instant timestamp, Instant outputTimestamp, TimeDomain timeDomain) {
+      // todo radoslaws@ should we pass draining bit to timerdata
       timerInternals.setTimer(TimerData.of(namespace, timestamp, outputTimestamp, timeDomain));
     }
 
     @Override
     public void deleteTimer(Instant timestamp, TimeDomain timeDomain) {
+      // todo radoslaws@ should we pass draining bit to timerdata
       timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timestamp, timeDomain));
     }
 
@@ -315,11 +327,13 @@ public <StateT extends State> Map<W, StateT> accessInEachMergingWindow(
   private class ContextImpl extends ReduceFn<K, InputT, OutputT, W>.Context {
     private final StateAccessorImpl<K, W> state;
     private final TimersImpl timers;
+    private final Boolean draining;
 
-    private ContextImpl(StateAccessorImpl<K, W> state) {
+    private ContextImpl(StateAccessorImpl<K, W> state, @Nullable Boolean draining) {
       reduceFn.super();
       this.state = state;
       this.timers = new TimersImpl(state.namespace());
+      this.draining = draining;
     }
 
     @Override
@@ -354,21 +368,28 @@ private class ProcessValueContextImpl
     private final Instant timestamp;
     private final StateAccessorImpl<K, W> state;
     private final TimersImpl timers;
+    private final Boolean draining;
 
     private ProcessValueContextImpl(
-        StateAccessorImpl<K, W> state, InputT value, Instant timestamp) {
+        StateAccessorImpl<K, W> state, InputT value, Instant timestamp, Boolean draining) {
       reduceFn.super();
       this.state = state;
       this.value = value;
       this.timestamp = timestamp;
       this.timers = new TimersImpl(state.namespace());
+      this.draining = draining;
     }
 
     @Override
     public K key() {
       return key;
     }
 
+    @Override
+    public Boolean draining() {
+      return draining;
+    }
+
     @Override
     public W window() {
       return state.window();
@@ -405,21 +426,31 @@ 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 Boolean draining;
 
     private OnTriggerContextImpl(
-        StateAccessorImpl<K, W> state, PaneInfo paneInfo, OnTriggerCallbacks<OutputT> callbacks) {
+        StateAccessorImpl<K, W> state,
+        PaneInfo paneInfo,
+        OnTriggerCallbacks<OutputT> callbacks,
+        Boolean draining) {
       reduceFn.super();
       this.state = state;
       this.paneInfo = paneInfo;
       this.callbacks = callbacks;
       this.timers = new TimersImpl(state.namespace());
+      this.draining = draining;
     }
 
     @Override
     public K key() {
       return key;
     }
 
+    @Override
+    public Boolean draining() {
+      return draining;
+    }
+
     @Override
     public W window() {
       return state.window();