This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 5e499388670 KeyCommitTooLarge logging improvements (#39316)
5e499388670 is described below
commit 5e499388670c269bea1aa52d584d1adfb76036c2
Author: Ryan Wigglesworth <[email protected]>
AuthorDate: Wed Jul 22 19:21:48 2026 -0700
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](https://github.com/actions/checkout/compare/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] <[email protected]>
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](https://github.com/googleapis/google-cloud-go/compare/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] <[email protected]>
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](https://github.com/aws/aws-sdk-go-v2/compare/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] <[email protected]>
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 <[email protected]>
* respond to comments
---------
Signed-off-by: dependabot[bot] <[email protected]>
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 <[email protected]>
Co-authored-by: Abdelrahman Ibrahim <[email protected]>
Co-authored-by: raman118 <[email protected]>
Co-authored-by: Arun Pandian <[email protected]>
---
.../worker/StreamingModeExecutionContext.java | 53 +++++++++++++++++++
.../streaming/KeyCommitTooLargeException.java | 20 ++++---
.../processing/ComputationWorkExecutorFactory.java | 14 ++++-
.../work/processing/StreamingWorkScheduler.java | 61 +++-------------------
.../worker/StreamingDataflowWorkerTest.java | 31 +++++++++--
.../worker/StreamingModeExecutionContextTest.java | 5 ++
.../dataflow/worker/WorkerCustomSourcesTest.java | 8 +++
7 files changed, 126 insertions(+), 66 deletions(-)
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java
index 9401cc5f8ed..db930d53f76 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java
+++
b/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.CounterFactory;
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.WindmillTagEncodin
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 class StreamingModeExecutionContext
HotKeyLogger hotKeyLogger,
boolean hotKeyLoggingEnabled,
String stepName,
+ String systemName,
+ StreamingCounters streamingCounters,
+ FailureTracker failureTracker,
String sourceBytesProcessCounterName,
SideInputStateFetcherFactory sideInputStateFetcherFactory) {
super(
@@ -231,6 +240,9 @@ public class 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 @@ public class StreamingModeExecutionContext
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) {
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 4d6ae8a208c..331b9a2a734 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,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 final class KeyCommitTooLargeException extends
Exception {
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/ComputationWorkExecutorFactory.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java
index 4a52d9fde77..0c3591102f9 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/ComputationWorkExecutorFactory.java
+++
b/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.ComputationState;
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 @@ final class ComputationWorkExecutorFactory {
hotKeyLogger,
hotKeyLoggingEnabled,
stepName,
+ stageInfo.systemName(),
+ streamingCounters,
+ failureTracker,
computationState.sourceBytesProcessCounterName(),
sideInputStateFetcherFactory);
}
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 e9b85d720d2..990393f51f8 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
@@ -43,7 +43,6 @@ import
org.apache.beam.runners.dataflow.worker.streaming.BoundedQueueExecutorWor
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 class StreamingWorkScheduler {
readerCache,
stateCacheFactory,
sampler,
- streamingCounters.pendingDeltaCounters(),
+ streamingCounters,
+ failureTracker,
idGenerator,
globalConfigHandle,
hotKeyLogger,
@@ -149,13 +143,11 @@ public class StreamingWorkScheduler {
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 @@ public class StreamingWorkScheduler {
.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 @@ public class StreamingWorkScheduler {
}
}
- 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 @@ public class StreamingWorkScheduler {
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(
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 23730bc5770..1e7f3b01f00 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 class StreamingDataflowWorkerTest {
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 class StreamingDataflowWorkerTest {
StreamingDataflowWorker worker =
makeWorker(
- defaultWorkerParams()
+ workerParams
.setInstructions(instructions)
.setStreamingGlobalConfig(
StreamingGlobalConfig.builder()
@@ -1336,7 +1337,6 @@ public class StreamingDataflowWorkerTest {
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 class StreamingDataflowWorkerTest {
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());
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
index 561596f68d0..3d11e509746 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java
@@ -63,6 +63,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.Work;
import
org.apache.beam.runners.dataflow.worker.streaming.config.FakeGlobalConfigHandle;
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.SideInputStateFetcherFactory;
import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
@@ -71,6 +72,7 @@ import
org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache
import
org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateReader;
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.work.processing.failures.FailureTracker;
import
org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.Coder;
@@ -142,6 +144,9 @@ public class StreamingModeExecutionContextTest {
new HotKeyLogger(),
/*hotKeyLoggingEnabled=*/ false,
/*stepName=*/ "stepName",
+ /*systemName=*/ "systemName",
+ StreamingCounters.create(),
+ mock(FailureTracker.class),
"sourceBytesProcessCounterName",
SideInputStateFetcherFactory.fromOptions(options));
}
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java
index 31ea1bab07a..2da52bea8eb 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java
@@ -94,6 +94,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.Work;
import
org.apache.beam.runners.dataflow.worker.streaming.config.FixedGlobalConfigHandle;
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.SideInputStateFetcherFactory;
import org.apache.beam.runners.dataflow.worker.testing.TestCountingSource;
import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader;
@@ -103,6 +104,7 @@ import
org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import
org.apache.beam.runners.dataflow.worker.windmill.client.getdata.FakeGetDataClient;
import
org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache;
import
org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateReader;
+import
org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker;
import
org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
@@ -641,6 +643,9 @@ public class WorkerCustomSourcesTest {
new HotKeyLogger(),
/*hotKeyLoggingEnabled=*/ false,
/*stepName=*/ "stepName",
+ /*systemName=*/ "systemName",
+ StreamingCounters.create(),
+ mock(FailureTracker.class),
"sourceBytesProcessCounterName",
SideInputStateFetcherFactory.fromOptions(options));
@@ -1014,6 +1019,9 @@ public class WorkerCustomSourcesTest {
new HotKeyLogger(),
/*hotKeyLoggingEnabled=*/ false,
/*stepName=*/ "stepName",
+ /*systemName=*/ "systemName",
+ StreamingCounters.create(),
+ mock(FailureTracker.class),
"sourceBytesProcessCounterName",
SideInputStateFetcherFactory.fromOptions(options));
