Drain Mode as WindowedValue extension

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto

@@ -1677,6 +1677,9 @@ message StandardProtocols {
     // Indicates whether the SDK supports ordered list state.
     ORDERED_LIST_STATE = 10
     [(beam_urn) = "beam:protocol:ordered_list_state:v1"];
+
+    EXTENDED_ELEMENT_METADATA = 11
+    [(beam_urn) = "beam:protocol:extended_element_metadata:v1"];
   }
 }
 
@@ -2028,3 +2031,14 @@ message StandardResourceHints {
     CPU_COUNT = 2 [(beam_urn) = "beam:resources:cpu_count:v1"];
   }
 }
+message DrainMode {
+  enum Enum {
+    DRAIN_MODE_UNSPECIFIED = 0;
+    DRAIN_MODE_NOT_DRAINING = 1;
+    DRAIN_MODE_DRAINING = 2;
+  }
+}
+
+message ElementMetadata {
+  optional DrainMode.Enum drain = 1;
+}

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

@@ -63,6 +63,7 @@
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.transforms.windowing.Window;
 import org.apache.beam.sdk.util.CoderUtils;
+import org.apache.beam.sdk.util.ElementMetadata;
 import org.apache.beam.sdk.util.SystemDoFnInternal;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
@@ -1364,16 +1365,28 @@ private static <T> WindowedValue<T> valueInEmptyWindows(T value) {
   private static class ValueInEmptyWindows<T> extends WindowedValue<T> {
 
     private final T value;
+    private final @Nullable ElementMetadata elementMetadata;
 
     private ValueInEmptyWindows(T value) {
       this.value = value;
+      this.elementMetadata = null;
+    }
+
+    private ValueInEmptyWindows(T value, @Nullable ElementMetadata elementMetadata) {
+      this.value = value;
+      this.elementMetadata = elementMetadata;
     }
 
     @Override
     public <NewT> WindowedValue<NewT> withValue(NewT value) {
       return new ValueInEmptyWindows<>(value);
     }
 
+    @Override
+    public WindowedValue<T> withElementMetadata(@Nullable ElementMetadata elementMetadata) {
+      return new ValueInEmptyWindows<>(this.getValue(), elementMetadata);
+    }
+
     @Override
     public T getValue() {
       return value;
@@ -1394,6 +1407,11 @@ public PaneInfo getPane() {
       return PaneInfo.NO_FIRING;
     }
 
+    @Override
+    public @Nullable ElementMetadata getElementMetadata() {
+      return elementMetadata;
+    }
+
     @Override
     public String toString() {
       return MoreObjects.toStringHelper(getClass()).add("value", getValue()).toString();

sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java

@@ -146,10 +146,10 @@ private static byte encodedByte(boolean isFirst, boolean isLast, Timing timing)
     ImmutableMap.Builder<Byte, PaneInfo> decodingBuilder = ImmutableMap.builder();
     for (Timing timing : Timing.values()) {
       long onTimeIndex = timing == Timing.EARLY ? -1 : 0;
-      register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex));
-      register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex));
-      register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex));
-      register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex));
+      register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex, false));
+      register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex, false));
+      register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex, false));
+      register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex, false));
     }
     BYTE_TO_PANE_INFO = decodingBuilder.build();
   }
@@ -159,7 +159,7 @@ private static void register(ImmutableMap.Builder<Byte, PaneInfo> builder, PaneI
   }
 
   private final byte encodedByte;
-
+  private final boolean containsElementMetadata;
   private final boolean isFirst;
   private final boolean isLast;
   private final Timing timing;
@@ -177,13 +177,20 @@ private static void register(ImmutableMap.Builder<Byte, PaneInfo> builder, PaneI
   public static final PaneInfo ON_TIME_AND_ONLY_FIRING =
       PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0);
 
-  private PaneInfo(boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) {
+  private PaneInfo(
+      boolean isFirst,
+      boolean isLast,
+      Timing timing,
+      long index,
+      long onTimeIndex,
+      boolean containsElementMetadata) {
     this.encodedByte = encodedByte(isFirst, isLast, timing);
     this.isFirst = isFirst;
     this.isLast = isLast;
     this.timing = timing;
     this.index = index;
     this.nonSpeculativeIndex = onTimeIndex;
+    this.containsElementMetadata = containsElementMetadata;
   }
 
   public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing) {
@@ -194,10 +201,21 @@ public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing
   /** Factory method to create a {@link PaneInfo} with the specified parameters. */
   public static PaneInfo createPane(
       boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) {
+    return createPane(isFirst, isLast, timing, index, onTimeIndex, false);
+  }
+
+  /** Factory method to create a {@link PaneInfo} with the specified parameters. */
+  public static PaneInfo createPane(
+      boolean isFirst,
+      boolean isLast,
+      Timing timing,
+      long index,
+      long onTimeIndex,
+      boolean containsElementMetadata) {
     if (isFirst || timing == Timing.UNKNOWN) {
       return checkNotNull(BYTE_TO_PANE_INFO.get(encodedByte(isFirst, isLast, timing)));
     } else {
-      return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex);
+      return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex, containsElementMetadata);
     }
   }
 
@@ -241,6 +259,15 @@ public long getIndex() {
     return index;
   }
 
+  public boolean isElementMetadata() {
+    return containsElementMetadata;
+  }
+
+  public PaneInfo withElementMetadata(boolean elementMetadata) {
+    return new PaneInfo(
+        this.isFirst, this.isLast, this.timing, index, nonSpeculativeIndex, elementMetadata);
+  }
+
   /**
    * The zero-based index of this trigger firing among non-speculative panes.
    *
@@ -360,7 +387,8 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce
       byte keyAndTag = (byte) inStream.read();
       PaneInfo base = Preconditions.checkNotNull(BYTE_TO_PANE_INFO.get((byte) (keyAndTag & 0x0F)));
       long index, onTimeIndex;
-      switch (Encoding.fromTag(keyAndTag)) {
+      boolean elementMetadata = (keyAndTag & 0x80) > 0;
+      switch (Encoding.fromTag((byte) (keyAndTag & 0x7F))) {
         case FIRST:
           return base;
         case ONE_INDEX:
@@ -374,7 +402,8 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce
         default:
           throw new CoderException("Unknown encoding " + (keyAndTag & 0xF0));
       }
-      return new PaneInfo(base.isFirst, base.isLast, base.timing, index, onTimeIndex);
+      return new PaneInfo(
+          base.isFirst, base.isLast, base.timing, index, onTimeIndex, elementMetadata);
     }
 
     @Override

sdks/java/core/src/main/java/org/apache/beam/sdk/util/ElementMetadata.java

@@ -0,0 +1,34 @@
+/*
+ * 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.util;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.model.pipeline.v1.RunnerApi.DrainMode;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+@AutoValue
+public abstract class ElementMetadata {
+
+  DrainMode.@Nullable Enum drainMode;
+
+  public static ElementMetadata create(DrainMode.@Nullable Enum drainMode) {
+    return new AutoValue_ElementMetadata(drainMode);
+  }
+
+  abstract DrainMode.@Nullable Enum drainMode();
+}