[Dataflow Streaming] [Multi Key] Drop failed work in BoundedQueueExecutor::pollWork

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

@@ -53,6 +53,7 @@
 import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
 import org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
 import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
 import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException;
 import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
 import org.apache.beam.runners.dataflow.worker.streaming.Work;
@@ -172,10 +173,7 @@ public class StreamingModeExecutionContext
   private @Nullable WorkExecutor workExecutor;
   private boolean finishKeyCalled = false;
 
-  @SuppressWarnings("UnusedVariable")
   private @Nullable BoundedQueueExecutor workQueueExecutor;
-
-  @SuppressWarnings("UnusedVariable")
   private @Nullable BoundedQueueExecutorWorkHandle budgetHandle;
 
   private final HotKeyLogger hotKeyLogger;
@@ -192,8 +190,8 @@ public interface KeyTransitionListener {
     void onKeyTransition(@Nullable Work oldWork, Work newWork);
   }
 
-  @SuppressWarnings("UnusedVariable")
   private @Nullable KeyTransitionListener keyTransitionListener;
+  private @Nullable FailedWorkHandler onFailedWorkHandler;
 
   private List<Work> executedWorks = Collections.emptyList();
   private List<Windmill.WorkItemCommitRequest.Builder> outputBuilders = Collections.emptyList();
