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

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();

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

@@ -137,6 +137,9 @@ public void finish() throws Exception {
     }
   }
 
+  /** Called when all elements for a specific key have been processed. */
+  public abstract void finishKey() throws Exception;
+
   /** Aborts this Operation's execution. */
   public void abort() throws Exception {
     synchronized (initializationStateLock) {

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

@@ -45,13 +45,18 @@ public void process(Object elem) throws Exception {
     }
   }
 
+  // Batch mode does not use this method and instead relies on BatchModeUngroupingParDoFn
+  // to process timers per key.
   @Override
-  public void finish() throws Exception {
+  public void finishKey() throws Exception {
     try (Closeable scope = context.enterProcessTimers()) {
       checkStarted();
       fn.processTimers();
     }
+  }
 
+  @Override
+  public void finish() throws Exception {
     try (Closeable scope = context.enterFinish()) {
       fn.finishBundle();
       super.finish();

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

@@ -271,6 +271,9 @@ public void finish() throws Exception {
     }
   }
 
+  @Override
+  public void finishKey() throws Exception {}
+
   @Override
   public void abort() throws Exception {
     if (readerIterator != null) {

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

@@ -34,6 +34,9 @@ public interface WorkExecutor extends AutoCloseable {
   /** Executes the task. */
   public abstract void execute() throws Exception;
 
+  /** Called when all elements for a specific key have been processed. */
+  void finishKey() throws Exception;
+
   /**
    * Returns the worker's current progress.
    *

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

@@ -105,6 +105,9 @@ public void finish() throws Exception {
     }
   }
 
+  @Override
+  public void finishKey() throws Exception {}
+
   @Override
   public void abort() throws Exception {
     if (writer == null) {

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/IntrinsicMapTaskExecutorTest.java

@@ -104,6 +104,9 @@ public void abort() throws Exception {
       aborted = true;
       super.abort();
     }
+
+    @Override
+    public void finishKey() throws Exception {}
   }
 
   // A mock ReadOperation fed to a MapTaskExecutor in test.
@@ -312,6 +315,9 @@ public void start() throws Exception {
                   Metrics.counter("TestMetric", "MetricCounter").inc(1L);
                 }
               }
+
+              @Override
+              public void finishKey() throws Exception {}
             },
             new Operation(new OutputReceiver[] {}, context2) {
               @Override
@@ -321,6 +327,9 @@ public void start() throws Exception {
                   Metrics.counter("TestMetric", "MetricCounter").inc(2L);
                 }
               }
+
+              @Override
+              public void finishKey() throws Exception {}
             },
             new Operation(new OutputReceiver[] {}, context3) {
               @Override
@@ -330,6 +339,9 @@ public void start() throws Exception {
                   Metrics.counter("TestMetric", "MetricCounter").inc(3L);
                 }
               }
+
+              @Override
+              public void finishKey() throws Exception {}
             });
 
     try (IntrinsicMapTaskExecutor executor =

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java

@@ -3657,8 +3657,8 @@ public void testActiveWorkFailure() throws Exception {
         server.waitForAndGetCommitsWithTimeout(1, Duration.standardSeconds(5));
     assertEquals(1, commits.size());
 
-    assertEquals(0, BlockingFn.teardownCounter.get());
-    assertEquals(1, BlockingFn.setupCounter.get());
+    assertEquals(1, BlockingFn.teardownCounter.get());
+    assertEquals(2, BlockingFn.setupCounter.get());
 
     worker.stop();
   }