From 4df59615d47a8d30a57b3fd6a5c0bf1b39674c83 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 20:21:01 +0200 Subject: [PATCH 1/4] [Dataflow Java] Extend work item lease while DoFn.setup runs in batch worker Batch work items failed when DoFn.setup took longer than the 3 minute lease, because progress reporting started only after the executor was created. Start reporting first, send lease-only updates until the executor exists, and apply a service abort received during setup once the executor is set. Fixes #34206 --- .../dataflow/worker/BatchDataflowWorker.java | 37 +++++++++++-- .../worker/DataflowWorkProgressUpdater.java | 36 +++++++++++-- .../dataflow/worker/WorkItemStatusClient.java | 13 ++++- .../worker/BatchDataflowWorkerTest.java | 54 +++++++++++++++++++ .../DataflowWorkProgressUpdaterTest.java | 16 ++++++ .../worker/WorkItemStatusClientTest.java | 14 +++-- 6 files changed, 156 insertions(+), 14 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java index fa6adaa75212..1b411c679b3d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java @@ -22,6 +22,7 @@ import java.io.Closeable; import java.io.IOException; import java.util.Collection; +import java.util.Objects; import java.util.function.Function; import org.apache.beam.runners.core.metrics.ExecutionStateSampler; import org.apache.beam.runners.dataflow.DataflowRunner; @@ -233,6 +234,9 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr LOG.debug("Executing: {}", workItem); DataflowWorkExecutor worker = null; + DataflowWorkProgressUpdater progressUpdater = null; + // once execution starts, executeStartedWork owns stopping the progress updater + boolean executionStarted = false; try { // Populate PipelineOptions with data from work unit. options.setProject(workItem.getProjectId()); @@ -263,6 +267,12 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr LOG.debug("Network as Graphviz .dot: {}", Networks.toDot(network)); } + // creating the executor runs every DoFn.setup, which can take longer than the 3 minute + // lease. Reporting starts first so the lease keeps being extended; without this the service + // fails the work item, and after 4 failures of the same item it fails the whole job + progressUpdater = + new DataflowWorkProgressUpdater(workItemStatusClient, workItem, null, options); + progressUpdater.startReportingProgress(); worker = mapTaskExecutorFactory.create( network, @@ -286,14 +296,30 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr } workItemStatusClient.setWorker(worker, executionContext); - DataflowWorkProgressUpdater progressUpdater = - new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker, options); - executeWork(worker, progressUpdater); + if (Objects.isNull(progressUpdater)) { + progressUpdater = + new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker, options); + executionStarted = true; + executeWork(worker, progressUpdater); + } else { + progressUpdater.setWorker(worker); + executionStarted = true; + executeStartedWork(worker, progressUpdater); + } workItemStatusClient.reportSuccess(); return true; } catch (OutOfMemoryError oom) { throw oom; } catch (Throwable e) { + // executor creation failed while progress was already being reported; stop it first so no + // lease extension is sent after the error report + if (Objects.nonNull(progressUpdater) && !executionStarted) { + try { + progressUpdater.stopReportingProgress(); + } catch (Exception stopExn) { + e.addSuppressed(stopExn); + } + } workItemStatusClient.reportError(e); return false; } finally { @@ -314,6 +340,11 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr void executeWork(DataflowWorkExecutor worker, DataflowWorkProgressUpdater progressUpdater) throws Exception { progressUpdater.startReportingProgress(); + executeStartedWork(worker, progressUpdater); + } + + private void executeStartedWork( + DataflowWorkExecutor worker, DataflowWorkProgressUpdater progressUpdater) throws Exception { // Blocks while executing the work. try { worker.execute(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java index b94759c239a3..a492818d260e 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java @@ -26,6 +26,7 @@ import com.google.api.services.dataflow.model.HotKeyDetection; import com.google.api.services.dataflow.model.WorkItem; import com.google.api.services.dataflow.model.WorkItemServiceState; +import java.util.Objects; import java.util.concurrent.ScheduledExecutorService; import javax.annotation.concurrent.NotThreadSafe; import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions; @@ -58,7 +59,13 @@ public class DataflowWorkProgressUpdater extends WorkProgressUpdater { private HotKeyLogger hotKeyLogger; - private boolean wasAskedToAbort = false; + private volatile boolean wasAskedToAbort = false; + + /** + * Guards the hand-off between an abort that arrives while the executor is still being created and + * {@link #setWorker}, so the abort reaches the executor whichever thread runs first. + */ + private final Object abortLock = new Object(); public DataflowWorkProgressUpdater( WorkItemStatusClient workItemStatusClient, @@ -93,6 +100,19 @@ public DataflowWorkProgressUpdater( this.options = options; } + @Override + public void setWorker(WorkExecutor worker) { + boolean abortRequested; + synchronized (abortLock) { + super.setWorker(worker); + abortRequested = wasAskedToAbort; + } + // a map task executor remembers the abort, so its read loop stops as soon as execute() starts + if (abortRequested) { + worker.abort(); + } + } + @Override protected String workString() { return workItemStatusClient.uniqueWorkId(); @@ -122,8 +142,15 @@ protected void reportProgressHelper() throws Exception { if (result.getCompleteWorkStatus() != null && result.getCompleteWorkStatus().getCode() != com.google.rpc.Code.OK.getNumber()) { LOG.info("Service asked worker to abort with status: {}", result.getCompleteWorkStatus()); - wasAskedToAbort = true; - worker.abort(); + WorkExecutor currentWorker; + synchronized (abortLock) { + wasAskedToAbort = true; + currentWorker = worker; + } + // the executor may still be under construction, in which case setWorker aborts it + if (Objects.nonNull(currentWorker)) { + currentWorker.abort(); + } return; } @@ -154,7 +181,8 @@ protected void reportProgressHelper() throws Exception { leaseRemainingTime(getLeaseExpirationTimestamp(result))); ApproximateSplitRequest suggestedStopPoint = result.getSplitRequest(); - if (suggestedStopPoint != null) { + // there is nothing to split until the executor exists, the service asks again later + if (suggestedStopPoint != null && Objects.nonNull(worker)) { LOG.info("Proposing dynamic split of work unit {} at {}", workString(), suggestedStopPoint); dynamicSplitResultToReport = worker.requestDynamicSplit( diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemStatusClient.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemStatusClient.java index 85cf877481a1..cf9659f58ad9 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemStatusClient.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemStatusClient.java @@ -36,6 +36,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.function.Consumer; import org.apache.beam.runners.core.metrics.ExecutionStateTracker; import org.apache.beam.runners.core.metrics.ExecutionStateTracker.ExecutionState; @@ -176,11 +177,14 @@ public synchronized void setWorker( return execute(status); } - /** Return the {@link WorkItemServiceState} resulting from sending a progress update. */ + /** + * Return the {@link WorkItemServiceState} resulting from sending a progress update. Before {@link + * #setWorker} is called the update carries no progress, counters or metrics and only extends the + * lease. + */ public synchronized @Nullable WorkItemServiceState reportUpdate( @Nullable DynamicSplitResult dynamicSplitResult, Duration requestedLeaseDuration) throws Exception { - checkState(worker != null, "setWorker should be called before reportUpdate"); checkState(!finalStateSent, "cannot reportUpdates after sending a final state"); checkArgument(requestedLeaseDuration != null, "requestLeaseDuration must be non-null"); if (wasAskedToAbort) { @@ -261,6 +265,11 @@ synchronized void populateSplitResult( @VisibleForTesting synchronized void populateProgress(WorkItemStatus status) throws Exception { + // updates sent while the executor is still being created only extend the lease + if (Objects.isNull(worker)) { + return; + } + Progress progress = worker.getWorkerProgress(); if (progress != null) { status.setReportedProgress(SourceTranslationUtils.readerProgressToCloudProgress(progress)); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorkerTest.java index e33412a19d94..e8c3b0d3b0d6 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorkerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorkerTest.java @@ -38,6 +38,8 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions; import org.apache.beam.runners.dataflow.util.TimeUtil; import org.apache.beam.sdk.options.PipelineOptionsFactory; @@ -47,6 +49,7 @@ import org.hamcrest.TypeSafeMatcher; import org.joda.time.Duration; import org.joda.time.Instant; +import org.junit.Assert; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -55,6 +58,7 @@ import org.junit.runners.JUnit4; import org.mockito.ArgumentCaptor; import org.mockito.Mock; +import org.mockito.Mockito; import org.mockito.MockitoAnnotations; import org.mockito.hamcrest.MockitoHamcrest; @@ -161,6 +165,56 @@ public void testStopProgressReportInCaseOfFailure() throws Exception { verify(mockProgressUpdater, times(1)).stopReportingProgress(); } + @Test + public void testProgressReportedWhileExecutorIsBeingCreated() throws Exception { + // creating the executor runs every DoFn.setup, which can take minutes; if no status reaches + // the service in that time the lease expires and the work item is failed + CountDownLatch statusReported = new CountDownLatch(1); + Mockito.when(mockWorkUnitClient.reportWorkItemStatus(Mockito.any())) + .thenAnswer( + invocation -> { + statusReported.countDown(); + return null; + }); + boolean[] reportedBeforeCreateReturned = new boolean[1]; + DataflowMapTaskExecutorFactory blockingFactory = + (network, + pipelineOptions, + stageName, + readerFactory, + sinkFactory, + executionContext, + counterSet, + idGenerator) -> { + try { + reportedBeforeCreateReturned[0] = statusReported.await(30, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return mock(DataflowMapTaskExecutor.class); + }; + BatchDataflowWorker worker = + new BatchDataflowWorker(mockWorkUnitClient, blockingFactory, options); + WorkItem workItem = + new WorkItem() + .setId(1L) + .setJobId("SlowExecutorCreation") + .setInitialReportIndex(1L) + .setMapTask( + new MapTask() + .setInstructions(new ArrayList()) + .setStageName("testStage")) + .setLeaseExpireTime( + TimeUtil.toCloudTime(Instant.now().plus(Duration.standardMinutes(3)))) + .setReportStatusInterval(TimeUtil.toCloudDuration(Duration.standardSeconds(1))); + + worker.doWork(workItem, new WorkItemStatusClient(mockWorkUnitClient, workItem)); + + Assert.assertTrue( + "no status was reported while the executor was being created", + reportedBeforeCreateReturned[0]); + } + @Test public void testIsSplitResponseTooLarge() throws IOException { SourceSplitResponse splitResponse = new SourceSplitResponse(); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java index 0ede8383818f..ea797c901e52 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java @@ -182,6 +182,22 @@ public void workProgressAskedToAbortWorker() throws Exception { progressUpdater.stopReportingProgress(); } + @Test + public void workProgressAbortBeforeWorkerIsSetAbortsItOnceSet() throws Exception { + // the batch worker starts reporting while the executor is still being created + progressUpdater.setWorker(null); + when(workItemStatusClient.reportUpdate(isNull(DynamicSplitResult.class), isA(Duration.class))) + .thenReturn(generateServiceAbort()); + progressUpdater.startReportingProgress(); + executor.runNextRunnable(); + Mockito.verify(worker, Mockito.never()).abort(); + + progressUpdater.setWorker(worker); + + Mockito.verify(worker, Mockito.times(1)).abort(); + progressUpdater.stopReportingProgress(); + } + /** Verifies that the update after a split is requested acknowledeges it. */ @Test public void workProgressSendsSplitResults() throws Exception { diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkItemStatusClientTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkItemStatusClientTest.java index 5d0195e9a5d4..494d0214ffef 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkItemStatusClientTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkItemStatusClientTest.java @@ -518,11 +518,15 @@ public void populateSplitResultNull() throws Exception { } @Test - public void reportUpdateBeforeSetWorker() throws Exception { - thrown.expect(IllegalStateException.class); - thrown.expectMessage("setWorker"); - thrown.expectMessage("reportUpdate"); - statusClient.reportUpdate(null, null); + public void reportUpdateBeforeSetWorkerOnlyExtendsLease() throws Exception { + statusClient.reportUpdate(null, Duration.standardMinutes(3)); + + verify(workUnitClient).reportWorkItemStatus(statusCaptor.capture()); + WorkItemStatus workStatus = statusCaptor.getValue(); + assertThat(workStatus.getRequestedLeaseDuration(), equalTo("180s")); + assertThat(workStatus.getReportedProgress(), nullValue()); + assertThat(workStatus.getCounterUpdates(), nullValue()); + assertThat(workStatus.getMetricUpdates(), nullValue()); } private static class DummyBoundedSource extends BoundedSource { From e82738cbbe3156fb9b2cd53482043bdb63105840 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 21:51:53 +0200 Subject: [PATCH 2/4] [Dataflow Java] Warn while DoFn.setup is still running in batch worker The lease is now extended during setup, so a stuck setup no longer fails the work item on its own. Log a warning after 3 minutes and every 5 minutes after that while the executor is still being created, so it shows in the worker log. Skip worker.abort() when the progress thread is interrupted before the executor exists, which would throw a NullPointerException. Reword the slow setup warning, which no longer leads to lease expiry. Addresses #34206 --- .../worker/DataflowWorkProgressUpdater.java | 39 +++++++++++++++++++ .../dataflow/worker/DoFnInstanceManagers.java | 9 +++-- .../common/worker/WorkProgressUpdater.java | 7 +++- .../DataflowWorkProgressUpdaterTest.java | 21 ++++++++++ 4 files changed, 71 insertions(+), 5 deletions(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java index a492818d260e..94e479fda614 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java @@ -50,6 +50,15 @@ public class DataflowWorkProgressUpdater extends WorkProgressUpdater { private static final Logger LOG = LoggerFactory.getLogger(DataflowWorkProgressUpdater.class); + /** + * How long executor creation may run before the first "still creating" warning. Matches the point + * where {@code DoFnInstanceManagers} calls a single DoFn.setup slow (3 minutes). + */ + private static final long SLOW_SETUP_FIRST_WARNING_MILLIS = 3 * 60 * 1000; + + /** Gap between repeated "still creating" warnings, so a stuck setup keeps showing in the log. */ + private static final long SLOW_SETUP_WARNING_INTERVAL_MILLIS = 5 * 60 * 1000; + private final WorkItemStatusClient workItemStatusClient; /** The WorkItem for which work progress updates are sent. */ @@ -67,6 +76,12 @@ public class DataflowWorkProgressUpdater extends WorkProgressUpdater { */ private final Object abortLock = new Object(); + /** When this updater was created, which for the batch worker is when executor creation starts. */ + private final long createdAtMillis; + + /** Only read and written by the progress reporting thread. */ + private long nextSlowSetupWarningAtMillis; + public DataflowWorkProgressUpdater( WorkItemStatusClient workItemStatusClient, WorkItem workItem, @@ -77,6 +92,8 @@ public DataflowWorkProgressUpdater( this.workItem = workItem; this.hotKeyLogger = new HotKeyLogger(); this.options = options; + this.createdAtMillis = clock.currentTimeMillis(); + this.nextSlowSetupWarningAtMillis = createdAtMillis + SLOW_SETUP_FIRST_WARNING_MILLIS; } /** @@ -98,6 +115,8 @@ public DataflowWorkProgressUpdater( this.workItem = workItem; this.hotKeyLogger = hotKeyLogger; this.options = options; + this.createdAtMillis = clock.currentTimeMillis(); + this.nextSlowSetupWarningAtMillis = createdAtMillis + SLOW_SETUP_FIRST_WARNING_MILLIS; } @Override @@ -134,6 +153,9 @@ protected void reportProgressHelper() throws Exception { LOG.info("Service already asked to abort work item, not reporting ignored progress."); return; } + if (Objects.isNull(worker)) { + warnIfSetupIsSlow(); + } WorkItemServiceState result = workItemStatusClient.reportUpdate( dynamicSplitResultToReport, Duration.millis(requestedLeaseDurationMs)); @@ -191,6 +213,23 @@ protected void reportProgressHelper() throws Exception { } } + private void warnIfSetupIsSlow() { + long now = clock.currentTimeMillis(); + if (now < nextSlowSetupWarningAtMillis) { + return; + } + + // the lease keeps being extended while setup runs, so a stuck setup never fails the work item + // on its own. Repeating the warning is what makes it visible in the worker log + nextSlowSetupWarningAtMillis = now + SLOW_SETUP_WARNING_INTERVAL_MILLIS; + LOG.warn( + "Work item {} is still creating its executor after {} seconds. Creating the executor runs" + + " DoFn.setup, so a slow or stuck setup keeps this work item leased without" + + " processing any elements.", + workString(), + (now - createdAtMillis) / 1000); + } + /** Returns the given work unit's lease expiration timestamp. */ private long getLeaseExpirationTimestamp(WorkItem workItem) { return fromCloudTime(workItem.getLeaseExpireTime()).getMillis(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DoFnInstanceManagers.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DoFnInstanceManagers.java index 32a8445d2211..a0a6c1585637 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DoFnInstanceManagers.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DoFnInstanceManagers.java @@ -91,11 +91,12 @@ private ConcurrentQueueInstanceManager(DoFnInfo info, PipelineOptions opti DoFnInvokers.tryInvokeSetupFor(fn.getDoFn(), options); long elapsed = System.currentTimeMillis() - startMillis; if (elapsed > 180_000) { // 3 min - // Work item could fail for long-running setup due to Dataflow worker lease timeout + // setup runs again on every new worker and for every thread that needs its own instance, + // so a slow setup delays processing each time LOG.warn( - "DoFn.setup for {} ran for {} seconds.\nThis could cause Dataflow worker " - + "lease expire and failing the job. DoFn.Setup should not contain long " - + "running operations.", + "DoFn.setup for {} ran for {} seconds.\nSetup runs again on each new worker and " + + "for each thread that needs its own instance, so it delays processing every " + + "time. DoFn.Setup should not contain long running operations.", fn.getDoFn().getClass(), elapsed / 1_000); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkProgressUpdater.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkProgressUpdater.java index 12593222c4e1..c784da722751 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkProgressUpdater.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkProgressUpdater.java @@ -19,6 +19,7 @@ import com.google.api.client.util.Clock; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.util.Objects; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; @@ -315,7 +316,11 @@ private void reportProgress() { throw oom; } catch (InterruptedException e) { LOG.info("Cancelling workitem execution: {}", workString(), e); - worker.abort(); + // the batch worker reports progress while the executor is still being created, before + // there is a worker to abort + if (Objects.nonNull(worker)) { + worker.abort(); + } } catch (Throwable e) { LOG.error("Error reporting workitem progress update to Dataflow service: ", e); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java index ea797c901e52..2ca95837c20d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java @@ -44,6 +44,7 @@ import org.apache.beam.runners.dataflow.worker.util.common.worker.StubbedExecutor; import org.apache.beam.sdk.extensions.gcp.util.Transport; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.testing.ExpectedLogs; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.joda.time.Instant; @@ -74,6 +75,9 @@ public class DataflowWorkProgressUpdaterTest { @Rule public final ExpectedException thrown = ExpectedException.none(); + @Rule + public final ExpectedLogs expectedLogs = ExpectedLogs.none(DataflowWorkProgressUpdater.class); + private StubbedExecutor executor; private DataflowWorkProgressUpdater progressUpdater; @@ -198,6 +202,23 @@ public void workProgressAbortBeforeWorkerIsSetAbortsItOnceSet() throws Exception progressUpdater.stopReportingProgress(); } + @Test + public void workProgressWarnsWhileExecutorCreationIsSlow() throws Exception { + // the batch worker starts reporting while the executor is still being created + progressUpdater.setWorker(null); + when(workItemStatusClient.reportUpdate(isNull(DynamicSplitResult.class), isA(Duration.class))) + .thenReturn(generateServiceState(null, 1000)); + progressUpdater.startReportingProgress(); + executor.runNextRunnable(); + expectedLogs.verifyNotLogged("is still creating its executor"); + + clock.setTime(clock.currentTimeMillis() + Duration.standardMinutes(3).getMillis()); + executor.runNextRunnable(); + + expectedLogs.verifyWarn("is still creating its executor"); + progressUpdater.stopReportingProgress(); + } + /** Verifies that the update after a split is requested acknowledeges it. */ @Test public void workProgressSendsSplitResults() throws Exception { From 3f39a9051e979b466a2087ee3b7d3cb21b1b3ede Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 21:55:28 +0200 Subject: [PATCH 3/4] Add CHANGES.md entry for #34206 --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index 59253c35bf87..7e765d57cd6e 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -147,6 +147,7 @@ * (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no longer fail on rows with NaN values or on categorical columns with unused categories, and `nunique(dropna=False)` now counts NaN ([#40570](https://github.com/apache/beam/issues/40570)). * (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)). * (Go) Fixed generic registration (for example `register.Function1x1` or `register.DoFn1x1`) panicking on functions and methods that take or return an anonymous struct. Unnamed struct types are now only registered with the schema registry, since they have no type key for the runtime type registry ([#21789](https://github.com/apache/beam/issues/21789)). +* (Java) Dataflow legacy batch worker: a `DoFn.setup` that takes longer than 3 minutes no longer loses the work item lease and fails the job. While setup is still running, the worker logs a warning every few minutes ([#34206](https://github.com/apache/beam/issues/34206)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes From 56fc3b2cfbd141c814b92f3b1f353879e07d3879 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Wed, 28 Oct 2026 08:12:12 +0200 Subject: [PATCH 4/4] [Dataflow Java] Log hot key without key before batch executor exists --- .../worker/DataflowWorkProgressUpdater.java | 5 ++++- .../worker/DataflowWorkProgressUpdaterTest.java | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java index 94e479fda614..9097f16f05ac 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java @@ -182,7 +182,10 @@ protected void reportProgressHelper() throws Exception { // The key set the in BatchModeExecutionContext is only set in the GroupingShuffleReader // which is the correct key. The key is also translated into a Java object in the reader. - if (options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging")) { + // the execution context is null until the executor exists, so a report sent while + // DoFn.setup still runs logs the detection without the key instead of throwing + if ((options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging")) + && Objects.nonNull(workItemStatusClient.getExecutionContext())) { hotKeyLogger.logHotKeyDetection( hotKeyDetection.getUserStepName(), TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge()), diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java index 2ca95837c20d..73fbcc8ff926 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdaterTest.java @@ -174,6 +174,21 @@ public void workProgressLogsHotKeyDetectionNotEnabled() throws Exception { progressUpdater.stopReportingProgress(); } + @Test + public void workProgressLogsHotKeyDetectionBeforeWorkerIsSet() throws Exception { + // the batch worker starts reporting before the execution context holding the key exists + progressUpdater.setWorker(null); + when(workItemStatusClient.getExecutionContext()).thenReturn(null); + when(workItemStatusClient.reportUpdate(isNull(DynamicSplitResult.class), isA(Duration.class))) + .thenReturn(generateServiceState(null, 1000)); + progressUpdater.startReportingProgress(); + executor.runNextRunnable(); + + verify(hotKeyLogger, atLeastOnce()).logHotKeyDetection(STEP_ID, HOT_KEY_AGE); + + progressUpdater.stopReportingProgress(); + } + @Test public void workProgressAskedToAbortWorker() throws Exception { when(workItemStatusClient.reportUpdate(isNull(DynamicSplitResult.class), isA(Duration.class)))