From c831c3363077e48942a661f02faa3893fe09cd3e Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Tue, 14 Jul 2026 00:22:40 +0000 Subject: [PATCH 01/10] 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 --- .../streaming/KeyCommitTooLargeException.java | 25 ++++++-- .../processing/StreamingWorkScheduler.java | 30 ++++++--- .../worker/StreamingDataflowWorkerTest.java | 61 +++++++++++++------ 3 files changed, 80 insertions(+), 36 deletions(-) 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..227dfa235218 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 @@ -17,17 +17,30 @@ */ package org.apache.beam.runners.dataflow.worker.streaming; +import com.google.protobuf.TextFormat; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; 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()); } @@ -49,3 +61,4 @@ private KeyCommitTooLargeException(String message) { super(message); } } + 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..eb713a35d150 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 @@ -17,10 +17,13 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.work.processing; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Preconditions.checkState; +import static org.apache.beam.runners.dataflow.DataflowRunner.hasExperiment; import com.google.api.services.dataflow.model.MapTask; import com.google.auto.value.AutoValue; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.ByteString; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentMap; @@ -29,7 +32,6 @@ import java.util.function.Function; import java.util.function.Supplier; import javax.annotation.concurrent.ThreadSafe; -import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair; import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; import org.apache.beam.runners.dataflow.worker.DataflowExecutionStateSampler; import org.apache.beam.runners.dataflow.worker.DataflowMapTaskExecutorFactory; @@ -61,8 +63,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.WorkFailureProcessor; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.fn.IdGenerator; -import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.commons.lang3.tuple.Pair; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import org.slf4j.Logger; @@ -89,6 +90,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 +102,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 +114,7 @@ public StreamingWorkScheduler( this.stageInfoMap = stageInfoMap; this.sampler = sampler; this.globalConfigHandle = globalConfigHandle; + this.hotKeyLoggingEnabled = hotKeyLoggingEnabled; } public static StreamingWorkScheduler create( @@ -145,6 +149,9 @@ public static StreamingWorkScheduler create( hotKeyLogger, sideInputStateFetcherFactory); + boolean hotKeyLoggingEnabled = + options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging"); + return new StreamingWorkScheduler( clock, workExecutor, @@ -155,7 +162,8 @@ public static StreamingWorkScheduler create( streamingCounters, stageInfoMap, sampler, - globalConfigHandle); + globalConfigHandle, + hotKeyLoggingEnabled); } private static long computeShuffleBytesRead(Windmill.WorkItem workItem) { @@ -307,7 +315,7 @@ private void processWork( private Windmill.WorkItemCommitRequest validateCommitRequestSize( Windmill.WorkItemCommitRequest commitRequest, - String computationId, + String stageName, Windmill.WorkItem workItem) { long byteLimit = globalConfigHandle.getConfig().operationalLimits().getMaxWorkItemCommitBytes(); int commitSize = commitRequest.getSerializedSize(); @@ -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 @@ -509,3 +518,4 @@ static ExecuteWorkResult create( abstract long stateBytesRead(); } } + 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..43eda2ae8106 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 @@ -58,6 +58,19 @@ import com.google.api.services.dataflow.model.WorkItemStatus; import com.google.api.services.dataflow.model.WriteInstruction; import com.google.auto.value.AutoValue; +import com.google.common.cache.CacheStats; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; +import com.google.common.collect.Lists; +import com.google.common.primitives.UnsignedLong; +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import com.google.common.util.concurrent.Uninterruptibles; +import com.google.protobuf.ByteString; +import com.google.protobuf.TextFormat; +import io.grpc.Server; +import io.grpc.ServerBuilder; +import io.grpc.testing.GrpcCleanupRule; import java.io.IOException; import java.io.InputStream; import java.net.ServerSocket; @@ -187,19 +200,6 @@ import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode; -import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; -import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Server; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ServerBuilder; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.testing.GrpcCleanupRule; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedLong; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles; import org.hamcrest.Matcher; import org.hamcrest.Matchers; import org.joda.time.Duration; @@ -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()); @@ -5105,3 +5125,4 @@ final Builder publishCounters() { } } } + From 25d7bbc34189611eabc781bab6f98a82cec81cbd Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Tue, 14 Jul 2026 11:01:49 -0700 Subject: [PATCH 02/10] 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> --- .../dataflow/worker/streaming/KeyCommitTooLargeException.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 227dfa235218..59b2d050d957 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 @@ -17,7 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker.streaming; -import com.google.protobuf.TextFormat; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; public final class KeyCommitTooLargeException extends Exception { From 2513ffebcb01177da6006a4e56c6de9cb5919c34 Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Tue, 14 Jul 2026 11:03:00 -0700 Subject: [PATCH 03/10] 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> --- .../windmill/work/processing/StreamingWorkScheduler.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) 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 eb713a35d150..7ad95ee7958c 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 @@ -17,16 +17,15 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.work.processing; -import static com.google.common.base.Preconditions.checkState; -import static org.apache.beam.runners.dataflow.DataflowRunner.hasExperiment; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.api.services.dataflow.model.MapTask; import com.google.auto.value.AutoValue; -import com.google.common.collect.ImmutableList; -import com.google.protobuf.ByteString; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentMap; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.function.Function; From 647a81542c8065daca4a9b43dfebacc95de06779 Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Tue, 14 Jul 2026 11:09:18 -0700 Subject: [PATCH 04/10] 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> --- .../worker/StreamingDataflowWorkerTest.java | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) 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 43eda2ae8106..e73bbbfc6dc3 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 @@ -58,19 +58,19 @@ import com.google.api.services.dataflow.model.WorkItemStatus; import com.google.api.services.dataflow.model.WriteInstruction; import com.google.auto.value.AutoValue; -import com.google.common.cache.CacheStats; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Iterables; -import com.google.common.collect.Lists; -import com.google.common.primitives.UnsignedLong; -import com.google.common.util.concurrent.ThreadFactoryBuilder; -import com.google.common.util.concurrent.Uninterruptibles; -import com.google.protobuf.ByteString; -import com.google.protobuf.TextFormat; -import io.grpc.Server; -import io.grpc.ServerBuilder; -import io.grpc.testing.GrpcCleanupRule; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Server; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ServerBuilder; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.testing.GrpcCleanupRule; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedLong; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles; import java.io.IOException; import java.io.InputStream; import java.net.ServerSocket; From 369aff8a9d0fe3cd6680eabab56ecac5ad7af846 Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Tue, 14 Jul 2026 20:26:38 +0000 Subject: [PATCH 05/10] gemini review responses --- .../streaming/KeyCommitTooLargeException.java | 2 +- .../processing/StreamingWorkScheduler.java | 11 +++++--- .../worker/StreamingDataflowWorkerTest.java | 26 +++++++++---------- 3 files changed, 21 insertions(+), 18 deletions(-) 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 59b2d050d957..227dfa235218 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 @@ -17,7 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker.streaming; -import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat; +import com.google.protobuf.TextFormat; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; public final class KeyCommitTooLargeException extends Exception { 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 7ad95ee7958c..84457942de06 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 @@ -17,15 +17,15 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.work.processing; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Preconditions.checkState; import com.google.api.services.dataflow.model.MapTask; import com.google.auto.value.AutoValue; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.ByteString; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentMap; -import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.function.Function; @@ -149,7 +149,10 @@ public static StreamingWorkScheduler create( sideInputStateFetcherFactory); boolean hotKeyLoggingEnabled = - options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging"); + options.isHotKeyLoggingEnabled() + || (options.getExperiments() != null + && options.getExperiments().stream() + .anyMatch("enable_hot_key_logging"::equalsIgnoreCase)); return new StreamingWorkScheduler( clock, 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 e73bbbfc6dc3..31aceb012eac 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 @@ -58,19 +58,6 @@ import com.google.api.services.dataflow.model.WorkItemStatus; import com.google.api.services.dataflow.model.WriteInstruction; import com.google.auto.value.AutoValue; -import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; -import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Server; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ServerBuilder; -import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.testing.GrpcCleanupRule; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedLong; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles; import java.io.IOException; import java.io.InputStream; import java.net.ServerSocket; @@ -200,6 +187,19 @@ import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode; +import com.google.protobuf.ByteString; +import com.google.protobuf.TextFormat; +import io.grpc.Server; +import io.grpc.ServerBuilder; +import io.grpc.testing.GrpcCleanupRule; +import com.google.common.cache.CacheStats; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; +import com.google.common.collect.Lists; +import com.google.common.primitives.UnsignedLong; +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import com.google.common.util.concurrent.Uninterruptibles; import org.hamcrest.Matcher; import org.hamcrest.Matchers; import org.joda.time.Duration; From 15751e2636bd7a00070115e193f6e4729cfb7b5c Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Tue, 14 Jul 2026 23:30:06 +0000 Subject: [PATCH 06/10] Fix imports --- .../streaming/KeyCommitTooLargeException.java | 2 +- .../processing/StreamingWorkScheduler.java | 6 ++--- .../worker/StreamingDataflowWorkerTest.java | 26 +++++++++---------- 3 files changed, 17 insertions(+), 17 deletions(-) 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 227dfa235218..59b2d050d957 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 @@ -17,7 +17,7 @@ */ package org.apache.beam.runners.dataflow.worker.streaming; -import com.google.protobuf.TextFormat; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; public final class KeyCommitTooLargeException extends Exception { 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 84457942de06..0a73c70843d9 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 @@ -17,12 +17,10 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.work.processing; -import static com.google.common.base.Preconditions.checkState; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.api.services.dataflow.model.MapTask; import com.google.auto.value.AutoValue; -import com.google.common.collect.ImmutableList; -import com.google.protobuf.ByteString; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentMap; @@ -62,6 +60,8 @@ import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.WorkFailureProcessor; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.fn.IdGenerator; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.commons.lang3.tuple.Pair; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; 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 31aceb012eac..645a3447562f 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 @@ -187,19 +187,19 @@ import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder; import org.apache.beam.sdk.values.WindowingStrategy; import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode; -import com.google.protobuf.ByteString; -import com.google.protobuf.TextFormat; -import io.grpc.Server; -import io.grpc.ServerBuilder; -import io.grpc.testing.GrpcCleanupRule; -import com.google.common.cache.CacheStats; -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Iterables; -import com.google.common.collect.Lists; -import com.google.common.primitives.UnsignedLong; -import com.google.common.util.concurrent.ThreadFactoryBuilder; -import com.google.common.util.concurrent.Uninterruptibles; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Server; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ServerBuilder; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.testing.GrpcCleanupRule; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheStats; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.UnsignedLong; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Uninterruptibles; import org.hamcrest.Matcher; import org.hamcrest.Matchers; import org.joda.time.Duration; From d3a4d1c6e7d2de41556f70f391fa28093c707eec Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Tue, 14 Jul 2026 23:32:55 +0000 Subject: [PATCH 07/10] one more --- .../worker/windmill/work/processing/StreamingWorkScheduler.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 0a73c70843d9..7fb37edf6162 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 @@ -29,6 +29,7 @@ import java.util.function.Function; import java.util.function.Supplier; import javax.annotation.concurrent.ThreadSafe; +import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair; import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; import org.apache.beam.runners.dataflow.worker.DataflowExecutionStateSampler; import org.apache.beam.runners.dataflow.worker.DataflowMapTaskExecutorFactory; @@ -62,7 +63,6 @@ import org.apache.beam.sdk.fn.IdGenerator; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; -import org.apache.commons.lang3.tuple.Pair; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; import org.slf4j.Logger; From d5babb02d6051888474aa1bd6efe28529a0192b2 Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Wed, 15 Jul 2026 04:20:03 +0000 Subject: [PATCH 08/10] one more attempt --- .../windmill/work/processing/StreamingWorkScheduler.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 7fb37edf6162..dbb5bc391cb9 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 @@ -148,11 +148,11 @@ public static StreamingWorkScheduler create( hotKeyLogger, sideInputStateFetcherFactory); + List experiments = options.getExperiments(); boolean hotKeyLoggingEnabled = options.isHotKeyLoggingEnabled() - || (options.getExperiments() != null - && options.getExperiments().stream() - .anyMatch("enable_hot_key_logging"::equalsIgnoreCase)); + || (experiments != null + && experiments.stream().anyMatch("enable_hot_key_logging"::equalsIgnoreCase)); return new StreamingWorkScheduler( clock, From 9f611b500893f95f725b951ccbce0a5f1fba6b58 Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Wed, 15 Jul 2026 07:02:07 +0000 Subject: [PATCH 09/10] one more attempt --- .../dataflow/worker/streaming/KeyCommitTooLargeException.java | 2 +- .../windmill/work/processing/StreamingWorkScheduler.java | 4 +--- .../runners/dataflow/worker/StreamingDataflowWorkerTest.java | 1 - 3 files changed, 2 insertions(+), 5 deletions(-) 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 59b2d050d957..950cf9b4ac41 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 @@ -17,8 +17,8 @@ */ package org.apache.beam.runners.dataflow.worker.streaming; -import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.TextFormat; 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 { 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 dbb5bc391cb9..4947d3769e88 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 @@ -316,9 +316,7 @@ private void processWork( } private Windmill.WorkItemCommitRequest validateCommitRequestSize( - Windmill.WorkItemCommitRequest commitRequest, - String stageName, - 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; 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 645a3447562f..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 @@ -5125,4 +5125,3 @@ final Builder publishCounters() { } } } - From a5bd7d0d38ad441d3d1e3bc94eae9263a7e44db6 Mon Sep 17 00:00:00 2001 From: Ryan Wigglesworth Date: Wed, 15 Jul 2026 17:36:55 +0000 Subject: [PATCH 10/10] formatting fix --- .../dataflow/worker/streaming/KeyCommitTooLargeException.java | 1 - .../worker/windmill/work/processing/StreamingWorkScheduler.java | 1 - 2 files changed, 2 deletions(-) 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 950cf9b4ac41..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 @@ -61,4 +61,3 @@ private KeyCommitTooLargeException(String message) { super(message); } } - 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 4947d3769e88..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 @@ -518,4 +518,3 @@ static ExecuteWorkResult create( abstract long stateBytesRead(); } } -