KeyCommitTooLarge logging improvements

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

@@ -52,10 +52,12 @@
 import org.apache.beam.runners.dataflow.worker.counters.NameContext;
 import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
 import org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
+import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException;
 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;
 import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle;
+import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters;
 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;
@@ -77,6 +79,7 @@
 import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV1;
 import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV2;
 import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTimerData;
+import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.CoderException;
@@ -177,6 +180,9 @@ public class StreamingModeExecutionContext
   private final HotKeyLogger hotKeyLogger;
   private final boolean hotKeyLoggingEnabled;
   private final String stepName;
+  private final String systemName;
+  private final StreamingCounters streamingCounters;
+  private final FailureTracker failureTracker;
   private @Nullable Coder<?> keyCoder;
 
   // Key switch listener to delegate MDC logging context and thread name updates
@@ -212,6 +218,9 @@ public StreamingModeExecutionContext(
       HotKeyLogger hotKeyLogger,
       boolean hotKeyLoggingEnabled,
       String stepName,
+      String systemName,
+      StreamingCounters streamingCounters,
+      FailureTracker failureTracker,
       String sourceBytesProcessCounterName,
       SideInputStateFetcherFactory sideInputStateFetcherFactory) {
     super(
@@ -231,6 +240,9 @@ public StreamingModeExecutionContext(
     this.hotKeyLogger = checkNotNull(hotKeyLogger);
     this.hotKeyLoggingEnabled = hotKeyLoggingEnabled;
     this.stepName = checkNotNull(stepName);
+    this.systemName = checkNotNull(systemName);
+    this.streamingCounters = checkNotNull(streamingCounters);
+    this.failureTracker = checkNotNull(failureTracker);
     this.sourceBytesProcessCounterName = checkNotNull(sourceBytesProcessCounterName);
     this.sideInputStateFetcherFactory = sideInputStateFetcherFactory;
     StreamingGlobalConfig config = globalConfigHandle.getConfig();
@@ -669,6 +681,47 @@ private void flushStateInternal() {
 
     getOutputBuilder()
         .setSourceBytesProcessed(computeSourceBytesProcessed(sourceBytesProcessCounterName));
+
+    validateCommitRequestSize();
+  }
+
+  private void validateCommitRequestSize() {
+    Windmill.WorkItemCommitRequest.Builder currentBuilder = getOutputBuilder();
+    Work currentWork = getWork();
+    long byteLimit = operationalLimits.getMaxWorkItemCommitBytes();
+    Windmill.WorkItemCommitRequest commitRequest = currentBuilder.build();
+    int commitSize = commitRequest.getSerializedSize();
+
+    // Detect overflow of integer serialized size or if the byte limit was exceeded.
+    // Commit is too large if overflow has occurred or the commitSize has exceeded the allowed
+    // commit byte limit.
+    int estimatedCommitSize = commitSize < 0 ? Integer.MAX_VALUE : commitSize;
+    streamingCounters.windmillMaxObservedWorkItemCommitBytes().addValue(estimatedCommitSize);
+    if (commitSize >= 0 && commitSize < byteLimit) {
+      return;
+    }
+
+    KeyCommitTooLargeException e =
+        KeyCommitTooLargeException.causedBy(
+            systemName, byteLimit, commitRequest, key, hotKeyLoggingEnabled);
+    failureTracker.trackFailure(systemName, currentWork.getWorkItem(), e);
+    LOG.error("{}", e.toString());
+
+    // Drop the current request in favor of a new, minimal one requesting truncation.
+    // Messages, timers, counters, and other commit content will not be used by the service
+    // so, we're purposefully dropping them here
+    Windmill.WorkItemCommitRequest.Builder truncationBuilder =
+        buildWorkItemTruncationRequestBuilder(currentWork, estimatedCommitSize);
+    currentBuilder.clear();
+    currentBuilder.mergeFrom(truncationBuilder.build());
+  }
+
+  private Windmill.WorkItemCommitRequest.Builder buildWorkItemTruncationRequestBuilder(
+      Work work, int estimatedCommitSize) {
+    Windmill.WorkItemCommitRequest.Builder outputBuilder = createOutputBuilder(work);
+    outputBuilder.setExceedsMaxWorkItemCommitBytes(true);
+    outputBuilder.setEstimatedWorkItemCommitBytes(estimatedCommitSize);
+    return outputBuilder;
   }
 
   private final long computeSourceBytesProcessed(String sourceBytesCounterName) {

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

@@ -18,16 +18,25 @@
 package org.apache.beam.runners.dataflow.worker.streaming;
 
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 public final class KeyCommitTooLargeException extends Exception {
 
   public static KeyCommitTooLargeException causedBy(
-      String computationId, long byteLimit, Windmill.WorkItemCommitRequest request) {
+      String stageName,
+      long byteLimit,
+      Windmill.WorkItemCommitRequest request,
+      @Nullable Object decodedKey,
+      boolean hotKeyLoggingEnabled) {
     StringBuilder message = new StringBuilder();
     message.append("Commit request for stage ");
-    message.append(computationId);
+    message.append(stageName);
     message.append(" and sharding key ");
-    message.append(request.getShardingKey());
+    message.append(Long.toUnsignedString(request.getShardingKey()));
+    if (decodedKey != null && hotKeyLoggingEnabled) {
+      message.append(" and key ");
+      message.append(decodedKey);
+    }
     if (request.getSerializedSize() > 0) {
       message.append(
           " has size "
@@ -38,9 +47,8 @@ public static KeyCommitTooLargeException causedBy(
       message.append(" is larger than 2GB and cannot be processed");
     }
     message.append(
-        ". This may be caused by grouping a very "
-            + "large amount of data in a single window without using Combine,"
-            + " or by producing a large amount of data from a single input element."
+        ". This may be caused by grouping a very large amount of data in a single window without"
+            + " using Combine, or by producing a large amount of data from a single input element."
             + " See https://cloud.google.com/dataflow/docs/guides/common-errors#key-commit-too-large-exception.");
     return new KeyCommitTooLargeException(message.toString());
   }

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

@@ -49,11 +49,13 @@
 import org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor;
 import org.apache.beam.runners.dataflow.worker.streaming.StageInfo;
 import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle;
+import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters;
 import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory;
 import org.apache.beam.runners.dataflow.worker.util.common.worker.MapTaskExecutor;
 import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter;
 import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation;
 import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
+import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.fn.IdGenerator;
@@ -84,6 +86,8 @@ final class ComputationWorkExecutorFactory {
   private final SinkRegistry sinkRegistry;
   private final DataflowExecutionStateSampler sampler;
   private final CounterSet pendingDeltaCounters;
+  private final StreamingCounters streamingCounters;
+  private final FailureTracker failureTracker;
 
   /**
    * Function which converts map tasks to their network representation for execution.
@@ -108,7 +112,8 @@ final class ComputationWorkExecutorFactory {
       ReaderCache readerCache,
       Function<String, WindmillStateCache.ForComputation> stateCacheFactory,
       DataflowExecutionStateSampler sampler,
-      CounterSet pendingDeltaCounters,
+      StreamingCounters streamingCounters,
+      FailureTracker failureTracker,
       IdGenerator idGenerator,
       StreamingGlobalConfigHandle globalConfigHandle,
       HotKeyLogger hotKeyLogger,
@@ -122,7 +127,9 @@ final class ComputationWorkExecutorFactory {
     this.readerRegistry = ReaderRegistry.defaultRegistry();
     this.sinkRegistry = SinkRegistry.defaultRegistry();
     this.sampler = sampler;
-    this.pendingDeltaCounters = pendingDeltaCounters;
+    this.streamingCounters = streamingCounters;
+    this.failureTracker = failureTracker;
+    this.pendingDeltaCounters = streamingCounters.pendingDeltaCounters();
     this.mapTaskToNetwork = new MapTaskToNetworkFunction(idGenerator);
     this.maxSinkBytes =
         hasExperiment(options, DISABLE_SINK_BYTE_LIMIT_EXPERIMENT)
@@ -286,6 +293,9 @@ private StreamingModeExecutionContext createExecutionContext(
         hotKeyLogger,
         hotKeyLoggingEnabled,
         stepName,
+        stageInfo.systemName(),
+        streamingCounters,
+        failureTracker,
         computationState.sourceBytesProcessCounterName(),
         sideInputStateFetcherFactory);
   }

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

@@ -43,7 +43,6 @@
 import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
 import org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor;
 import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
-import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException;
 import org.apache.beam.runners.dataflow.worker.streaming.StageInfo;
 import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
 import org.apache.beam.runners.dataflow.worker.streaming.Work;
@@ -81,36 +80,30 @@ public class StreamingWorkScheduler {
 
   private final Supplier<Instant> clock;
   private final ComputationWorkExecutorFactory computationWorkExecutorFactory;
-  private final FailureTracker failureTracker;
   private final WorkFailureProcessor workFailureProcessor;
   private final StreamingCommitFinalizer commitFinalizer;
   private final StreamingCounters streamingCounters;
   private final ConcurrentMap<String, StageInfo> stageInfoMap;
   private final DataflowExecutionStateSampler sampler;
-  private final StreamingGlobalConfigHandle globalConfigHandle;
   private final BoundedQueueExecutor workExecutor;
 
   public StreamingWorkScheduler(
       Supplier<Instant> clock,
       BoundedQueueExecutor workExecutor,
       ComputationWorkExecutorFactory computationWorkExecutorFactory,
-      FailureTracker failureTracker,
       WorkFailureProcessor workFailureProcessor,
       StreamingCommitFinalizer commitFinalizer,
       StreamingCounters streamingCounters,
       ConcurrentMap<String, StageInfo> stageInfoMap,
-      DataflowExecutionStateSampler sampler,
-      StreamingGlobalConfigHandle globalConfigHandle) {
+      DataflowExecutionStateSampler sampler) {
     this.clock = clock;
     this.workExecutor = workExecutor;
     this.computationWorkExecutorFactory = computationWorkExecutorFactory;
-    this.failureTracker = failureTracker;
     this.workFailureProcessor = workFailureProcessor;
     this.commitFinalizer = commitFinalizer;
     this.streamingCounters = streamingCounters;
     this.stageInfoMap = stageInfoMap;
     this.sampler = sampler;
-    this.globalConfigHandle = globalConfigHandle;
   }
 
   public static StreamingWorkScheduler create(
@@ -139,7 +132,8 @@ public static StreamingWorkScheduler create(
             readerCache,
             stateCacheFactory,
             sampler,
-            streamingCounters.pendingDeltaCounters(),
+            streamingCounters,
+            failureTracker,
             idGenerator,
             globalConfigHandle,
             hotKeyLogger,
@@ -149,13 +143,11 @@ public static StreamingWorkScheduler create(
         clock,
         workExecutor,
         computationWorkExecutorFactory,
-        failureTracker,
         workFailureProcessor,
         StreamingCommitFinalizer.create(workExecutor, commitFinalizerCleanupExecutor),
         streamingCounters,
         stageInfoMap,
-        sampler,
-        globalConfigHandle);
+        sampler);
   }
 
   private static long computeShuffleBytesRead(Windmill.WorkItem workItem) {
@@ -175,14 +167,6 @@ private static Windmill.WorkItemCommitRequest.Builder initializeOutputBuilder(
         .setCacheToken(workItem.getCacheToken());
   }
 
-  private static Windmill.WorkItemCommitRequest buildWorkItemTruncationRequest(
-      ByteString key, Windmill.WorkItem workItem, int estimatedCommitSize) {
-    Windmill.WorkItemCommitRequest.Builder outputBuilder = initializeOutputBuilder(key, workItem);
-    outputBuilder.setExceedsMaxWorkItemCommitBytes(true);
-    outputBuilder.setEstimatedWorkItemCommitBytes(estimatedCommitSize);
-    return outputBuilder.build();
-  }
-
   /** Sets the stage name and workId of the Thread executing the {@link Work} for logging. */
   private static void setUpWorkLoggingContext(String workLatencyTrackingId, String computationId) {
     setLoggingContextWorkId(workLatencyTrackingId);
@@ -305,33 +289,6 @@ private void processWork(
     }
   }
 
-  private Windmill.WorkItemCommitRequest validateCommitRequestSize(
-      Windmill.WorkItemCommitRequest commitRequest,
-      String computationId,
-      Windmill.WorkItem workItem) {
-    long byteLimit = globalConfigHandle.getConfig().operationalLimits().getMaxWorkItemCommitBytes();
-    int commitSize = commitRequest.getSerializedSize();
-    int estimatedCommitSize = commitSize < 0 ? Integer.MAX_VALUE : commitSize;
-
-    // Detect overflow of integer serialized size or if the byte limit was exceeded.
-    // Commit is too large if overflow has occurred or the commitSize has exceeded the allowed
-    // commit byte limit.
-    streamingCounters.windmillMaxObservedWorkItemCommitBytes().addValue(estimatedCommitSize);
-    if (commitSize >= 0 && commitSize < byteLimit) {
-      return commitRequest;
-    }
-
-    KeyCommitTooLargeException e =
-        KeyCommitTooLargeException.causedBy(computationId, byteLimit, commitRequest);
-    failureTracker.trackFailure(computationId, workItem, e);
-    LOG.error("{}", e.toString());
-
-    // Drop the current request in favor of a new, minimal one requesting truncation.
-    // Messages, timers, counters, and other commit content will not be used by the service
-    // so, we're purposefully dropping them here
-    return buildWorkItemTruncationRequest(workItem.getKey(), workItem, estimatedCommitSize);
-  }
-
   private void recordProcessingStats(
       List<Work> workBatch,
       List<Windmill.WorkItemCommitRequest> workItemCommits,
@@ -448,17 +405,13 @@ private void commitWorkBatch(
 
   private void commitSingleKeyWork(
       ComputationState computationState, Work work, Windmill.WorkItemCommitRequest commitRequest) {
-    // Validate the commit request, possibly requesting truncation if the commitSize is too large.
-    Windmill.WorkItemCommitRequest validatedCommitRequest =
-        validateCommitRequestSize(
-            commitRequest, computationState.getComputationId(), work.getWorkItem());
     work.setState(Work.State.COMMIT_QUEUED);
-    validatedCommitRequest =
-        validatedCommitRequest
+    Windmill.WorkItemCommitRequest commitRequestWithAttributions =
+        commitRequest
             .toBuilder()
             .addAllPerWorkItemLatencyAttributions(work.getLatencyAttributions(sampler))
             .build();
-    work.queueCommit(validatedCommitRequest, computationState);
+    work.queueCommit(commitRequestWithAttributions, computationState);
   }
 
   private void recordProcessingTime(

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

@@ -1290,8 +1290,9 @@ public void testKeyTokenInvalidException() throws Exception {
     worker.stop();
   }
 
-  @Test
-  public void testKeyCommitTooLargeException() throws Exception {
+  private void runKeyCommitTooLargeExceptionTest(
+      StreamingDataflowWorkerTestParams.Builder workerParams, boolean expectKeyInErrorMessage)
+      throws Exception {
     KvCoder<String, String> kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of());
 
     List<ParallelInstruction> instructions =
@@ -1304,7 +1305,7 @@ public void testKeyCommitTooLargeException() throws Exception {
 
     StreamingDataflowWorker worker =
         makeWorker(
-            defaultWorkerParams()
+            workerParams
                 .setInstructions(instructions)
                 .setStreamingGlobalConfig(
                     StreamingGlobalConfig.builder()
@@ -1336,7 +1337,6 @@ public void testKeyCommitTooLargeException() throws Exception {
                 1, "large_key", DEFAULT_SHARDING_KEY, largeCommit.getEstimatedWorkItemCommitBytes())
             .build(),
         removeDynamicFields(largeCommit));
-
     // Check this explicitly since the estimated commit bytes weren't actually
     // checked against an expected value in the previous step
     assertTrue(largeCommit.getEstimatedWorkItemCommitBytes() > 1000);
@@ -1360,12 +1360,35 @@ public void testKeyCommitTooLargeException() throws Exception {
         foundErrors = true;
         String errorMessage = status.getErrors().get(0).getMessage();
         assertThat(errorMessage, Matchers.containsString("KeyCommitTooLargeException"));
+        assertThat(errorMessage, Matchers.containsString("Commit request for stage computation"));
+        if (expectKeyInErrorMessage) {
+          assertThat(errorMessage, Matchers.containsString("and key large_key"));
+        } else {
+          assertThat(errorMessage, Matchers.not(Matchers.containsString("and key large_key")));
+        }
       }
     }
     assertTrue(foundErrors);
     worker.stop();
   }
 
+  @Test
+  public void testKeyCommitTooLargeException() throws Exception {
+    runKeyCommitTooLargeExceptionTest(defaultWorkerParams(), /* expectKeyInErrorMessage= */ false);
+  }
+
+  @Test
+  public void testKeyCommitTooLargeException_withHotKeyLoggingEnabled() throws Exception {
+    runKeyCommitTooLargeExceptionTest(
+        defaultWorkerParams("--hotKeyLoggingEnabled=true"), /* expectKeyInErrorMessage= */ true);
+  }
+
+  @Test
+  public void testKeyCommitTooLargeException_withHotKeyLoggingDisabled() throws Exception {
+    runKeyCommitTooLargeExceptionTest(
+        defaultWorkerParams("--hotKeyLoggingEnabled=false"), /* expectKeyInErrorMessage= */ false);
+  }
+
   @Test
   public void testOutputKeyTooLargeException() throws Exception {
     KvCoder<String, String> kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of());