[Dataflow Streaming] Add Operation::finishKey() and move processTimers into it (#38430)

* [Dataflow] Add Operation::finishKey() and move timers into it

Moving the processTimers call from finish() to finishKey().

In upcoming changes there'll be multiple streaming work items
in a single beam bundle. With multiple work items, we've to process
elements and timers of each work item before moving to the next work
items.

finishKey() will be called by the NativeIterator classes after iterating
through all elements from a work item.

Batch processes timers in BatchModeUngroupingParDoFn and does not rely
on the processTimers() in ParDoOperation::finish(). So removing the
processTimers() call from ParDoOperation::finish() is safe. Batch also
does not use the new finishKey() method.

Author: arunpandianp

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

@@ -117,20 +117,12 @@ public NativeReader<?> create(
 
   @Override
   public NativeReaderIterator<WindowedValue<T>> iterator() throws IOException {
-    return new PubsubReaderIterator(context.getWorkItem());
+    return new PubsubReaderIterator();
   }
 
   class PubsubReaderIterator extends WindmillReaderIteratorBase<T> {
-    protected PubsubReaderIterator(Windmill.WorkItem work) {
-      super(work, skipUndecodableElements);
-    }
-
-    @Override
-    public boolean advance() throws IOException {
-      if (context.workIsFailed()) {
-        return false;
-      }
-      return super.advance();
+    protected PubsubReaderIterator() {
+      super(context, skipUndecodableElements);
     }
 
     @Override

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

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.runners.dataflow.worker;
 
+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.checkNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
@@ -56,6 +57,7 @@
 import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInput;
 import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputState;
 import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher;
+import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataId;
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
@@ -157,6 +159,9 @@ public class StreamingModeExecutionContext extends DataflowExecutionContext<Step
    */
   private @Nullable UnboundedReader<?> activeReader;
 
+  private @Nullable WorkExecutor workExecutor;
+  private boolean finishKeyCalled = false;
+
   public StreamingModeExecutionContext(
       CounterFactory counterFactory,
       String computationId,
@@ -240,9 +245,12 @@ public void start(
       Work work,
       WindmillStateReader stateReader,
       SideInputStateFetcher sideInputStateFetcher,
-      Windmill.WorkItemCommitRequest.Builder outputBuilder) {
+      Windmill.WorkItemCommitRequest.Builder outputBuilder,
+      WorkExecutor workExecutor) {
     this.key = key;
     this.work = work;
+    this.workExecutor = workExecutor;
+    this.finishKeyCalled = false;
     this.computationKey = WindmillComputationKey.create(computationId, work.getShardedKey());
     this.sideInputStateFetcher = sideInputStateFetcher;
     StreamingGlobalConfig config = globalConfigHandle.getConfig();
@@ -271,6 +279,17 @@ public void start(
     }
   }
 
+  public void finishKey() {
+    checkState(!finishKeyCalled, "finishKey was already called");
+    checkStateNotNull(workExecutor, "workExecutor must be set before calling finishKey()");
+    try {
+      workExecutor.finishKey();
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+    this.finishKeyCalled = true;
+  }
+
   /**
    * Ensure that the processing time is greater than any fired processing time timers. Otherwise, a
    * trigger could ignore the timer and orphan the window.
@@ -452,6 +471,7 @@ public void invalidateCache() {
   }
 
   public Map<Long, Pair<Instant, Runnable>> flushState() {
+    checkState(finishKeyCalled, "finishKey must be called before flushState");
     Map<Long, Pair<Instant, Runnable>> callbacks = new HashMap<>();
 
     for (StepContext stepContext : getAllStepContexts()) {

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

@@ -110,20 +110,12 @@ public NativeReader<?> create(
 
   @Override
   public NativeReaderIterator<WindowedValue<T>> iterator() throws IOException {
-    return new UngroupedWindmillReaderIterator(context.getWorkItem());
+    return new UngroupedWindmillReaderIterator();
   }
 
   class UngroupedWindmillReaderIterator extends WindmillReaderIteratorBase<T> {
-    UngroupedWindmillReaderIterator(Windmill.WorkItem work) {
-      super(work, skipUndecodableElements);
-    }
-
-    @Override
-    public boolean advance() throws IOException {
-      if (context.workIsFailed()) {
-        return false;
-      }
-      return super.advance();
+    UngroupedWindmillReaderIterator() {
+      super(context, skipUndecodableElements);
     }
 
     @Override

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

@@ -34,6 +34,7 @@
  */
 public abstract class WindmillReaderIteratorBase<T>
     extends NativeReader.NativeReaderIterator<WindowedValue<T>> {
+  private final StreamingModeExecutionContext context;
   private final Windmill.WorkItem work;
   private int bundleIndex = 0;
   private int messageIndex = -1;
@@ -42,9 +43,10 @@ public abstract class WindmillReaderIteratorBase<T>
   private static final Logger LOG = LoggerFactory.getLogger(WindmillReaderIteratorBase.class);
 
   protected WindmillReaderIteratorBase(
-      Windmill.WorkItem work, ValueProvider<Boolean> skipUndecodableElements) {
+      StreamingModeExecutionContext context, ValueProvider<Boolean> skipUndecodableElements) {
+    this.context = context;
     this.skipUndecodableElements = skipUndecodableElements;
-    this.work = work;
+    this.work = context.getWorkItem();
   }
 
   @Override
@@ -54,9 +56,14 @@ public boolean start() throws IOException {
 
   @Override
   public boolean advance() throws IOException {
+    if (context.workIsFailed()) {
+      throw new WorkItemCancelledException(context.getWorkItem().getShardingKey());
+    }
+
     while (true) {
       if (bundleIndex >= work.getMessageBundlesCount()) {
         current = null;
+        context.finishKey();
         return false;
       }
       Windmill.InputMessageBundle bundle = work.getMessageBundles(bundleIndex);

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

@@ -156,6 +156,7 @@ public NativeReaderIterator<WindowedValue<KeyedWorkItem<K, T>>> iterator() throw
       return new NativeReaderIterator<WindowedValue<KeyedWorkItem<K, T>>>() {
         @Override
         public boolean start() throws IOException {
+          context.finishKey();
           return false;
         }
 
@@ -182,6 +183,7 @@ public boolean start() throws IOException {
         @Override
         public boolean advance() throws IOException {
           current = null;
+          context.finishKey();
           return false;
         }
 

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

@@ -89,6 +89,9 @@ public void execute() throws Exception {
     LOG.debug("Source operation execution complete");
   }
 
+  @Override
+  public void finishKey() throws Exception {}
+
   @Override
   public SourceOperationResponse getResponse() {
     return response;

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

@@ -824,6 +824,7 @@ public boolean start() throws IOException {
       }
       try {
         if (!reader.start()) {
+          context.finishKey();
           return false;
         }
       } catch (Exception e) {
@@ -841,10 +842,13 @@ public boolean advance() throws IOException {
       // that there are regular checkpoints and that state does not become too large.
       BackOff backoff = backoffFactory.backoff();
       while (true) {
+        if (context.workIsFailed()) {
+          throw new WorkItemCancelledException(context.getWorkItem().getShardingKey());
+        }
         if (elemsRead >= maxElems
             || Instant.now().isAfter(endTime)
-            || context.isSinkFullHintSet()
-            || context.workIsFailed()) {
+            || context.isSinkFullHintSet()) {
+          context.finishKey();
           return false;
         }
         try {
@@ -857,6 +861,7 @@ public boolean advance() throws IOException {
         }
         long nextBackoff = backoff.nextBackOffMillis();
         if (nextBackoff == BackOff.STOP) {
+          context.finishKey();
           return false;
         }
         Uninterruptibles.sleepUninterruptibly(nextBackoff, TimeUnit.MILLISECONDS);

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

@@ -74,7 +74,7 @@ public final void executeWork(
       SideInputStateFetcher sideInputStateFetcher,
       Windmill.WorkItemCommitRequest.Builder outputBuilder)
       throws Exception {
-    context().start(key, work, stateReader, sideInputStateFetcher, outputBuilder);
+    context().start(key, work, stateReader, sideInputStateFetcher, outputBuilder, workExecutor());
     workExecutor().execute();
   }
 

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

@@ -43,6 +43,9 @@ public void process(Object elem) throws Exception {
     }
   }
 
+  @Override
+  public void finishKey() throws Exception {}
+
   @Override
   public boolean supportsRestart() {
     return true;

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

@@ -112,6 +112,13 @@ public void execute() throws Exception {
     // TODO: support for success / failure ports?
   }
 
+  @Override
+  public void finishKey() throws Exception {
+    for (Operation op : operations) {
+      op.finishKey();
+    }
+  }
+
   @Override
   public NativeReader.Progress getWorkerProgress() throws Exception {
     return getReadOperation().getProgress();