KeyCommitTooLarge logging improvements

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

@@ -18,16 +18,29 @@
 package org.apache.beam.runners.dataflow.worker.streaming;
 
 import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
+import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat;
 
 public final class KeyCommitTooLargeException extends Exception {
 
   public static KeyCommitTooLargeException causedBy(
-      String computationId, long byteLimit, Windmill.WorkItemCommitRequest request) {
+      String stageName, long byteLimit, Windmill.WorkItemCommitRequest request) {
+    return causedBy(stageName, byteLimit, request, false);
+  }
+
+  public static KeyCommitTooLargeException causedBy(
+      String stageName,
+      long byteLimit,
+      Windmill.WorkItemCommitRequest request,
+      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 (hotKeyLoggingEnabled && !request.getKey().isEmpty()) {
+      message.append(" and key ");
+      message.append(TextFormat.escapeBytes(request.getKey()));
+    }
     if (request.getSerializedSize() > 0) {
       message.append(
           " has size "
@@ -38,9 +51,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/StreamingWorkScheduler.java

@@ -89,6 +89,7 @@ public class StreamingWorkScheduler {
   private final DataflowExecutionStateSampler sampler;
   private final StreamingGlobalConfigHandle globalConfigHandle;
   private final BoundedQueueExecutor workExecutor;
+  private final boolean hotKeyLoggingEnabled;
 
   public StreamingWorkScheduler(
       Supplier<Instant> clock,
@@ -100,7 +101,8 @@ public StreamingWorkScheduler(
       StreamingCounters streamingCounters,
       ConcurrentMap<String, StageInfo> stageInfoMap,
       DataflowExecutionStateSampler sampler,
-      StreamingGlobalConfigHandle globalConfigHandle) {
+      StreamingGlobalConfigHandle globalConfigHandle,
+      boolean hotKeyLoggingEnabled) {
     this.clock = clock;
     this.workExecutor = workExecutor;
     this.computationWorkExecutorFactory = computationWorkExecutorFactory;
@@ -111,6 +113,7 @@ public StreamingWorkScheduler(
     this.stageInfoMap = stageInfoMap;
     this.sampler = sampler;
     this.globalConfigHandle = globalConfigHandle;
+    this.hotKeyLoggingEnabled = hotKeyLoggingEnabled;
   }
 
   public static StreamingWorkScheduler create(
@@ -145,6 +148,12 @@ public static StreamingWorkScheduler create(
             hotKeyLogger,
             sideInputStateFetcherFactory);
 
+    List<String> experiments = options.getExperiments();
+    boolean hotKeyLoggingEnabled =
+        options.isHotKeyLoggingEnabled()
+            || (experiments != null
+                && experiments.stream().anyMatch("enable_hot_key_logging"::equalsIgnoreCase));
+
     return new StreamingWorkScheduler(
         clock,
         workExecutor,
@@ -155,7 +164,8 @@ public static StreamingWorkScheduler create(
         streamingCounters,
         stageInfoMap,
         sampler,
-        globalConfigHandle);
+        globalConfigHandle,
+        hotKeyLoggingEnabled);
   }
 
   private static long computeShuffleBytesRead(Windmill.WorkItem workItem) {
@@ -306,9 +316,7 @@ private void processWork(
   }
 
   private Windmill.WorkItemCommitRequest validateCommitRequestSize(
-      Windmill.WorkItemCommitRequest commitRequest,
-      String computationId,
-      Windmill.WorkItem workItem) {
+      Windmill.WorkItemCommitRequest commitRequest, String stageName, Windmill.WorkItem workItem) {
     long byteLimit = globalConfigHandle.getConfig().operationalLimits().getMaxWorkItemCommitBytes();
     int commitSize = commitRequest.getSerializedSize();
     int estimatedCommitSize = commitSize < 0 ? Integer.MAX_VALUE : commitSize;
@@ -322,8 +330,9 @@ private Windmill.WorkItemCommitRequest validateCommitRequestSize(
     }
 
     KeyCommitTooLargeException e =
-        KeyCommitTooLargeException.causedBy(computationId, byteLimit, commitRequest);
-    failureTracker.trackFailure(computationId, workItem, e);
+        KeyCommitTooLargeException.causedBy(
+            stageName, byteLimit, commitRequest, hotKeyLoggingEnabled);
+    failureTracker.trackFailure(stageName, workItem, e);
     LOG.error("{}", e.toString());
 
     // Drop the current request in favor of a new, minimal one requesting truncation.
@@ -451,7 +460,7 @@ private void commitSingleKeyWork(
     // Validate the commit request, possibly requesting truncation if the commitSize is too large.
     Windmill.WorkItemCommitRequest validatedCommitRequest =
         validateCommitRequestSize(
-            commitRequest, computationState.getComputationId(), work.getWorkItem());
+            commitRequest, computationState.getMapTask().getSystemName(), work.getWorkItem());
     work.setState(Work.State.COMMIT_QUEUED);
     validatedCommitRequest =
         validatedCommitRequest

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()
@@ -1337,18 +1338,14 @@ public void testKeyCommitTooLargeException() throws Exception {
             .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);
 
-    // Spam worker updates a few times.
     int maxTries = 10;
     while (--maxTries > 0) {
       worker.reportPeriodicWorkerUpdatesForTest();
       Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
     }
 
-    // We should see an exception reported for the large commit but not the small one.
     ArgumentCaptor<WorkItemStatus> workItemStatusCaptor =
         ArgumentCaptor.forClass(WorkItemStatus.class);
     verify(mockWorkUnitClient, atLeast(2)).reportWorkItemStatus(workItemStatusCaptor.capture());
@@ -1360,12 +1357,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());