Add DoFnRunner::finishKey method.

In upcoming changes there'll be multiple dataflow 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.

The new finishKey method allows the DoFnRunners to
cleanup/persist state (that should not be carried over) before switching work items
on multi key bundles.

Streaming SideInputDoFnRunners are the only classes using the finishKey
method right now.

The finishKey() is not exposed to DoFns and is not visible in user apis.

Author: arunpandianp

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,11 @@ public void finishBundle() {
     }
   }
 
+  @Override
+  public void finishKey() {
+    // Do nothing by default.
+  }
+
   @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/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/AssignWindowsParDoFnFactory.java

@@ -119,6 +119,11 @@ public void processTimers() throws Exception {
       // Nothing.
     }
 
+    @Override
+    public void finishKey() throws Exception {
+      // Nothing.
+    }
+
     @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,11 @@ 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 {
+    // Nothing.
+  }
+
   @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,11 @@ public void processElement(Object untypedElem) throws Exception {
     @Override
     public void processTimers() {}
 
+    @Override
+    public void finishKey() throws Exception {
+      // Nothing.
+    }
+
     @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();