fix NPE and other exceptions thrown when work fails to process before work execution

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

@@ -73,7 +73,6 @@
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier;
 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.HashBasedTable;
@@ -129,7 +128,7 @@ public class StreamingModeExecutionContext extends DataflowExecutionContext<Step
    */
   private @Nullable Object key = null;
 
-  private Work work;
+  private @Nullable Work work;
   private WindmillComputationKey computationKey;
   private SideInputStateFetcher sideInputStateFetcher;
   // OperationalLimits is updated in start() because a StreamingModeExecutionContext can
@@ -304,9 +303,9 @@ private <T> SideInput<T> fetchSideInput(
     return fetchSideInputFromWindmill(
         view,
         sideInputWindow,
-        Preconditions.checkNotNull(stateFamily),
+        checkNotNull(stateFamily),
         state,
-        Preconditions.checkNotNull(scopedReadStateSupplier),
+        checkNotNull(scopedReadStateSupplier),
         tagCache);
   }
 
@@ -329,27 +328,27 @@ private <T> SideInput<T> fetchSideInputFromWindmill(
   }
 
   public Iterable<Windmill.GlobalDataId> getSideInputNotifications() {
-    return work.getWorkItem().getGlobalDataIdNotificationsList();
+    return getWorkItem().getGlobalDataIdNotificationsList();
   }
 
   private List<Timer> getFiredTimers() {
-    return work.getWorkItem().getTimers().getTimersList();
+    return getWorkItem().getTimers().getTimersList();
   }
 
   public @Nullable ByteString getSerializedKey() {
-    return work.getWorkItem().getKey();
+    return work == null ? null : work.getWorkItem().getKey();
   }
 
   public WindmillComputationKey getComputationKey() {
     return computationKey;
   }
 
   public long getWorkToken() {
-    return work.getWorkItem().getWorkToken();
+    return getWorkItem().getWorkToken();
   }
 
   public Windmill.WorkItem getWorkItem() {
-    return work.getWorkItem();
+    return checkNotNull(work).getWorkItem();
   }
 
   public Windmill.WorkItemCommitRequest.Builder getOutputBuilder() {
@@ -390,7 +389,7 @@ public void invalidateCache() {
   public UnboundedSource.@Nullable CheckpointMark getReaderCheckpoint(
       Coder<? extends UnboundedSource.CheckpointMark> coder) {
     try {
-      ByteString sourceStateState = work.getWorkItem().getSourceState().getState();
+      ByteString sourceStateState = getWorkItem().getSourceState().getState();
       if (sourceStateState.isEmpty()) {
         return null;
       }
@@ -737,7 +736,7 @@ public void start(
               key,
               stateFamily,
               stateReader,
-              work.getWorkItem().getIsNewKey(),
+              checkNotNull(work).getWorkItem().getIsNewKey(),
               cacheForKey.forFamily(stateFamily),
               scopedReadStateSupplier);
 

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

@@ -83,11 +83,11 @@ public final void executeWork(
    * cannot be reused.
    */
   public final void invalidate() {
-    context().invalidateCache();
     try {
+      context().invalidateCache();
       workExecutor().close();
     } catch (Exception e) {
-      LOG.warn("Failed to close map task executor: ", e);
+      LOG.warn("Failed to invalidate ComputationWorkExecutor: ", e);
     }
   }
 

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationWorkExecutorTest.java

@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming;
+
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verifyNoInteractions;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.beam.runners.core.metrics.ExecutionStateTracker;
+import org.apache.beam.runners.dataflow.worker.DataflowWorkExecutor;
+import org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ComputationWorkExecutorTest {
+
+  private final DataflowWorkExecutor dataflowWorkExecutor = mock(DataflowWorkExecutor.class);
+  private final StreamingModeExecutionContext context = mock(StreamingModeExecutionContext.class);
+  private ComputationWorkExecutor computationWorkExecutor;
+
+  @Before
+  public void setUp() {
+    computationWorkExecutor =
+        ComputationWorkExecutor.builder()
+            .setWorkExecutor(dataflowWorkExecutor)
+            .setContext(context)
+            .setExecutionStateTracker(mock(ExecutionStateTracker.class))
+            .build();
+  }
+
+  @Test
+  public void testInvalidate_handlesException() {
+    NullPointerException npe = new NullPointerException("something bad happened");
+    doThrow(npe).when(context).invalidateCache();
+    computationWorkExecutor.invalidate();
+    verifyNoInteractions(dataflowWorkExecutor);
+    AtomicBoolean verifyContextInvalidated = new AtomicBoolean(false);
+    Throwable e = new RuntimeException("something bad happened 2");
+    doThrow(e).when(dataflowWorkExecutor).close();
+    doAnswer(
+            ignored -> {
+              verifyContextInvalidated.set(true);
+              return null;
+            })
+        .when(context)
+        .invalidateCache();
+    computationWorkExecutor.invalidate();
+    assertTrue(verifyContextInvalidated.get());
+  }
+}