Add DoFnRunner::finishKey() method

.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/core-java/src/main/java/org/apache/beam/runners/core/DoFnRunner.java

@@ -63,6 +63,12 @@ public interface DoFnRunner<InputT extends @Nullable Object, OutputT extends @Nu
   <KeyT extends @Nullable Object> void onWindowExpiration(
       BoundedWindow window, Instant timestamp, KeyT key);
 
+  /**
+   * Performs per-key cleanup or processing after all elements and timers for a key have been
+   * processed.
+   */
+  void finishKey();
+
   /**
    * Returns the underlying fn instance.
    *

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

@@ -101,6 +101,11 @@ public void finishBundle() {
     doFnRunner.finishBundle();
   }
 
+  @Override
+  public void finishKey() {
+    doFnRunner.finishKey();
+  }
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     doFnRunner.onWindowExpiration(window, timestamp, key);

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

@@ -229,6 +229,9 @@ public void finishBundle() {
     }
   }
 
+  @Override
+  public void finishKey() {}
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     invoker.invokeOnWindowExpiration(

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

@@ -131,6 +131,11 @@ public void finishBundle() {
     doFnRunner.finishBundle();
   }
 
+  @Override
+  public void finishKey() {
+    doFnRunner.finishKey();
+  }
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     doFnRunner.onWindowExpiration(window, timestamp, key);

runners/core-java/src/test/java/org/apache/beam/runners/core/SimplePushbackSideInputDoFnRunnerTest.java

@@ -379,6 +379,9 @@ public void finishBundle() {
       finished = true;
     }
 
+    @Override
+    public void finishKey() {}
+
     @Override
     public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {}
   }

runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/DoFnRunnerWithMetricsUpdate.java

@@ -105,6 +105,9 @@ public void finishBundle() {
     container.updateMetrics(stepName);
   }
 
+  @Override
+  public void finishKey() {}
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     delegate.onWindowExpiration(window, timestamp, key);

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java

@@ -1064,6 +1064,9 @@ public void finishBundle() {
       }
     }
 
+    @Override
+    public void finishKey() {}
+
     @Override
     public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {}
 

runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/stableinput/BufferingDoFnRunner.java

@@ -255,6 +255,9 @@ public void finishBundle() {
     Optional.ofNullable(finishBundleCallback).ifPresent(Runnable::run);
   }
 
+  @Override
+  public void finishKey() {}
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {}
 

runners/flink/src/test/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperatorTest.java

@@ -523,6 +523,9 @@ public void finishBundle() {
                 wrappedRunner.finishBundle();
               }
 
+              @Override
+              public void finishKey() {}
+
               @Override
               public <KeyT> void onWindowExpiration(
                   BoundedWindow window, Instant timestamp, KeyT key) {

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

@@ -119,6 +119,9 @@ public void processTimers() throws Exception {
       // Nothing.
     }
 
+    @Override
+    public void finishKey() throws Exception {}
+
     @Override
     public void finishBundle() throws Exception {
       receiver = null;

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

@@ -73,6 +73,9 @@ public void processTimers() throws Exception {
     // The timers for the underlying ParDoFn are processed at the end of each element
   }
 
+  @Override
+  public void finishKey() throws Exception {}
+
   @Override
   public void finishBundle() throws Exception {
     underlyingParDoFn.finishBundle();

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

@@ -114,6 +114,9 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey() throws Exception {}
+
     @Override
     public void finishBundle() {}
 

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

@@ -128,6 +128,11 @@ public void finishBundle() {
     simpleRunner.finishBundle();
   }
 
+  @Override
+  public void finishKey() {
+    simpleRunner.finishKey();
+  }
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {
     simpleRunner.onWindowExpiration(window, timestamp, key);

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

@@ -47,6 +47,11 @@ public void processTimers() throws Exception {
     delegate.processTimers();
   }
 
+  @Override
+  public void finishKey() throws Exception {
+    delegate.finishKey();
+  }
+
   @Override
   public void finishBundle() throws Exception {
     delegate.finishBundle();

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

@@ -102,6 +102,9 @@ private void invokeProcessElement(WindowedValue<InputT> elem) {
   @Override
   public void finishBundle() {}
 
+  @Override
+  public void finishKey() {}
+
   @Override
   public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, KeyT key) {}
 

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

@@ -142,6 +142,13 @@ public void processTimers() throws Exception {
     // it here to build a KeyedWorkItem
   }
 
+  @Override
+  public void finishKey() throws Exception {
+    if (fnRunner != null) {
+      fnRunner.finishKey();
+    }
+  }
+
   @Override
   public void finishBundle() throws Exception {
     checkState(fnRunner != null);

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

@@ -98,6 +98,9 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey() throws Exception {}
+
     @Override
     public void finishBundle() {}
 

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

@@ -317,6 +317,9 @@ public void processElement(Object elem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey() throws Exception {}
+
     @Override
     public void finishBundle() throws Exception {
       groupingTable.flush(receiver);
@@ -377,10 +380,14 @@ public void processElement(Object elem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey() throws Exception {
+      sideInputFetcher.persist();
+    }
+
     @Override
     public void finishBundle() throws Exception {
       groupingTable.flush(receiver);
-      sideInputFetcher.persist();
     }
 
     @Override

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/ReifyTimestampAndWindowsParDoFnFactory.java

@@ -86,6 +86,9 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey() throws Exception {}
+
     @Override
     public void finishBundle() throws Exception {
       this.receiver = null;

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

@@ -363,6 +363,13 @@ public void processTimers() throws Exception {
     processTimers(TimerType.SYSTEM, stepContext, windowCoder);
   }
 
+  @Override
+  public void finishKey() throws Exception {
+    if (fnRunner != null) {
+      fnRunner.finishKey();
+    }
+  }
+
   private void processUserTimer(TimerData timer) throws Exception {
     if (fnSignature.timerDeclarations().containsKey(timer.getTimerId())
         || fnSignature.timerFamilyDeclarations().containsKey(timer.getTimerFamilyId())) {

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

@@ -152,6 +152,11 @@ public <KeyT> void onTimer(
   @Override
   public void finishBundle() {
     simpleDoFnRunner.finishBundle();
+  }
+
+  @Override
+  public void finishKey() {
+    simpleDoFnRunner.finishKey();
     sideInputFetcher.persist();
   }
 

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,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 +244,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();
@@ -270,6 +277,17 @@ public void start(
     }
   }
 
+  public void finishKey() {
+    checkState(!finishKeyCalled, "finishKey was already called");
+    checkNotNull(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.
@@ -451,6 +469,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/StreamingPCollectionViewWriterParDoFn.java

@@ -73,6 +73,9 @@ public void processElement(Object element) throws Exception {
   @Override
   public void processTimers() {}
 
+  @Override
+  public void finishKey() throws Exception {}
+
   @Override
   public void finishBundle() throws Exception {}
 

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

@@ -94,6 +94,11 @@ public <KeyT> void onTimer(
   @Override
   public void finishBundle() {
     simpleDoFnRunner.finishBundle();
+  }
+
+  @Override
+  public void finishKey() {
+    simpleDoFnRunner.finishKey();
     sideInputFetcher.persist();
   }
 

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

@@ -149,6 +149,9 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey() throws Exception {}
+
     @Override
     public void finishBundle() {}
 

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