Fix propagation of metadata

Author: stankiewicz

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

@@ -145,7 +145,13 @@ public <K, InputT> Iterable<WindowedValue<InputT>> filter(
           } else {
             nonLateElements.add(
                 WindowedValues.of(
-                    element.getValue(), element.getTimestamp(), window, element.getPaneInfo()));
+                    element.getValue(),
+                    element.getTimestamp(),
+                    window,
+                    element.getPaneInfo(),
+                    element.getRecordId(),
+                    element.getRecordOffset(),
+                    element.causedByDrain()));
           }
         }
       }

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

@@ -39,6 +39,7 @@
 import org.apache.beam.sdk.transforms.windowing.PaneInfo;
 import org.apache.beam.sdk.transforms.windowing.PaneInfo.PaneInfoCoder;
 import org.apache.beam.sdk.util.ByteStringOutputStream;
+import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.ValueWithRecordId;
 import org.apache.beam.sdk.values.ValueWithRecordId.ValueWithRecordIdCoder;
@@ -220,7 +221,12 @@ public long add(WindowedValue<T> data) throws IOException {
       ByteString id = ByteString.EMPTY;
       // todo #33176 specify additional metadata in the future
       BeamFnApi.Elements.ElementMetadata additionalMetadata =
-          BeamFnApi.Elements.ElementMetadata.newBuilder().build();
+          BeamFnApi.Elements.ElementMetadata.newBuilder()
+              .setDrain(
+                  data.causedByDrain() == CausedByDrain.CAUSED_BY_DRAIN
+                      ? BeamFnApi.Elements.DrainMode.Enum.DRAINING
+                      : BeamFnApi.Elements.DrainMode.Enum.NOT_DRAINING)
+              .build();
       ByteString metadata =
           encodeMetadata(
               stream, windowsCoder, data.getWindows(), data.getPaneInfo(), additionalMetadata);