Skip to content

Commit 5e49938

Browse files
rwigglesgemini-code-assist[bot]dependabot[bot]AbacnaIbrahiim
authored
KeyCommitTooLarge logging improvements (#39316)
* Support logging the key during a KeyCommitTooLarge if EnableHotKeyLogging is turned on. Also switch to logging the dfe name instead of the computation name and fix an overflow in logging the sharding key * Update runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/KeyCommitTooLargeException.java Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * gemini review responses * Fix imports * one more * one more attempt * one more attempt * formatting fix * Bump actions/checkout from 6 to 7 (#39335) Bumps [actions/checkout](https://github.com/actions/checkout) from 6 to 7. - [Release notes](https://github.com/actions/checkout/releases) - [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md) - [Commits](actions/checkout@v6...v7) --- updated-dependencies: - dependency-name: actions/checkout dependency-version: '7' dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * Bump cloud.google.com/go/datastore from 1.24.0 to 1.25.0 in /sdks (#39333) Bumps [cloud.google.com/go/datastore](https://github.com/googleapis/google-cloud-go) from 1.24.0 to 1.25.0. - [Release notes](https://github.com/googleapis/google-cloud-go/releases) - [Changelog](https://github.com/googleapis/google-cloud-go/blob/main/documentai/CHANGES.md) - [Commits](googleapis/google-cloud-go@kms/v1.24.0...kms/v1.25.0) --- updated-dependencies: - dependency-name: cloud.google.com/go/datastore dependency-version: 1.25.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * Bump github.com/aws/aws-sdk-go-v2/feature/s3/manager in /sdks (#39334) Bumps [github.com/aws/aws-sdk-go-v2/feature/s3/manager](https://github.com/aws/aws-sdk-go-v2) from 1.22.32 to 1.22.33. - [Release notes](https://github.com/aws/aws-sdk-go-v2/releases) - [Commits](aws/aws-sdk-go-v2@feature/s3/manager/v1.22.32...feature/s3/manager/v1.22.33) --- updated-dependencies: - dependency-name: github.com/aws/aws-sdk-go-v2/feature/s3/manager dependency-version: 1.22.33 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * Fix DataflowV1 test failure by fixing getSimpleName access (#39330) * add jdkAddOpenModules for Dataflow V2 (#39306) * Add Spark JVM --add-opens for (Nexmark, TPC-DS, PortableJar) (#39337) * add Spark job-server --add-opens (#39339) * fix: add retries and query parameter encoding for GitHub API requests (closes #39188) * fix: add Apache license header to test_sending.py * Move validation to StreamingModeExecutionContext * remove unused import * respond to comments * bugfix * Update runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java Co-authored-by: Arun Pandian <arunpandianp@gmail.com> * respond to comments --------- Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Yi Hu <yathu@google.com> Co-authored-by: Abdelrahman Ibrahim <abdoibrahim1017@gmail.com> Co-authored-by: raman118 <rcho38706@gmail.com> Co-authored-by: Arun Pandian <arunpandianp@gmail.com>
1 parent adbb11d commit 5e49938

7 files changed

Lines changed: 126 additions & 66 deletions

File tree

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,10 +52,12 @@
5252
import org.apache.beam.runners.dataflow.worker.counters.NameContext;
5353
import org.apache.beam.runners.dataflow.worker.profiler.ScopedProfiler.ProfileScope;
5454
import org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWorkHandle;
55+
import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException;
5556
import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
5657
import org.apache.beam.runners.dataflow.worker.streaming.Work;
5758
import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig;
5859
import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle;
60+
import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters;
5961
import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInput;
6062
import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputState;
6163
import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcher;
@@ -77,6 +79,7 @@
7779
import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV1;
7880
import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV2;
7981
import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTimerData;
82+
import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker;
8083
import org.apache.beam.sdk.annotations.Internal;
8184
import org.apache.beam.sdk.coders.Coder;
8285
import org.apache.beam.sdk.coders.CoderException;
@@ -177,6 +180,9 @@ public class StreamingModeExecutionContext
177180
private final HotKeyLogger hotKeyLogger;
178181
private final boolean hotKeyLoggingEnabled;
179182
private final String stepName;
183+
private final String systemName;
184+
private final StreamingCounters streamingCounters;
185+
private final FailureTracker failureTracker;
180186
private @Nullable Coder<?> keyCoder;
181187

182188
// Key switch listener to delegate MDC logging context and thread name updates
@@ -212,6 +218,9 @@ public StreamingModeExecutionContext(
212218
HotKeyLogger hotKeyLogger,
213219
boolean hotKeyLoggingEnabled,
214220
String stepName,
221+
String systemName,
222+
StreamingCounters streamingCounters,
223+
FailureTracker failureTracker,
215224
String sourceBytesProcessCounterName,
216225
SideInputStateFetcherFactory sideInputStateFetcherFactory) {
217226
super(
@@ -231,6 +240,9 @@ public StreamingModeExecutionContext(
231240
this.hotKeyLogger = checkNotNull(hotKeyLogger);
232241
this.hotKeyLoggingEnabled = hotKeyLoggingEnabled;
233242
this.stepName = checkNotNull(stepName);
243+
this.systemName = checkNotNull(systemName);
244+
this.streamingCounters = checkNotNull(streamingCounters);
245+
this.failureTracker = checkNotNull(failureTracker);
234246
this.sourceBytesProcessCounterName = checkNotNull(sourceBytesProcessCounterName);
235247
this.sideInputStateFetcherFactory = sideInputStateFetcherFactory;
236248
StreamingGlobalConfig config = globalConfigHandle.getConfig();
@@ -669,6 +681,47 @@ private void flushStateInternal() {
669681

670682
getOutputBuilder()
671683
.setSourceBytesProcessed(computeSourceBytesProcessed(sourceBytesProcessCounterName));
684+
685+
validateCommitRequestSize();
686+
}
687+
688+
private void validateCommitRequestSize() {
689+
Windmill.WorkItemCommitRequest.Builder currentBuilder = getOutputBuilder();
690+
Work currentWork = getWork();
691+
long byteLimit = operationalLimits.getMaxWorkItemCommitBytes();
692+
Windmill.WorkItemCommitRequest commitRequest = currentBuilder.build();
693+
int commitSize = commitRequest.getSerializedSize();
694+
695+
// Detect overflow of integer serialized size or if the byte limit was exceeded.
696+
// Commit is too large if overflow has occurred or the commitSize has exceeded the allowed
697+
// commit byte limit.
698+
int estimatedCommitSize = commitSize < 0 ? Integer.MAX_VALUE : commitSize;
699+
streamingCounters.windmillMaxObservedWorkItemCommitBytes().addValue(estimatedCommitSize);
700+
if (commitSize >= 0 && commitSize < byteLimit) {
701+
return;
702+
}
703+
704+
KeyCommitTooLargeException e =
705+
KeyCommitTooLargeException.causedBy(
706+
systemName, byteLimit, commitRequest, key, hotKeyLoggingEnabled);
707+
failureTracker.trackFailure(systemName, currentWork.getWorkItem(), e);
708+
LOG.error("{}", e.toString());
709+
710+
// Drop the current request in favor of a new, minimal one requesting truncation.
711+
// Messages, timers, counters, and other commit content will not be used by the service
712+
// so, we're purposefully dropping them here
713+
Windmill.WorkItemCommitRequest.Builder truncationBuilder =
714+
buildWorkItemTruncationRequestBuilder(currentWork, estimatedCommitSize);
715+
currentBuilder.clear();
716+
currentBuilder.mergeFrom(truncationBuilder.build());
717+
}
718+
719+
private Windmill.WorkItemCommitRequest.Builder buildWorkItemTruncationRequestBuilder(
720+
Work work, int estimatedCommitSize) {
721+
Windmill.WorkItemCommitRequest.Builder outputBuilder = createOutputBuilder(work);
722+
outputBuilder.setExceedsMaxWorkItemCommitBytes(true);
723+
outputBuilder.setEstimatedWorkItemCommitBytes(estimatedCommitSize);
724+
return outputBuilder;
672725
}
673726

674727
private final long computeSourceBytesProcessed(String sourceBytesCounterName) {

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

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,16 +18,25 @@
1818
package org.apache.beam.runners.dataflow.worker.streaming;
1919

2020
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
21+
import org.checkerframework.checker.nullness.qual.Nullable;
2122

2223
public final class KeyCommitTooLargeException extends Exception {
2324

2425
public static KeyCommitTooLargeException causedBy(
25-
String computationId, long byteLimit, Windmill.WorkItemCommitRequest request) {
26+
String stageName,
27+
long byteLimit,
28+
Windmill.WorkItemCommitRequest request,
29+
@Nullable Object decodedKey,
30+
boolean hotKeyLoggingEnabled) {
2631
StringBuilder message = new StringBuilder();
2732
message.append("Commit request for stage ");
28-
message.append(computationId);
33+
message.append(stageName);
2934
message.append(" and sharding key ");
30-
message.append(request.getShardingKey());
35+
message.append(Long.toUnsignedString(request.getShardingKey()));
36+
if (decodedKey != null && hotKeyLoggingEnabled) {
37+
message.append(" and key ");
38+
message.append(decodedKey);
39+
}
3140
if (request.getSerializedSize() > 0) {
3241
message.append(
3342
" has size "
@@ -38,9 +47,8 @@ public static KeyCommitTooLargeException causedBy(
3847
message.append(" is larger than 2GB and cannot be processed");
3948
}
4049
message.append(
41-
". This may be caused by grouping a very "
42-
+ "large amount of data in a single window without using Combine,"
43-
+ " or by producing a large amount of data from a single input element."
50+
". This may be caused by grouping a very large amount of data in a single window without"
51+
+ " using Combine, or by producing a large amount of data from a single input element."
4452
+ " See https://cloud.google.com/dataflow/docs/guides/common-errors#key-commit-too-large-exception.");
4553
return new KeyCommitTooLargeException(message.toString());
4654
}

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,11 +49,13 @@
4949
import org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor;
5050
import org.apache.beam.runners.dataflow.worker.streaming.StageInfo;
5151
import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfigHandle;
52+
import org.apache.beam.runners.dataflow.worker.streaming.harness.StreamingCounters;
5253
import org.apache.beam.runners.dataflow.worker.streaming.sideinput.SideInputStateFetcherFactory;
5354
import org.apache.beam.runners.dataflow.worker.util.common.worker.MapTaskExecutor;
5455
import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputObjectAndByteCounter;
5556
import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation;
5657
import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
58+
import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker;
5759
import org.apache.beam.sdk.coders.Coder;
5860
import org.apache.beam.sdk.coders.KvCoder;
5961
import org.apache.beam.sdk.fn.IdGenerator;
@@ -84,6 +86,8 @@ final class ComputationWorkExecutorFactory {
8486
private final SinkRegistry sinkRegistry;
8587
private final DataflowExecutionStateSampler sampler;
8688
private final CounterSet pendingDeltaCounters;
89+
private final StreamingCounters streamingCounters;
90+
private final FailureTracker failureTracker;
8791

8892
/**
8993
* Function which converts map tasks to their network representation for execution.
@@ -108,7 +112,8 @@ final class ComputationWorkExecutorFactory {
108112
ReaderCache readerCache,
109113
Function<String, WindmillStateCache.ForComputation> stateCacheFactory,
110114
DataflowExecutionStateSampler sampler,
111-
CounterSet pendingDeltaCounters,
115+
StreamingCounters streamingCounters,
116+
FailureTracker failureTracker,
112117
IdGenerator idGenerator,
113118
StreamingGlobalConfigHandle globalConfigHandle,
114119
HotKeyLogger hotKeyLogger,
@@ -122,7 +127,9 @@ final class ComputationWorkExecutorFactory {
122127
this.readerRegistry = ReaderRegistry.defaultRegistry();
123128
this.sinkRegistry = SinkRegistry.defaultRegistry();
124129
this.sampler = sampler;
125-
this.pendingDeltaCounters = pendingDeltaCounters;
130+
this.streamingCounters = streamingCounters;
131+
this.failureTracker = failureTracker;
132+
this.pendingDeltaCounters = streamingCounters.pendingDeltaCounters();
126133
this.mapTaskToNetwork = new MapTaskToNetworkFunction(idGenerator);
127134
this.maxSinkBytes =
128135
hasExperiment(options, DISABLE_SINK_BYTE_LIMIT_EXPERIMENT)
@@ -286,6 +293,9 @@ private StreamingModeExecutionContext createExecutionContext(
286293
hotKeyLogger,
287294
hotKeyLoggingEnabled,
288295
stepName,
296+
stageInfo.systemName(),
297+
streamingCounters,
298+
failureTracker,
289299
computationState.sourceBytesProcessCounterName(),
290300
sideInputStateFetcherFactory);
291301
}

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

Lines changed: 7 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,6 @@
4343
import org.apache.beam.runners.dataflow.worker.streaming.ComputationState;
4444
import org.apache.beam.runners.dataflow.worker.streaming.ComputationWorkExecutor;
4545
import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork;
46-
import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException;
4746
import org.apache.beam.runners.dataflow.worker.streaming.StageInfo;
4847
import org.apache.beam.runners.dataflow.worker.streaming.Watermarks;
4948
import org.apache.beam.runners.dataflow.worker.streaming.Work;
@@ -81,36 +80,30 @@ public class StreamingWorkScheduler {
8180

8281
private final Supplier<Instant> clock;
8382
private final ComputationWorkExecutorFactory computationWorkExecutorFactory;
84-
private final FailureTracker failureTracker;
8583
private final WorkFailureProcessor workFailureProcessor;
8684
private final StreamingCommitFinalizer commitFinalizer;
8785
private final StreamingCounters streamingCounters;
8886
private final ConcurrentMap<String, StageInfo> stageInfoMap;
8987
private final DataflowExecutionStateSampler sampler;
90-
private final StreamingGlobalConfigHandle globalConfigHandle;
9188
private final BoundedQueueExecutor workExecutor;
9289

9390
public StreamingWorkScheduler(
9491
Supplier<Instant> clock,
9592
BoundedQueueExecutor workExecutor,
9693
ComputationWorkExecutorFactory computationWorkExecutorFactory,
97-
FailureTracker failureTracker,
9894
WorkFailureProcessor workFailureProcessor,
9995
StreamingCommitFinalizer commitFinalizer,
10096
StreamingCounters streamingCounters,
10197
ConcurrentMap<String, StageInfo> stageInfoMap,
102-
DataflowExecutionStateSampler sampler,
103-
StreamingGlobalConfigHandle globalConfigHandle) {
98+
DataflowExecutionStateSampler sampler) {
10499
this.clock = clock;
105100
this.workExecutor = workExecutor;
106101
this.computationWorkExecutorFactory = computationWorkExecutorFactory;
107-
this.failureTracker = failureTracker;
108102
this.workFailureProcessor = workFailureProcessor;
109103
this.commitFinalizer = commitFinalizer;
110104
this.streamingCounters = streamingCounters;
111105
this.stageInfoMap = stageInfoMap;
112106
this.sampler = sampler;
113-
this.globalConfigHandle = globalConfigHandle;
114107
}
115108

116109
public static StreamingWorkScheduler create(
@@ -139,7 +132,8 @@ public static StreamingWorkScheduler create(
139132
readerCache,
140133
stateCacheFactory,
141134
sampler,
142-
streamingCounters.pendingDeltaCounters(),
135+
streamingCounters,
136+
failureTracker,
143137
idGenerator,
144138
globalConfigHandle,
145139
hotKeyLogger,
@@ -149,13 +143,11 @@ public static StreamingWorkScheduler create(
149143
clock,
150144
workExecutor,
151145
computationWorkExecutorFactory,
152-
failureTracker,
153146
workFailureProcessor,
154147
StreamingCommitFinalizer.create(workExecutor, commitFinalizerCleanupExecutor),
155148
streamingCounters,
156149
stageInfoMap,
157-
sampler,
158-
globalConfigHandle);
150+
sampler);
159151
}
160152

161153
private static long computeShuffleBytesRead(Windmill.WorkItem workItem) {
@@ -175,14 +167,6 @@ private static Windmill.WorkItemCommitRequest.Builder initializeOutputBuilder(
175167
.setCacheToken(workItem.getCacheToken());
176168
}
177169

178-
private static Windmill.WorkItemCommitRequest buildWorkItemTruncationRequest(
179-
ByteString key, Windmill.WorkItem workItem, int estimatedCommitSize) {
180-
Windmill.WorkItemCommitRequest.Builder outputBuilder = initializeOutputBuilder(key, workItem);
181-
outputBuilder.setExceedsMaxWorkItemCommitBytes(true);
182-
outputBuilder.setEstimatedWorkItemCommitBytes(estimatedCommitSize);
183-
return outputBuilder.build();
184-
}
185-
186170
/** Sets the stage name and workId of the Thread executing the {@link Work} for logging. */
187171
private static void setUpWorkLoggingContext(String workLatencyTrackingId, String computationId) {
188172
setLoggingContextWorkId(workLatencyTrackingId);
@@ -305,33 +289,6 @@ private void processWork(
305289
}
306290
}
307291

308-
private Windmill.WorkItemCommitRequest validateCommitRequestSize(
309-
Windmill.WorkItemCommitRequest commitRequest,
310-
String computationId,
311-
Windmill.WorkItem workItem) {
312-
long byteLimit = globalConfigHandle.getConfig().operationalLimits().getMaxWorkItemCommitBytes();
313-
int commitSize = commitRequest.getSerializedSize();
314-
int estimatedCommitSize = commitSize < 0 ? Integer.MAX_VALUE : commitSize;
315-
316-
// Detect overflow of integer serialized size or if the byte limit was exceeded.
317-
// Commit is too large if overflow has occurred or the commitSize has exceeded the allowed
318-
// commit byte limit.
319-
streamingCounters.windmillMaxObservedWorkItemCommitBytes().addValue(estimatedCommitSize);
320-
if (commitSize >= 0 && commitSize < byteLimit) {
321-
return commitRequest;
322-
}
323-
324-
KeyCommitTooLargeException e =
325-
KeyCommitTooLargeException.causedBy(computationId, byteLimit, commitRequest);
326-
failureTracker.trackFailure(computationId, workItem, e);
327-
LOG.error("{}", e.toString());
328-
329-
// Drop the current request in favor of a new, minimal one requesting truncation.
330-
// Messages, timers, counters, and other commit content will not be used by the service
331-
// so, we're purposefully dropping them here
332-
return buildWorkItemTruncationRequest(workItem.getKey(), workItem, estimatedCommitSize);
333-
}
334-
335292
private void recordProcessingStats(
336293
List<Work> workBatch,
337294
List<Windmill.WorkItemCommitRequest> workItemCommits,
@@ -448,17 +405,13 @@ private void commitWorkBatch(
448405

449406
private void commitSingleKeyWork(
450407
ComputationState computationState, Work work, Windmill.WorkItemCommitRequest commitRequest) {
451-
// Validate the commit request, possibly requesting truncation if the commitSize is too large.
452-
Windmill.WorkItemCommitRequest validatedCommitRequest =
453-
validateCommitRequestSize(
454-
commitRequest, computationState.getComputationId(), work.getWorkItem());
455408
work.setState(Work.State.COMMIT_QUEUED);
456-
validatedCommitRequest =
457-
validatedCommitRequest
409+
Windmill.WorkItemCommitRequest commitRequestWithAttributions =
410+
commitRequest
458411
.toBuilder()
459412
.addAllPerWorkItemLatencyAttributions(work.getLatencyAttributions(sampler))
460413
.build();
461-
work.queueCommit(validatedCommitRequest, computationState);
414+
work.queueCommit(commitRequestWithAttributions, computationState);
462415
}
463416

464417
private void recordProcessingTime(

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

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1290,8 +1290,9 @@ public void testKeyTokenInvalidException() throws Exception {
12901290
worker.stop();
12911291
}
12921292

1293-
@Test
1294-
public void testKeyCommitTooLargeException() throws Exception {
1293+
private void runKeyCommitTooLargeExceptionTest(
1294+
StreamingDataflowWorkerTestParams.Builder workerParams, boolean expectKeyInErrorMessage)
1295+
throws Exception {
12951296
KvCoder<String, String> kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of());
12961297

12971298
List<ParallelInstruction> instructions =
@@ -1304,7 +1305,7 @@ public void testKeyCommitTooLargeException() throws Exception {
13041305

13051306
StreamingDataflowWorker worker =
13061307
makeWorker(
1307-
defaultWorkerParams()
1308+
workerParams
13081309
.setInstructions(instructions)
13091310
.setStreamingGlobalConfig(
13101311
StreamingGlobalConfig.builder()
@@ -1336,7 +1337,6 @@ public void testKeyCommitTooLargeException() throws Exception {
13361337
1, "large_key", DEFAULT_SHARDING_KEY, largeCommit.getEstimatedWorkItemCommitBytes())
13371338
.build(),
13381339
removeDynamicFields(largeCommit));
1339-
13401340
// Check this explicitly since the estimated commit bytes weren't actually
13411341
// checked against an expected value in the previous step
13421342
assertTrue(largeCommit.getEstimatedWorkItemCommitBytes() > 1000);
@@ -1360,12 +1360,35 @@ public void testKeyCommitTooLargeException() throws Exception {
13601360
foundErrors = true;
13611361
String errorMessage = status.getErrors().get(0).getMessage();
13621362
assertThat(errorMessage, Matchers.containsString("KeyCommitTooLargeException"));
1363+
assertThat(errorMessage, Matchers.containsString("Commit request for stage computation"));
1364+
if (expectKeyInErrorMessage) {
1365+
assertThat(errorMessage, Matchers.containsString("and key large_key"));
1366+
} else {
1367+
assertThat(errorMessage, Matchers.not(Matchers.containsString("and key large_key")));
1368+
}
13631369
}
13641370
}
13651371
assertTrue(foundErrors);
13661372
worker.stop();
13671373
}
13681374

1375+
@Test
1376+
public void testKeyCommitTooLargeException() throws Exception {
1377+
runKeyCommitTooLargeExceptionTest(defaultWorkerParams(), /* expectKeyInErrorMessage= */ false);
1378+
}
1379+
1380+
@Test
1381+
public void testKeyCommitTooLargeException_withHotKeyLoggingEnabled() throws Exception {
1382+
runKeyCommitTooLargeExceptionTest(
1383+
defaultWorkerParams("--hotKeyLoggingEnabled=true"), /* expectKeyInErrorMessage= */ true);
1384+
}
1385+
1386+
@Test
1387+
public void testKeyCommitTooLargeException_withHotKeyLoggingDisabled() throws Exception {
1388+
runKeyCommitTooLargeExceptionTest(
1389+
defaultWorkerParams("--hotKeyLoggingEnabled=false"), /* expectKeyInErrorMessage= */ false);
1390+
}
1391+
13691392
@Test
13701393
public void testOutputKeyTooLargeException() throws Exception {
13711394
KvCoder<String, String> kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of());

0 commit comments

Comments
 (0)