Merge pull request #36524: Drain - model part + windowedValue changes

Author: kennknowles

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/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/BatchViewOverrides.java

@@ -1378,6 +1378,11 @@ public T getValue() {
       return value;
     }
 
+    @Override
+    public boolean causedByDrain() {
+      return false;
+    }
+
     @Override
     public Instant getTimestamp() {
       return BoundedWindow.TIMESTAMP_MIN_VALUE;

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

@@ -59,6 +59,11 @@ public PaneInfo getPaneInfo() {
     return null;
   }
 
+  @Override
+  public boolean causedByDrain() {
+    return false;
+  }
+
   @Override
   public Iterable<WindowedValue<T>> explodeWindows() {
     return Collections.emptyList();

runners/spark/src/main/java/org/apache/beam/runners/spark/util/TimerUtils.java

@@ -115,6 +115,11 @@ public PaneInfo getPaneInfo() {
       return null;
     }
 
+    @Override
+    public boolean causedByDrain() {
+      return false;
+    }
+
     @Override
     public @Nullable Long getRecordOffset() {
       return null;

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

@@ -48,5 +48,7 @@ public interface OutputBuilder<T> extends WindowedValue<T> {
 
   OutputBuilder<T> setRecordOffset(@Nullable Long recordOffset);
 
+  OutputBuilder<T> setCausedByDrain(boolean causedByDrain);
+
   void output();
 }

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

@@ -52,6 +52,8 @@ public interface WindowedValue<T> {
   @Nullable
   Long getRecordOffset();
 
+  boolean causedByDrain();
+
   /**
    * A representation of each of the actual values represented by this compressed {@link
    * WindowedValue}, one per window.