diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/KeyCommitTooLargeException.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/KeyCommitTooLargeException.java index 4d6ae8a208c1..1069bee4f325 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/KeyCommitTooLargeException.java +++ b/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()); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java index e9b85d720d2b..45b9fbe93396 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java +++ b/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 clock, @@ -100,7 +101,8 @@ public StreamingWorkScheduler( StreamingCounters streamingCounters, ConcurrentMap 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 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 diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java index 23730bc57705..fc8f12c45cb3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java +++ b/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 kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); List 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 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 kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of());