[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

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json

@@ -1,5 +1,5 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run!",
-  "modification": 1,
+  "modification": 2,
 }
 

.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow_Streaming.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run!",
-  "modification":  1,
+  "modification":  2,
 }

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

@@ -56,6 +56,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 +158,8 @@ public class StreamingModeExecutionContext extends DataflowExecutionContext<Step
    */
   private @Nullable UnboundedReader<?> activeReader;
 
+  private @Nullable WorkExecutor workExecutor;
+
   public StreamingModeExecutionContext(
       CounterFactory counterFactory,
       String computationId,
@@ -240,9 +243,11 @@ 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.computationKey = WindmillComputationKey.create(computationId, work.getShardedKey());
     this.sideInputStateFetcher = sideInputStateFetcher;
     StreamingGlobalConfig config = globalConfigHandle.getConfig();
@@ -270,6 +275,11 @@ public void start(
     }
   }
 
+  public void finishKey() throws Exception {
+    checkNotNull(workExecutor, "workExecutor must be set before calling finishKey()");
+    workExecutor.finishKey();
+  }
+
   /**
    * Ensure that the processing time is greater than any fired processing time timers. Otherwise, a
    * trigger could ignore the timer and orphan the window.

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

@@ -109,20 +109,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,18 @@ 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;
+        try {
+          context.finishKey();
+        } catch (Exception e) {
+          throw new RuntimeException(e);
+        }
         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

@@ -182,6 +182,11 @@ public boolean start() throws IOException {
         @Override
         public boolean advance() throws IOException {
           current = null;
+          try {
+            context.finishKey();
+          } catch (Exception e) {
+            throw new RuntimeException(e);
+          }
           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/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;