find npe

Author: stankiewicz

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

@@ -126,7 +126,7 @@ public static <K> Timer<K> cleared(
    */
   public abstract @Nullable PaneInfo getPaneInfo();
 
-  public abstract @Nullable CausedByDrain causedByDrain();
+  public abstract CausedByDrain causedByDrain();
 
   @Override
   public final boolean equals(@Nullable Object other) {

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

@@ -148,6 +148,7 @@ public Builder<T> setRecordOffset(@Nullable Long recordOffset) {
 
     @Override
     public Builder<T> setCausedByDrain(CausedByDrain causedByDrain) {
+      checkStateNotNull(causedByDrain, "CausedByDrain is null");
       this.causedByDrain = causedByDrain;
       return this;
     }
@@ -202,6 +203,7 @@ public PaneInfo getPaneInfo() {
 
     @Override
     public CausedByDrain causedByDrain() {
+      checkStateNotNull(causedByDrain, "CausedByDrain not set");
       return causedByDrain;
     }
 

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

@@ -624,15 +624,18 @@ private void startBundle() {
 
   private void processElementForParDo(WindowedValue<InputT> elem) {
     currentElement = elem;
+    causedByDrain = currentElement.causedByDrain();
     try {
       doFnInvoker.invokeProcessElement(processContext);
     } finally {
       currentElement = null;
+      causedByDrain = null;
     }
   }
 
   private void processElementForWindowObservingParDo(WindowedValue<InputT> elem) {
     currentElement = elem;
+    causedByDrain = currentElement.causedByDrain();
     try {
       Iterator<BoundedWindow> windowIterator =
           (Iterator<BoundedWindow>) elem.getWindows().iterator();
@@ -643,12 +646,14 @@ private void processElementForWindowObservingParDo(WindowedValue<InputT> elem) {
     } finally {
       currentElement = null;
       currentWindow = null;
+      causedByDrain = null;
     }
   }
 
   private void processElementForWindowObservingSizedElementAndRestriction(
       WindowedValue<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> elem) {
     currentElement = elem.withValue(elem.getValue().getKey().getKey());
+    causedByDrain = elem.causedByDrain();
     windowCurrentIndex = -1;
     windowStopIndex = currentElement.getWindows().size();
     currentWindows = ImmutableList.copyOf(currentElement.getWindows());
@@ -660,6 +665,7 @@ private void processElementForWindowObservingSizedElementAndRestriction(
           windowCurrentIndex = -1;
           windowStopIndex = 0;
           currentElement = null;
+          causedByDrain = null;
           currentWindows = null;
           currentRestriction = null;
           currentWatermarkEstimatorState = null;
@@ -1202,7 +1208,8 @@ private <K> void processTimer(
     checkNotNull(timerBundleTracker);
     try {
       currentKey = timer.getUserKey();
-      causedByDrain = timer.causedByDrain();
+      causedByDrain = Preconditions.checkNotNull(timer.causedByDrain());
+
       // add drain
       Iterator<BoundedWindow> windowIterator =
           (Iterator<BoundedWindow>) timer.getWindows().iterator();
@@ -1286,6 +1293,7 @@ private <K> void processOnWindowExpiration(Timer<K> timer) {
     try {
       currentKey = timer.getUserKey();
       currentTimer = timer;
+      causedByDrain = timer.causedByDrain();
       Iterator<BoundedWindow> windowIterator =
           (Iterator<BoundedWindow>) timer.getWindows().iterator();
       while (windowIterator.hasNext()) {
@@ -1296,6 +1304,7 @@ private <K> void processOnWindowExpiration(Timer<K> timer) {
       currentKey = null;
       currentTimer = null;
       currentWindow = null;
+      causedByDrain = null;
     }
   }
 
@@ -2315,7 +2324,7 @@ public OutputBuilder<OutputT> builder(OutputT value) {
             .setWindow(currentWindow)
             .setTimestamp(currentTimer.getHoldTimestamp())
             .setPaneInfo(currentTimer.getPaneInfo())
-            .setCausedByDrain(causedByDrain)
+            .setCausedByDrain(currentTimer.causedByDrain())
             .setReceiver(
                 windowedValue -> {
                   checkOnWindowExpirationTimestamp(windowedValue.getTimestamp());
@@ -2336,7 +2345,7 @@ public <T> void output(TupleTag<T> tag, T output) {
                 output,
                 currentTimer.getHoldTimestamp(),
                 currentWindow,
-                currentTimer.getPaneInfo()));
+                currentTimer.getPaneInfo(), null, null, currentTimer.causedByDrain()));
       }
 
       @Override
@@ -2786,7 +2795,7 @@ public OutputBuilder<T> builder(T value) {
                     .setValue(value)
                     .setTimestamp(currentTimer.getHoldTimestamp())
                     .setWindow(currentWindow)
-                    .setCausedByDrain(causedByDrain)
+                    .setCausedByDrain(currentTimer.causedByDrain())
                     .setPaneInfo(currentTimer.getPaneInfo())
                     .setReceiver(windowedValue -> context.outputWindowedValue(windowedValue));
               }
@@ -2819,7 +2828,7 @@ public OutputBuilder<Row> builder(Row value) {
                     .setTimestamp(currentTimer.getHoldTimestamp())
                     .setWindow(currentWindow)
                     .setPaneInfo(currentTimer.getPaneInfo())
-                    .setCausedByDrain(causedByDrain)
+                    .setCausedByDrain(currentTimer.causedByDrain())
                     .setReceiver(
                         windowedValue ->
                             context.outputWindowedValue(