Fix nullness errors in SimpleDoFnRunner and DoFnInvoker

Author: kennknowles

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

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.core;
 
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 
@@ -66,6 +67,9 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.FluentIterable;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
+import org.checkerframework.checker.initialization.qual.Initialized;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
@@ -81,11 +85,6 @@
  * @param <InputT> the type of the {@link DoFn} (main) input elements
  * @param <OutputT> the type of the {@link DoFn} (main) output elements
  */
-@SuppressWarnings({
-  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
-  "nullness",
-  "keyfor"
-}) // TODO(https://github.com/apache/beam/issues/20497)
 public class SimpleDoFnRunner<InputT, OutputT> implements DoFnRunner<InputT, OutputT> {
 
   private final PipelineOptions options;
@@ -527,13 +526,21 @@ public Object key() {
     }
 
     @Override
-    public Object sideInput(String tagId) {
-      return sideInput(sideInputMapping.get(tagId));
+    public @Nullable Object sideInput(String tagId) {
+      PCollectionView<?> view =
+          checkStateNotNull(sideInputMapping.get(tagId), "Side input tag %s not found", tagId);
+      return sideInput(view);
     }
 
     @Override
     public Object schemaElement(int index) {
-      SerializableFunction converter = doFnSchemaInformation.getElementConverters().get(index);
+      checkStateNotNull(
+          doFnSchemaInformation,
+          "attempt to access element via schema when no schema information provided");
+
+      SerializableFunction<InputT, Object> converter =
+          (SerializableFunction<InputT, Object>)
+              doFnSchemaInformation.getElementConverters().get(index);
       return converter.apply(element());
     }
 
@@ -561,6 +568,7 @@ public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {
 
     @Override
     public OutputReceiver<Row> outputRowReceiver(DoFn<InputT, OutputT> doFn) {
+      checkStateNotNull(mainOutputSchemaCoder, "cannot provide row receiver without schema coder");
       return DoFnOutputReceivers.rowReceiver(
           this, builderSupplier, mainOutputTag, mainOutputSchemaCoder);
     }
@@ -601,14 +609,25 @@ public WatermarkEstimator<?> watermarkEstimator() {
     @Override
     public State state(String stateId, boolean alwaysFetched) {
       try {
+        DoFnSignature.StateDeclaration stateDeclaration =
+            checkStateNotNull(
+                signature.stateDeclarations().get(stateId), "state not found: %s", stateId);
+
         StateSpec<?> spec =
-            (StateSpec<?>) signature.stateDeclarations().get(stateId).field().get(fn);
+            checkStateNotNull(
+                (StateSpec<?>) stateDeclaration.field().get(fn),
+                "Field %s corresponding to state id %s contained null",
+                stateDeclaration.field(),
+                stateId);
+
+        @NonNull
+        @Initialized // unclear why checkerframework needs this help
         State state =
             stepContext
                 .stateInternals()
-                .state(getNamespace(), StateTags.tagForSpec(stateId, (StateSpec) spec));
+                .state(getNamespace(), StateTags.tagForSpec(stateId, (StateSpec<?>) spec));
         if (alwaysFetched) {
-          return (State) ((ReadableState) state).readLater();
+          return (State) ((ReadableState<?>) state).readLater();
         } else {
           return state;
         }
@@ -620,7 +639,16 @@ public State state(String stateId, boolean alwaysFetched) {
     @Override
     public Timer timer(String timerId) {
       try {
-        TimerSpec spec = (TimerSpec) signature.timerDeclarations().get(timerId).field().get(fn);
+        DoFnSignature.TimerDeclaration timerDeclaration =
+            checkStateNotNull(
+                signature.timerDeclarations().get(timerId), "timer not found: %s", timerId);
+        TimerSpec spec =
+            (TimerSpec)
+                checkStateNotNull(
+                    timerDeclaration.field().get(fn),
+                    "Field %s corresponding to timer id %s contained null",
+                    timerDeclaration.field(),
+                    timerId);
         return new TimerInternalsTimer(
             window(), getNamespace(), timerId, spec, timestamp(), stepContext.timerInternals());
       } catch (IllegalAccessException e) {
@@ -631,8 +659,19 @@ public Timer timer(String timerId) {
     @Override
     public TimerMap timerFamily(String timerFamilyId) {
       try {
+        DoFnSignature.TimerFamilyDeclaration timerFamilyDeclaration =
+            checkStateNotNull(
+                signature.timerFamilyDeclarations().get(timerFamilyId),
+                "timer family not found: %s",
+                timerFamilyId);
+
         TimerSpec spec =
-            (TimerSpec) signature.timerFamilyDeclarations().get(timerFamilyId).field().get(fn);
+            (TimerSpec)
+                checkStateNotNull(
+                    timerFamilyDeclaration.field().get(fn),
+                    "Field %s corresponding to timer family id %s contained null",
+                    timerFamilyDeclaration.field(),
+                    timerFamilyId);
         return new TimerInternalsTimerMap(
             timerFamilyId,
             window(),
@@ -794,6 +833,7 @@ public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {
 
     @Override
     public OutputReceiver<Row> outputRowReceiver(DoFn<InputT, OutputT> doFn) {
+      checkStateNotNull(mainOutputSchemaCoder, "cannot provide row receiver without schema coder");
       return DoFnOutputReceivers.rowReceiver(
           this, builderSupplier, mainOutputTag, mainOutputSchemaCoder);
     }
@@ -833,8 +873,18 @@ public WatermarkEstimator<?> watermarkEstimator() {
     @Override
     public State state(String stateId, boolean alwaysFetched) {
       try {
+        DoFnSignature.StateDeclaration stateDeclaration =
+            checkStateNotNull(
+                signature.stateDeclarations().get(stateId), "state not found: %s", stateId);
+
         StateSpec<?> spec =
-            (StateSpec<?>) signature.stateDeclarations().get(stateId).field().get(fn);
+            checkStateNotNull(
+                (StateSpec<?>) stateDeclaration.field().get(fn),
+                "Field %s corresponding to state id %s contained null",
+                stateDeclaration.field(),
+                stateId);
+
+        @NonNull
         State state =
             stepContext
                 .stateInternals()
@@ -852,7 +902,16 @@ public State state(String stateId, boolean alwaysFetched) {
     @Override
     public Timer timer(String timerId) {
       try {
-        TimerSpec spec = (TimerSpec) signature.timerDeclarations().get(timerId).field().get(fn);
+        DoFnSignature.TimerDeclaration timerDeclaration =
+            checkStateNotNull(
+                signature.timerDeclarations().get(timerId), "timer not found: %s", timerId);
+        TimerSpec spec =
+            (TimerSpec)
+                checkStateNotNull(
+                    timerDeclaration.field().get(fn),
+                    "Field %s corresponding to timer id %s contained null",
+                    timerDeclaration.field(),
+                    timerId);
         return new TimerInternalsTimer(
             window, getNamespace(), timerId, spec, timestamp(), stepContext.timerInternals());
       } catch (IllegalAccessException e) {
@@ -863,8 +922,18 @@ public Timer timer(String timerId) {
     @Override
     public TimerMap timerFamily(String timerFamilyId) {
       try {
+        DoFnSignature.TimerFamilyDeclaration timerFamilyDeclaration =
+            checkStateNotNull(
+                signature.timerFamilyDeclarations().get(timerFamilyId),
+                "timer family not found: %s",
+                timerFamilyId);
         TimerSpec spec =
-            (TimerSpec) signature.timerFamilyDeclarations().get(timerFamilyId).field().get(fn);
+            (TimerSpec)
+                checkStateNotNull(
+                    timerFamilyDeclaration.field().get(fn),
+                    "Field %s corresponding to timer family id %s contained null",
+                    timerFamilyDeclaration.field(),
+                    timerFamilyId);
         return new TimerInternalsTimerMap(
             timerFamilyId,
             window(),
@@ -1058,6 +1127,7 @@ public OutputReceiver<OutputT> outputReceiver(DoFn<InputT, OutputT> doFn) {
 
     @Override
     public OutputReceiver<Row> outputRowReceiver(DoFn<InputT, OutputT> doFn) {
+      checkStateNotNull(mainOutputSchemaCoder, "cannot provide row receiver without schema coder");
       return DoFnOutputReceivers.rowReceiver(
           this, builderSupplier, mainOutputTag, mainOutputSchemaCoder);
     }
@@ -1096,14 +1166,23 @@ public WatermarkEstimator<?> watermarkEstimator() {
     @Override
     public State state(String stateId, boolean alwaysFetched) {
       try {
+        DoFnSignature.StateDeclaration stateDeclaration =
+            checkStateNotNull(
+                signature.stateDeclarations().get(stateId), "state not found: %s", stateId);
         StateSpec<?> spec =
-            (StateSpec<?>) signature.stateDeclarations().get(stateId).field().get(fn);
+            checkStateNotNull(
+                (StateSpec<?>) stateDeclaration.field().get(fn),
+                "Field %s corresponding to state id %s contained null",
+                stateDeclaration.field(),
+                stateId);
+        @NonNull
+        @Initialized // unclear why checkerframework needs this help
         State state =
             stepContext
                 .stateInternals()
-                .state(getNamespace(), StateTags.tagForSpec(stateId, (StateSpec) spec));
+                .state(getNamespace(), StateTags.tagForSpec(stateId, (StateSpec<?>) spec));
         if (alwaysFetched) {
-          return (State) ((ReadableState) state).readLater();
+          return (State) ((ReadableState<?>) state).readLater();
         } else {
           return state;
         }
@@ -1195,7 +1274,7 @@ private class TimerInternalsTimer implements Timer {
     private final String timerId;
     private final String timerFamilyId;
     private final TimerSpec spec;
-    private Instant target;
+    private @MonotonicNonNull Instant target;
     private @Nullable Instant outputTimestamp;
     private boolean noOutputTimestamp;
     private final Instant elementInputTimestamp;
@@ -1313,23 +1392,26 @@ public Timer withNoOutputTimestamp() {
      * <li>The current element timestamp for other time domains.
      */
     private void setAndVerifyOutputTimestamp() {
+      checkStateNotNull(target, "attempt to set outputTimestamp before setting target firing time");
       if (outputTimestamp != null) {
+        // setting to local var so checkerframework knows that method calls will not mutate it
+        Instant timestampToValidate = outputTimestamp;
         Instant lowerBound;
         try {
           lowerBound = elementInputTimestamp.minus(fn.getAllowedTimestampSkew());
         } catch (ArithmeticException e) {
           lowerBound = BoundedWindow.TIMESTAMP_MIN_VALUE;
         }
-        if (outputTimestamp.isBefore(lowerBound)
-            || outputTimestamp.isAfter(BoundedWindow.TIMESTAMP_MAX_VALUE)) {
+        if (timestampToValidate.isBefore(lowerBound)
+            || timestampToValidate.isAfter(BoundedWindow.TIMESTAMP_MAX_VALUE)) {
           throw new IllegalArgumentException(
               String.format(
                   "Cannot output timer with output timestamp %s. Output timestamps must be no "
                       + "earlier than the timestamp of the current input or timer (%s) minus the "
                       + "allowed skew (%s) and no later than %s. See the "
                       + "DoFn#getAllowedTimestampSkew() Javadoc for details on changing the "
                       + "allowed skew.",
-                  outputTimestamp,
+                  timestampToValidate,
                   elementInputTimestamp,
                   fn.getAllowedTimestampSkew().getMillis() >= Integer.MAX_VALUE
                       ? fn.getAllowedTimestampSkew()
@@ -1346,6 +1428,9 @@ private void setAndVerifyOutputTimestamp() {
         // the element (or timer) setting this timer.
         outputTimestamp = elementInputTimestamp;
       }
+
+      // Now it has been set for all cases other than this.noOutputTimestamp == true, and there are
+      // further validations
       if (outputTimestamp != null) {
         Instant windowExpiry = LateDataUtils.garbageCollectionTime(window, allowedLateness);
         if (TimeDomain.EVENT_TIME.equals(spec.getTimeDomain())) {
@@ -1380,6 +1465,12 @@ private void setAndVerifyOutputTimestamp() {
      * user has no way to compute a good choice of time.
      */
     private void setUnderlyingTimer() {
+      checkStateNotNull(
+          outputTimestamp,
+          "internal error: null outputTimestamp: must be populated by setAndVerifyOutputTimestamp()");
+      checkStateNotNull(
+          target,
+          "internal error: attempt to set internal timer when target timestamp not yet set");
       timerInternals.setTimer(
           namespace, timerId, timerFamilyId, target, outputTimestamp, spec.getTimeDomain());
     }
@@ -1396,7 +1487,9 @@ private Instant getCurrentTime(TimeDomain timeDomain) {
         case PROCESSING_TIME:
           return timerInternals.currentProcessingTime();
         case SYNCHRONIZED_PROCESSING_TIME:
-          return timerInternals.currentSynchronizedProcessingTime();
+          return checkStateNotNull(
+              timerInternals.currentSynchronizedProcessingTime(),
+              "internal error: requested SYNCHRONIZED_PROCESSING_TIME but it was null");
         default:
           throw new IllegalStateException(
               String.format("Timer created for unknown time domain %s", spec.getTimeDomain()));
@@ -1446,19 +1539,17 @@ public void set(String timerId, Instant absoluteTime) {
 
     @Override
     public Timer get(String timerId) {
-      if (timers.get(timerId) == null) {
-        Timer timer =
-            new TimerInternalsTimer(
-                window,
-                namespace,
-                timerId,
-                timerFamilyId,
-                spec,
-                elementInputTimestamp,
-                timerInternals);
-        timers.put(timerId, timer);
-      }
-      return timers.get(timerId);
+      return timers.computeIfAbsent(
+          timerId,
+          id ->
+              new TimerInternalsTimer(
+                  window,
+                  namespace,
+                  id,
+                  timerFamilyId,
+                  spec,
+                  elementInputTimestamp,
+                  timerInternals));
     }
   }
 }