@@ -335,6 +333,7 @@ public void reset() {
     this.workQueueExecutor = null;
     this.budgetHandle = null;
     this.keyTransitionListener = null;
+    this.onFailedWorkHandler = null;
     this.work = null;
     this.key = null;
     this.outputBuilder = null;
@@ -350,7 +349,8 @@ public void start(
       BoundedQueueExecutor workQueueExecutor,
       BoundedQueueExecutorWorkHandle budgetHandle,
       @Nullable Coder<?> keyCoder,
-      KeyTransitionListener keyTransitionListener)
+      KeyTransitionListener keyTransitionListener,
+      FailedWorkHandler onFailedWorkHandler)
       throws CoderException {
     reset();
     this.executedWorks = new ArrayList<>();
@@ -361,6 +361,7 @@ public void start(
     this.workQueueExecutor = workQueueExecutor;
     this.budgetHandle = budgetHandle;
     this.keyTransitionListener = keyTransitionListener;
+    this.onFailedWorkHandler = checkStateNotNull(onFailedWorkHandler);
 
     this.workItemsPolled = 1;
     this.bundleStartTimeNanos = System.nanoTime();
@@ -779,7 +780,11 @@ public boolean advance() throws CoderException {
 
     @Nullable
     ExecutableWork additionalWork =
-        executor.pollWork(computationId, activeWork.getKeyGroup(), handle);
+        executor.pollWork(
+            computationId,
+            activeWork.getKeyGroup(),
+            handle,
+            checkStateNotNull(onFailedWorkHandler));
     if (additionalWork != null) {
       flushStateInternal();
       Work newWork = additionalWork.work();

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

@@ -65,7 +65,8 @@ public final StreamingModeExecutionContext executeWork(
       Work work,
       BoundedQueueExecutor workQueueExecutor,
       BoundedQueueExecutorWorkHandle budgetHandle,
-      KeyTransitionListener keyTransitionListener)
+      KeyTransitionListener keyTransitionListener,
+      FailedWorkHandler onFailedWorkHandler)
       throws Exception {
     context()
         .start(
@@ -74,7 +75,8 @@ public final StreamingModeExecutionContext executeWork(
             workQueueExecutor,
             budgetHandle,
             keyCoder().orElse(null),
-            keyTransitionListener);
+            keyTransitionListener,
+            onFailedWorkHandler);
     workExecutor().execute();
     return context();
   }

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

@@ -0,0 +1,24 @@
+/*
+ * 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;
+
+/** Handler for failed {@link Work}. */
+@FunctionalInterface
+public interface FailedWorkHandler {
+  void onFailedWork(Work work);
+}

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

@@ -18,6 +18,7 @@
 package org.apache.beam.runners.dataflow.worker.util;
 
 import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+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 java.util.ArrayList;
@@ -31,6 +32,7 @@
 import javax.annotation.concurrent.GuardedBy;
 import org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
 import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
 import org.apache.beam.runners.dataflow.worker.streaming.Work;
 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;
@@ -387,20 +389,32 @@ BoundedQueueExecutorWorkHandleImpl createBudgetHandle(Work work, long bytes) {
   }
 
   public @Nullable ExecutableWork pollWork(
-      String computationId, Work.KeyGroup keyGroup, BoundedQueueExecutorWorkHandle handle) {
+      String computationId,
+      Work.KeyGroup keyGroup,
+      BoundedQueueExecutorWorkHandle handle,
+      FailedWorkHandler onFailedWorkHandler) {
     checkArgument(
         computationId != null && keyGroup != null && !keyGroup.equals(Work.KeyGroup.DEFAULT));
     checkArgument(handle instanceof BoundedQueueExecutorWorkHandleImpl);
+    checkStateNotNull(onFailedWorkHandler);
     BoundedQueueExecutorWorkHandleImpl internalHandle = (BoundedQueueExecutorWorkHandleImpl) handle;
     if (keyGroupWorkQueue == null) {
       return null;
     }
-    @Nullable QueuedWork queuedWork = keyGroupWorkQueue.pollWork(computationId, keyGroup);
-    if (queuedWork == null) {
-      return null;
+    while (true) {
+      @Nullable QueuedWork queuedWork = keyGroupWorkQueue.pollWork(computationId, keyGroup);
+      if (queuedWork == null) {
+        return null;
+      }
+      Work work = queuedWork.getWork().work();
+      if (work.isFailed()) {
+        queuedWork.getHandle().close();
+        onFailedWorkHandler.onFailedWork(work);
+        continue;
+      }
+      internalHandle.merge(queuedWork.getHandle());
+      return queuedWork.getWork();
     }
-    internalHandle.merge(queuedWork.getHandle());
-    return queuedWork.getWork();
   }
 
   private void decrementCounters(int elements, long bytes) {

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java

@@ -45,6 +45,7 @@
 import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
 import org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor;
 import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
 import org.apache.beam.runners.dataflow.worker.streaming.StageInfo;
 import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
 import org.apache.beam.runners.dataflow.worker.streaming.Work;
@@ -321,8 +322,11 @@ private ExecuteWorkResult executeWork(
     try {
       StreamingModeExecutionContext context = computationWorkExecutor.context();
 
+      FailedWorkHandler onFailedWorkHandler = getFailedWorkHandler(computationState);
+
       // Blocks while executing work.
-      computationWorkExecutor.executeWork(work, workExecutor, handle, keyTransitionListener);
+      computationWorkExecutor.executeWork(
+          work, workExecutor, handle, keyTransitionListener, onFailedWorkHandler);
 
       List<Work> workBatch;
       List<Windmill.WorkItemCommitRequest> workItemCommits;
@@ -465,14 +469,10 @@ private void handleProcessWorkFailure(
             ExecutableWork.create(w, (retry, h) -> processWork(computationState, retry, h)));
       }
 
+      FailedWorkHandler onFailedWorkHandler = getFailedWorkHandler(computationState);
+
       workFailureProcessor.logAndProcessFailureBatch(
-          computationId,
-          systemName,
-          executableWorks,
-          t,
-          invalidWork ->
-              computationState.completeWorkAndScheduleNextWorkForKey(
-                  invalidWork.getShardedKey(), invalidWork.id()));
+          computationId, systemName, executableWorks, t, onFailedWorkHandler);
     } catch (OutOfMemoryError oom) {
       throw oom;
     } catch (Throwable t2) {
@@ -481,6 +481,12 @@ private void handleProcessWorkFailure(
     }
   }
 
+  private static FailedWorkHandler getFailedWorkHandler(ComputationState computationState) {
+    return failedWork ->
+        computationState.completeWorkAndScheduleNextWorkForKey(
+            failedWork.getShardedKey(), failedWork.id());
+  }
+
   private void recordProcessingTime(
       StageInfo stageInfo, List<Work> workBatch, long processingStartTimeNanos) {
     long processingTimeMsecs =

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java

@@ -19,12 +19,12 @@
 
 import java.util.List;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Consumer;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.ThreadSafe;
 import org.apache.beam.runners.dataflow.worker.status.LastExceptionDataProvider;
 import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
+import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
 import org.apache.beam.runners.dataflow.worker.streaming.Work;
 import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
 import org.apache.beam.sdk.annotations.Internal;
@@ -103,7 +103,7 @@ public void logAndProcessFailureBatch(
       String systemName,
       List<ExecutableWork> executableWorks,
       Throwable t,
-      Consumer<Work> onInvalidWork)
+      FailedWorkHandler onFailedWorkHandler)
       throws Throwable {
     List<ExecutableWork> worksToRetryLocally = new java.util.ArrayList<>();
 
@@ -112,7 +112,7 @@ public void logAndProcessFailureBatch(
         case DO_NOT_RETRY:
           // Consider the item invalid. It will eventually be retried by Windmill if it still needs
           // to be processed.
-          onInvalidWork.accept(executableWork.work());
+          onFailedWorkHandler.onFailedWork(executableWork.work());
           break;
         case RETRY_LOCALLY:
           // Try again after some delay and at the end of the queue to avoid a tight loop.