[Dataflow Streaming] Commit size validation for multi key commits

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

@@ -55,6 +55,7 @@
 import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
 import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
 import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException;
+import org.apache.beam.runners.dataflow.worker.streaming.MultiKeyCommitValidationException;
 import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
 import org.apache.beam.runners.dataflow.worker.streaming.Work;
 import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig;
@@ -193,7 +194,6 @@ public interface KeyTransitionListener {
   private @Nullable KeyTransitionListener keyTransitionListener;
   private @Nullable FailedWorkHandler onFailedWorkHandler;
 
-  private List<Work> executedWorks = Collections.emptyList();
   private List<Windmill.WorkItemCommitRequest.Builder> outputBuilders = Collections.emptyList();
 
   // Map<finalizerId, Pair<callbackExpiration, callback>>
@@ -319,7 +319,6 @@ public byte[] getCurrentRecordOffset() {
   public void reset() {
     // these lists and maps are returned to callers after processing
     // don't clear and reuse, instead reset the reference.
-    this.executedWorks = Collections.emptyList();
     this.outputBuilders = Collections.emptyList();
     this.finalizationCallbacks = Collections.emptyMap();
     // Work from prior bundles might have a reference to the old workBatchFailed.
@@ -353,7 +352,6 @@ public void start(
       FailedWorkHandler onFailedWorkHandler)
       throws CoderException {
     reset();
-    this.executedWorks = new ArrayList<>();
     this.outputBuilders = new ArrayList<>();
     this.finalizationCallbacks = new HashMap<>();
     this.keyCoder = keyCoder;
@@ -578,11 +576,13 @@ public void setActiveReader(UnboundedReader<?> reader) {
 
   /** Invalidate the state and reader caches for this computation and key. */
   public void invalidateCache() {
-    for (Work w : executedWorks) {
-      WindmillComputationKey compKey =
-          WindmillComputationKey.create(computationId, w.getShardedKey());
-      readerCache.invalidateReader(compKey);
-      stateCache.invalidate(w.getShardedKey());
+    if (budgetHandle != null) {
+      for (Work w : budgetHandle.getWorkBatch()) {
+        WindmillComputationKey compKey =
+            WindmillComputationKey.create(computationId, w.getShardedKey());
+        readerCache.invalidateReader(compKey);
+        stateCache.invalidate(w.getShardedKey());
+      }
     }
     if (activeReader != null) {
       try {
@@ -718,6 +718,23 @@ private void validateCommitRequestSize() {
       return;
     }
 
+    // If this is a multi-key work item, then we need to retry all of the individual work items
+    // without merging so that we can identify large commits to truncate.
+    // TODO: Can we request truncation without retrying if the first commit exceed the limits?
+    BoundedQueueExecutorWorkHandle handle = checkNotNull(budgetHandle);
+    List<Work> currentBatch = handle.getWorkBatch();
+    checkState(!currentBatch.isEmpty());
+    if (currentBatch.size() > 1) {
+      LOG.warn(
+          "Windmill Commit limit exceeded on a multi key bundle. Retrying without batching. Batch size: {}",
+          currentBatch.size());
+      for (Work w : currentBatch) {
+        w.setMultiKeyBatchingDisabled(true);
+      }
+      throw new MultiKeyCommitValidationException(
+          "Commit size validation failed for batch. Retrying individually.");
+    }
+
     KeyCommitTooLargeException e =
         KeyCommitTooLargeException.causedBy(
             systemName, byteLimit, commitRequest, key, hotKeyLoggingEnabled);
@@ -731,11 +748,6 @@ private void validateCommitRequestSize() {
         buildWorkItemTruncationRequestBuilder(currentWork, estimatedCommitSize);
     currentBuilder.clear();
     currentBuilder.mergeFrom(truncationBuilder.build());
-
-    // TODO: throw and retry when truncation is not on a single key bundle.
-    checkState(
-        !multiKeyBundleOptions.multiKeyBundleEnabled(),
-        "Commit truncation not implemented for multikey bundles");
   }
 
   private Windmill.WorkItemCommitRequest.Builder buildWorkItemTruncationRequestBuilder(
@@ -774,7 +786,9 @@ public boolean advance() throws CoderException {
       throw new WorkItemCancelledException(activeWork.getWorkItem().getShardingKey());
     }
 
-    if (activeWork.getKeyGroup().equals(Work.KeyGroup.DEFAULT) || shouldStopBatching()) {
+    if (activeWork.getKeyGroup().equals(Work.KeyGroup.DEFAULT)
+        || activeWork.isMultiKeyBatchingDisabled()
+        || shouldStopBatching()) {
       return false;
     }
 
@@ -797,7 +811,6 @@ public boolean advance() throws CoderException {
   }
 
   private boolean shouldStopBatching() {
-    // TODO: stop batching if the previous work item requested truncation
     if (workItemsPolled >= multiKeyBundleOptions.maxKeyGroupBatchSize()) {
       return true;
     }
@@ -821,7 +834,6 @@ private void startForNewKey(Work newWork) throws CoderException {
     this.outputBuilder = createOutputBuilder(newWork);
     this.outputBuilders.add(this.outputBuilder);
     newWork.setOnFailureListener(this.workBatchFailed);
-    this.executedWorks.add(newWork);
 
     logHotKeyIfDetected(newWork, this.key);
 
@@ -862,11 +874,6 @@ public List<Windmill.WorkItemCommitRequest> getWorkItemCommits() {
     return commits;
   }
 
-  // Returns list of Work that was executed in the bundle
-  public List<Work> getExecutedWorks() {
-    return executedWorks;
-  }
-
   // Returns finalization callbacks recorded during the bundle execution
   public Map<Long, Pair<Instant, Runnable>> getFinalizationCallbacks() {
     return finalizationCallbacks;

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

@@ -17,13 +17,15 @@
  */
 package org.apache.beam.runners.dataflow.worker.streaming;
 
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import java.util.List;
 
 /**
  * A handle to use when requesting pulling more work from @BoundedQueueExecutor
  * via @BoundedQueueExecutor.pollWork
  */
 public interface BoundedQueueExecutorWorkHandle {
-  // Returns all work that are tracked by the handle
-  ImmutableList<Work> getWorkBatch();
+  // Returns all work that are tracked by the handle.
+  // Returned list cannot be modified. Copying the list is fine.
+  // Don't keep reference to the returned list after the processing exits the harness threads.
+  List<Work> getWorkBatch();
 }

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

@@ -82,4 +82,12 @@ public String getComputationId() {
   public Work.KeyGroup getKeyGroup() {
     return work().getKeyGroup();
   }
+
+  /**
+   * Returns true if multi-key batching is disabled for this work item (e.g. after a prior batch
+   * commit size validation failure).
+   */
+  public boolean isMultiKeyBatchingDisabled() {
+    return work().isMultiKeyBatchingDisabled();
+  }
 }

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

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.dataflow.worker.streaming;
+
+/**
+ * Thrown when a multi-key bundle exceeds commit size limits, triggering unbatching and local retry
+ * of individual work items.
+ */
+public final class MultiKeyCommitValidationException extends RuntimeException {
+  public MultiKeyCommitValidationException(String message) {
+    super(message);
+  }
+}

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

@@ -83,6 +83,10 @@ public final class Work implements RefreshableWork {
   private final long serializedWorkItemSize;
   private volatile TimedState currentState;
   private volatile boolean isFailed;
+  // If true, this work item will not be batched with other work items in a multi-key bundle.
+  // This is used to isolate work items that failed validation (e.g. commit size limit exceeded)
+  // so they can be retried individually and potentially truncated.
+  private volatile boolean disableMultiKeyBatching = false;
   private volatile String processingThreadName = "";
   private final AtomicReference<@Nullable AtomicBoolean> onFailureListener =
       new AtomicReference<>(null);
@@ -399,6 +403,22 @@ public boolean isFailed() {
     return isFailed;
   }
 
+  /**
+   * Sets whether multi-key batching should be disabled for this work item. When true, this work
+   * item will not be batched with other work items upon local retry.
+   */
+  public void setMultiKeyBatchingDisabled(boolean disableMultiKeyBatching) {
+    this.disableMultiKeyBatching = disableMultiKeyBatching;
+  }
+
+  /**
+   * Returns true if multi-key batching is disabled for this work item (e.g. after a prior batch
+   * commit size validation failure).
+   */
+  public boolean isMultiKeyBatchingDisabled() {
+    return disableMultiKeyBatching;
+  }
+
   boolean isStuckCommittingAt(Instant stuckCommitDeadline) {
     return currentState.state() == Work.State.COMMITTING
         && currentState.startTime().isBefore(stuckCommitDeadline);

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

@@ -22,6 +22,7 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -36,7 +37,6 @@
 import org.apache.beam.runners.dataflow.worker.streaming.Work;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor.Guard;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -306,8 +306,13 @@ public synchronized boolean isClosed() {
     }
 
     @Override
-    public synchronized ImmutableList<Work> getWorkBatch() {
-      return ImmutableList.copyOf(workBatch);
+    /*
+     * Returns an unmodifiable view over the underlying list.
+     * It is unsafe to use the returned list with concurrent calls to mutating methods
+     * like merge/close
+     */
+    public synchronized List<Work> getWorkBatch() {
+      return Collections.unmodifiableList(workBatch);
     }
 
     @VisibleForTesting

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

@@ -20,6 +20,7 @@
 import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
 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.checkArgument;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import java.util.AbstractQueue;
 import java.util.Collection;
@@ -67,9 +68,14 @@ static class Node {
     @Nullable Node prevKeyGroupNode;
     @Nullable Node nextKeyGroupNode;
 
+    private static boolean isMultiKeyBatchingDisabled(Runnable task) {
+      return !(task instanceof QueuedWork)
+          || ((QueuedWork) task).getWork().isMultiKeyBatchingDisabled();
+    }
+
     Node(Runnable task) {
       this.task = task;
-      if (task instanceof QueuedWork) {
+      if (!isMultiKeyBatchingDisabled(task)) {
         this.computationId = ((QueuedWork) task).getWork().getComputationId();
         this.keyGroup = ((QueuedWork) task).getWork().getKeyGroup();
       } else {
@@ -193,6 +199,10 @@ private void unlinkNode(Node node) {
       if (firstNode == keyGroupWorkList.tail) {
         return null;
       }
+
+      // MultiKeyBatchingDisabled items should not be in keyGroupWorkList
+      checkState(!Node.isMultiKeyBatchingDisabled(firstNode.task));
+
       unlinkNode(firstNode);
 
       return (QueuedWork) firstNode.task;

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java

@@ -242,7 +242,6 @@ private void processWork(
     long processingStartTimeNanos = System.nanoTime();
     StageInfo stageInfo = getStageInfo(computationState);
 
-    @Nullable List<Work> workBatch = null;
     try {
       if (work.isFailed()) {
         throw new WorkItemCancelledException(workItem.getShardingKey());
@@ -251,7 +250,7 @@ private void processWork(
       // Execute the user code for the Work batch.
       ExecuteWorkResult executeWorkResult =
           executeWork(work, stageInfo, computationState, handle, keyTransitionListener);
-      workBatch = executeWorkResult.workBatch();
+      List<Work> workBatch = handle.getWorkBatch();
       List<Windmill.WorkItemCommitRequest> workItemCommits = executeWorkResult.workItemCommits();
 
       commitFinalizer.cacheCommitFinalizers(executeWorkResult.finalizationCallbacks());
@@ -264,7 +263,7 @@ private void processWork(
       handleProcessWorkFailure(
           computationState, handle.getWorkBatch(), computationId, systemName, work, t);
     } finally {
-      List<Work> processedWorkBatch = workBatch != null ? workBatch : ImmutableList.of(work);
+      List<Work> processedWorkBatch = handle.getWorkBatch();
       // Update total processing time counters. Updating in finally clause ensures that
       // work items causing exceptions are also accounted in time spent.
       recordProcessingTime(stageInfo, processedWorkBatch, processingStartTimeNanos);
@@ -328,7 +327,6 @@ private ExecuteWorkResult executeWork(
       computationWorkExecutor.executeWork(
           work, workExecutor, handle, keyTransitionListener, onFailedWorkHandler);
 
-      List<Work> workBatch;
       List<Windmill.WorkItemCommitRequest> workItemCommits;
       Map<Long, Pair<Instant, Runnable>> finalizationCallbacks;
       long stateBytesRead;
@@ -338,9 +336,6 @@ private ExecuteWorkResult executeWork(
         }
         context.flushState();
 
-        // Retrieve executed works, work item commits, and accumulated callbacks from execution
-        // context
-        workBatch = context.getExecutedWorks();
         workItemCommits = context.getWorkItemCommits();
         finalizationCallbacks = context.getFinalizationCallbacks();
         stateBytesRead = context.getStateBytesRead();
@@ -351,8 +346,7 @@ private ExecuteWorkResult executeWork(
       computationState.releaseComputationWorkExecutor(computationWorkExecutor);
       computationWorkExecutor = null;
 
-      return ExecuteWorkResult.create(
-          workBatch, workItemCommits, finalizationCallbacks, stateBytesRead);
+      return ExecuteWorkResult.create(workItemCommits, finalizationCallbacks, stateBytesRead);
     } catch (Throwable t) {
       if (computationWorkExecutor != null) {
         // If processing failed due to a thrown exception, close the executionState. Do not
@@ -419,10 +413,6 @@ private void commitMultiKeyWorkBatch(
     }
     for (int i = 0; i < workBatch.size(); i++) {
       Windmill.WorkItemCommitRequest commit = workItemCommits.get(i);
-      // TODO: Retry on commit truncations
-      checkState(
-          !commit.getExceedsMaxWorkItemCommitBytes(),
-          "Commit truncation with multikey bundles not implemented");
       Work w = workBatch.get(i);
       multiKeyBuilder.addRequests(
           commit
@@ -523,16 +513,13 @@ private KeyTransitionListener createKeyTransitionListener() {
   @AutoValue
   abstract static class ExecuteWorkResult {
     static ExecuteWorkResult create(
-        List<Work> workBatch,
         List<Windmill.WorkItemCommitRequest> workItemCommits,
         Map<Long, Pair<Instant, Runnable>> finalizationCallbacks,
         long stateBytesRead) {
       return new AutoValue_StreamingWorkScheduler_ExecuteWorkResult(
-          workBatch, workItemCommits, finalizationCallbacks, stateBytesRead);
+          workItemCommits, finalizationCallbacks, stateBytesRead);
     }
 
-    abstract List<Work> workBatch();
-
     abstract List<Windmill.WorkItemCommitRequest> workItemCommits();
 
     // Map<finalizerId, Pair<callbackExpiration, callback>>

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java

@@ -25,6 +25,7 @@
 import org.apache.beam.runners.dataflow.worker.status.LastExceptionDataProvider;
 import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
 import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler;
+import org.apache.beam.runners.dataflow.worker.streaming.MultiKeyCommitValidationException;
 import org.apache.beam.runners.dataflow.worker.streaming.Work;
 import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor;
 import org.apache.beam.sdk.annotations.Internal;
@@ -162,6 +163,16 @@ private RetryEvaluation evaluateRetry(
     @Nullable final Throwable cause = t.getCause();
     Throwable parsedException = (t instanceof UserCodeException && cause != null) ? cause : t;
 
+    if (parsedException instanceof MultiKeyCommitValidationException) {
+      LOG.info(
+          "Execution of work for computation '{}' on sharding key '{}' for work token '{}' exceeded commit size limits. "
+              + "Work will be retried locally in smaller batches.",
+          computationId,
+          work.getWorkItem().getShardingKey(),
+          work.getWorkItem().getWorkToken());
+      return RetryEvaluation.RETRY_LOCALLY;
+    }
+
     LastExceptionDataProvider.reportException(parsedException);
     LOG.debug("Failed work: {}", work);
     Duration elapsedTimeSinceStart = new Duration(work.getStartTime(), clock.get());