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();