fix npe in ReduceFnRunner, add additional tests, add timestamp validation

Author: stankiewicz

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

@@ -376,7 +376,7 @@ public void processElements(Iterable<WindowedValue<InputT>> values) throws Excep
       emit(
           contextFactory.base(window, StateStyle.DIRECT),
           contextFactory.base(window, StateStyle.RENAMED),
-          null);
+          CausedByDrain.NORMAL);
     }
 
     // We're all done with merging and emitting elements so can compress the activeWindow state.

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

@@ -445,11 +445,13 @@ public <T> void output(TupleTag<T> tag, T output) {
 
     @Override
     public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+      checkTimestamp(elem.getTimestamp(), windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(mainOutputTag, windowedValue);
     }
 
     @Override
     public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
+      checkTimestamp(elem.getTimestamp(), windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(tag, windowedValue);
     }
 
@@ -1308,11 +1310,13 @@ public <T> void outputWindowedValue(
 
     @Override
     public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+      checkTimestamp(this.timestamp, windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(mainOutputTag, windowedValue);
     }
 
     @Override
     public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
+      checkTimestamp(this.timestamp, windowedValue.getTimestamp());
       SimpleDoFnRunner.this.outputWindowedValue(tag, windowedValue);
     }
 

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

@@ -322,6 +322,7 @@ public static <T> WindowedValue<T> of(
       @Nullable Long currentRecordOffset,
       CausedByDrain causedByDrain) {
     checkArgument(paneInfo != null, "WindowedValue requires PaneInfo, but it was null");
+    checkArgument(causedByDrain != null, "WindowedValue requires causedByDrain, but it was null");
 
     boolean isGlobal = GlobalWindow.INSTANCE.equals(window);
     if (isGlobal && BoundedWindow.TIMESTAMP_MIN_VALUE.equals(timestamp)) {

sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/MetadataPropagationTest.java

@@ -23,6 +23,7 @@
 import org.apache.beam.sdk.values.CausedByDrain;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.WindowedValues;
+import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
@@ -70,6 +71,22 @@ public void testMetadataPropagationAcrossShuffleParameter() {
       pipeline.run();
     }
 
+    @Test
+    @Category(NeedsRunner.class)
+    public void testDefaultMetadataPropagationAcrossShuffleParameter() {
+      Assert.assertFalse(WindowedValues.WindowedValueCoder.isMetadataSupported());
+      PCollection<String> results =
+          pipeline
+              .apply(Create.of(1))
+              .apply(ParDo.of(new CausedByDrainSettingDoFn()))
+              .apply(Redistribute.arbitrarily())
+              .apply(ParDo.of(new CausedByDrainExtractingDoFn()));
+
+      PAssert.that(results).containsInAnyOrder("NORMAL");
+
+      pipeline.run();
+    }
+
     @Test
     @Category(NeedsRunner.class)
     public void testMetadataPropagationParameter() {

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

@@ -1954,11 +1954,13 @@ public <T> void outputWindowedValue(
 
     @Override
     public void outputWindowedValue(WindowedValue<OutputT> windowedValue) {
+      checkTimestamp(windowedValue.getTimestamp());
       outputTo(mainOutputConsumer, windowedValue);
     }
 
     @Override
     public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T> windowedValue) {
+      checkTimestamp(windowedValue.getTimestamp());
       outputTo((FnDataReceiver) localNameToConsumer.get(tag.getId()), windowedValue);
     }