From 14a6c0a2848fce5a5d557a5493165557a5a30b95 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 20:21:01 +0200 Subject: [PATCH 1/7] [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 fd559923efe1293604a6b3c890e403f7ddb7bf4d Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 21:51:53 +0200 Subject: [PATCH 2/7] [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 3739d322475546911009f99b098e9094d11c7882 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 21:55:28 +0200 Subject: [PATCH 3/7] Add CHANGES.md entry for #34206 --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index 44e7568cb8ee..7b12910bb21a 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -80,6 +80,7 @@ ## Bugfixes +* (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 5d3b3b7dc97a4c6ea01f91cc31eb5e82b0800bbd Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Wed, 28 Oct 2026 08:12:12 +0200 Subject: [PATCH 4/7] [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))) From 64b15625dfae2ddc70b1cc47c2532cc82efc50a8 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Wed, 28 Oct 2026 21:57:32 +0200 Subject: [PATCH 5/7] [Dataflow Java] Check nullness in batch worker progress reporting Remove the nullness suppression from BatchDataflowWorker, DataflowWorkProgressUpdater, WorkItemStatusClient and WorkProgressUpdater, now that the progress updater runs before the executor exists and the worker can be null. Mark the worker, the pending split result and the execution context as nullable, read fields into a local before null checks, and wrap cloud time and duration parsing in checkStateNotNull. Replace @NotThreadSafe on DataflowWorkProgressUpdater with a description of the two threads that use it. --- .../dataflow/worker/BatchDataflowWorker.java | 15 ++--- .../worker/DataflowWorkProgressUpdater.java | 62 +++++++++++-------- .../dataflow/worker/WorkItemStatusClient.java | 35 ++++++----- .../common/worker/WorkProgressUpdater.java | 27 ++++---- 4 files changed, 74 insertions(+), 65 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 1b411c679b3d..6a622a999492 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,7 +22,7 @@ import java.io.Closeable; import java.io.IOException; import java.util.Collection; -import java.util.Objects; +import java.util.Optional; import java.util.function.Function; import org.apache.beam.runners.core.metrics.ExecutionStateSampler; import org.apache.beam.runners.dataflow.DataflowRunner; @@ -42,10 +42,10 @@ import org.apache.beam.sdk.util.Weighted; import org.apache.beam.sdk.util.WeightedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Optional; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.graph.MutableNetwork; +import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -56,9 +56,6 @@ *

DataflowWorker presents one public interface, getAndPerformWork(), which uses the * WorkUnitClient to get work, execute it, and update the work. */ -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public class BatchDataflowWorker implements Closeable { private static final Logger LOG = LoggerFactory.getLogger(BatchDataflowWorker.class); @@ -134,7 +131,7 @@ public class BatchDataflowWorker implements Closeable { private final WorkerStatusPages statusPages; /** Periodic sender of debug information to the debug capture service. */ - private DebugCapture.Manager debugCaptureManager = null; + private DebugCapture.@Nullable Manager debugCaptureManager = null; protected BatchDataflowWorker( WorkUnitClient workUnitClient, @@ -216,7 +213,7 @@ private static Thread startMemoryMonitorThread(MemoryMonitor memoryMonitor) { */ public boolean getAndPerformWork() throws IOException { while (true) { - Optional work = Optional.fromJavaUtil(workUnitClient.getWorkItem()); + Optional work = workUnitClient.getWorkItem(); if (work.isPresent()) { WorkItemStatusClient statusProvider = new WorkItemStatusClient(workUnitClient, work.get()); return doWork(work.get(), statusProvider); @@ -296,7 +293,7 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr } workItemStatusClient.setWorker(worker, executionContext); - if (Objects.isNull(progressUpdater)) { + if (progressUpdater == null) { progressUpdater = new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker, options); executionStarted = true; @@ -313,7 +310,7 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr } 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) { + if (progressUpdater != null && !executionStarted) { try { progressUpdater.stopReportingProgress(); } catch (Exception stopExn) { 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 9097f16f05ac..ec901d2faf53 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,14 +26,14 @@ 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; import org.apache.beam.runners.dataflow.util.TimeUtil; import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkExecutor; import org.apache.beam.runners.dataflow.worker.util.common.worker.WorkProgressUpdater; +import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -41,11 +41,12 @@ /** * DataflowWorkProgressUpdater implements the WorkProgressUpdater interface for the Cloud Dataflow * system. + * + *

Two threads use it. The thread running the work item starts and stops reporting and calls + * {@link #setWorker} once the executor exists, while the progress reporting thread sends the + * updates and receives aborts from the service. An abort can arrive while the executor is still + * being created, so handing that abort over to {@link #setWorker} is guarded by a lock. */ -@NotThreadSafe -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public class DataflowWorkProgressUpdater extends WorkProgressUpdater { private static final Logger LOG = LoggerFactory.getLogger(DataflowWorkProgressUpdater.class); @@ -85,7 +86,7 @@ public class DataflowWorkProgressUpdater extends WorkProgressUpdater { public DataflowWorkProgressUpdater( WorkItemStatusClient workItemStatusClient, WorkItem workItem, - WorkExecutor worker, + @Nullable WorkExecutor worker, DataflowPipelineOptions options) { super(worker, Integer.MAX_VALUE); this.workItemStatusClient = workItemStatusClient; @@ -105,7 +106,7 @@ public DataflowWorkProgressUpdater( DataflowWorkProgressUpdater( WorkItemStatusClient workItemStatusClient, WorkItem workItem, - WorkExecutor worker, + @Nullable WorkExecutor worker, ScheduledExecutorService executor, Clock clock, HotKeyLogger hotKeyLogger, @@ -144,7 +145,8 @@ protected long getWorkUnitLeaseExpirationTimestamp() { @Override protected long getWorkUnitSuggestedReportingInterval() { - return fromCloudDuration(workItem.getReportStatusInterval()).getMillis(); + return Preconditions.checkStateNotNull(fromCloudDuration(workItem.getReportStatusInterval())) + .getMillis(); } @Override @@ -153,7 +155,7 @@ protected void reportProgressHelper() throws Exception { LOG.info("Service already asked to abort work item, not reporting ignored progress."); return; } - if (Objects.isNull(worker)) { + if (worker == null) { warnIfSetupIsSlow(); } WorkItemServiceState result = @@ -170,7 +172,7 @@ protected void reportProgressHelper() throws Exception { currentWorker = worker; } // the executor may still be under construction, in which case setWorker aborts it - if (Objects.nonNull(currentWorker)) { + if (currentWorker != null) { currentWorker.abort(); } return; @@ -179,21 +181,25 @@ protected void reportProgressHelper() throws Exception { if (result.getHotKeyDetection() != null && result.getHotKeyDetection().getUserStepName() != null) { HotKeyDetection hotKeyDetection = result.getHotKeyDetection(); + Duration hotKeyAge = + Preconditions.checkStateNotNull( + TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge())); // 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. - // 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 + // the execution context is null until the executor exists, and the key is null until the + // reader reads one, so a report sent while DoFn.setup still runs logs the detection + // without the key instead of throwing + Object hotKey = null; + BatchModeExecutionContext executionContext = workItemStatusClient.getExecutionContext(); if ((options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging")) - && Objects.nonNull(workItemStatusClient.getExecutionContext())) { - hotKeyLogger.logHotKeyDetection( - hotKeyDetection.getUserStepName(), - TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge()), - workItemStatusClient.getExecutionContext().getKey()); + && executionContext != null) { + hotKey = executionContext.getKey(); + } + if (hotKey != null) { + hotKeyLogger.logHotKeyDetection(hotKeyDetection.getUserStepName(), hotKeyAge, hotKey); } else { - hotKeyLogger.logHotKeyDetection( - hotKeyDetection.getUserStepName(), - TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge())); + hotKeyLogger.logHotKeyDetection(hotKeyDetection.getUserStepName(), hotKeyAge); } } @@ -202,15 +208,17 @@ protected void reportProgressHelper() throws Exception { progressReportIntervalMs = nextProgressReportInterval( - fromCloudDuration(result.getReportStatusInterval()).getMillis(), + Preconditions.checkStateNotNull(fromCloudDuration(result.getReportStatusInterval())) + .getMillis(), leaseRemainingTime(getLeaseExpirationTimestamp(result))); ApproximateSplitRequest suggestedStopPoint = result.getSplitRequest(); + WorkExecutor currentWorker = worker; // there is nothing to split until the executor exists, the service asks again later - if (suggestedStopPoint != null && Objects.nonNull(worker)) { + if (suggestedStopPoint != null && currentWorker != null) { LOG.info("Proposing dynamic split of work unit {} at {}", workString(), suggestedStopPoint); dynamicSplitResultToReport = - worker.requestDynamicSplit( + currentWorker.requestDynamicSplit( SourceTranslationUtils.toDynamicSplitRequest(suggestedStopPoint)); } } @@ -235,11 +243,13 @@ private void warnIfSetupIsSlow() { /** Returns the given work unit's lease expiration timestamp. */ private long getLeaseExpirationTimestamp(WorkItem workItem) { - return fromCloudTime(workItem.getLeaseExpireTime()).getMillis(); + return Preconditions.checkStateNotNull(fromCloudTime(workItem.getLeaseExpireTime())) + .getMillis(); } /** Returns the given work unit service state lease expiration timestamp. */ private long getLeaseExpirationTimestamp(WorkItemServiceState workItemServiceState) { - return fromCloudTime(workItemServiceState.getLeaseExpireTime()).getMillis(); + return Preconditions.checkStateNotNull(fromCloudTime(workItemServiceState.getLeaseExpireTime())) + .getMillis(); } } 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 cf9659f58ad9..d0cde832a9aa 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,7 +36,6 @@ 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; @@ -63,9 +62,6 @@ */ // Very likely real potential for bugs - https://github.com/apache/beam/issues/19270 @SuppressFBWarnings("IS2_INCONSISTENT_SYNC") -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public class WorkItemStatusClient { private static final Logger LOG = LoggerFactory.getLogger(WorkItemStatusClient.class); @@ -74,7 +70,7 @@ public class WorkItemStatusClient { private @Nullable DataflowWorkExecutor worker; private Long nextReportIndex; - private transient String uniqueWorkId = null; + private transient @Nullable String uniqueWorkId = null; private boolean finalStateSent = false; private boolean wasAskedToAbort = false; @@ -145,7 +141,12 @@ public synchronized void setWorker( "{}: Uncaught exception occurred during work unit execution. This will be retried.", logPrefix, t); - error.setMessage(DataflowWorkerLoggingHandler.formatException(t)); + // t is null only for a UserCodeException without a cause, and then the error is sent + // without a message + String formattedException = DataflowWorkerLoggingHandler.formatException(t); + if (formattedException != null) { + error.setMessage(formattedException); + } } status.setErrors(ImmutableList.of(error)); @@ -200,7 +201,7 @@ public synchronized void setWorker( return execute(status); } - private static boolean isOutOfMemoryError(Throwable t) { + private static boolean isOutOfMemoryError(@Nullable Throwable t) { while (t != null) { if (t instanceof OutOfMemoryError) { return true; @@ -210,7 +211,7 @@ private static boolean isOutOfMemoryError(Throwable t) { return false; } - private static boolean isReadLoopAbortedError(Throwable t) { + private static boolean isReadLoopAbortedError(@Nullable Throwable t) { while (t != null) { if (t instanceof ReadOperation.ReadLoopAbortedException) { return true; @@ -247,7 +248,7 @@ private static boolean isReadLoopAbortedError(Throwable t) { @VisibleForTesting synchronized void populateSplitResult( - WorkItemStatus status, DynamicSplitResult dynamicSplitResult) { + WorkItemStatus status, @Nullable DynamicSplitResult dynamicSplitResult) { if (dynamicSplitResult instanceof NativeReader.DynamicSplitResultWithPosition) { NativeReader.DynamicSplitResultWithPosition asPosition = (NativeReader.DynamicSplitResultWithPosition) dynamicSplitResult; @@ -266,7 +267,7 @@ 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)) { + if (worker == null) { return; } @@ -279,8 +280,9 @@ synchronized void populateProgress(WorkItemStatus status) throws Exception { @VisibleForTesting synchronized void populateMetricUpdates(WorkItemStatus status) { List updates = new ArrayList<>(); - if (executionContext != null && executionContext.getExecutionStateTracker() != null) { - ExecutionStateTracker tracker = executionContext.getExecutionStateTracker(); + BatchModeExecutionContext currentContext = executionContext; + if (currentContext != null && currentContext.getExecutionStateTracker() != null) { + ExecutionStateTracker tracker = currentContext.getExecutionStateTracker(); MetricUpdate update = new MetricUpdate(); update.setKind("internal"); @@ -319,7 +321,8 @@ private synchronized WorkItemStatus createStatusUpdate(boolean isFinal) { @VisibleForTesting synchronized void populateCounterUpdates(WorkItemStatus status) { - if (worker == null) { + DataflowWorkExecutor currentWorker = worker; + if (currentWorker == null) { return; } @@ -337,7 +340,7 @@ synchronized void populateCounterUpdates(WorkItemStatus status) { x); // Output counters - extractCounters(worker.getOutputCounters()).forEach(appendCounterUpdate); + extractCounters(currentWorker.getOutputCounters()).forEach(appendCounterUpdate); // User metrics reported in Worker extractMetrics(isFinalUpdate).forEach(appendCounterUpdate); @@ -349,7 +352,7 @@ synchronized void populateCounterUpdates(WorkItemStatus status) { // This includes all different kinds of metrics coming from SDK. // Keep in mind that these metrics might contain different types of counter names: // i.e. structuredNameAndMetadata and nameAndKind - worker.extractMetricUpdates().forEach(appendCounterUpdate); + currentWorker.extractMetricUpdates().forEach(appendCounterUpdate); status.setCounterUpdates(ImmutableList.copyOf(counterUpdatesMap.values())); } @@ -404,7 +407,7 @@ synchronized void commitMetrics() { executionContext.commitMetricUpdates(); } - public BatchModeExecutionContext getExecutionContext() { + public @Nullable BatchModeExecutionContext getExecutionContext() { return this.executionContext; } } 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 c784da722751..048668fae2c4 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,7 +19,6 @@ 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; @@ -27,6 +26,7 @@ import javax.annotation.concurrent.NotThreadSafe; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.ThreadFactoryBuilder; +import org.checkerframework.checker.nullness.qual.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -43,9 +43,6 @@ // Very likely real potential for bugs - https://github.com/apache/beam/issues/19274 @SuppressFBWarnings("JLM_JSR166_UTILCONCURRENT_MONITORENTER") @NotThreadSafe -@SuppressWarnings({ - "nullness" // TODO(https://github.com/apache/beam/issues/20497) -}) public abstract class WorkProgressUpdater { private static final Logger LOG = LoggerFactory.getLogger(WorkProgressUpdater.class); @@ -68,10 +65,11 @@ public abstract class WorkProgressUpdater { private static final long DEFAULT_MAX_REPORTING_INTERVAL_MILLIS = 10 * 60 * 1000; /** - * Worker providing the work progress updates. This is a volatile variable because the worker - * thread sets it while the progress updater thread reads it. + * Worker providing the work progress updates, or {@code null} while the batch worker is still + * creating it. This is a volatile variable because the worker thread sets it while the progress + * updater thread reads it. */ - protected volatile WorkExecutor worker = null; + protected volatile @Nullable WorkExecutor worker = null; /** Requested periodic checkpoint period. */ private final int checkpointPeriodSec; @@ -112,13 +110,13 @@ protected enum CheckpointState { * update, or {@code null} if there is nothing to report (if no dynamic split happened since the * last progress update). */ - protected NativeReader.DynamicSplitResult dynamicSplitResultToReport; + protected NativeReader.@Nullable DynamicSplitResult dynamicSplitResultToReport; /** * @param checkpointPeriodSec the desired amount of time in seconds between periodic checkpoints; * if no periodic checkpoints are desired then pass {@link Integer#MAX_VALUE} */ - public WorkProgressUpdater(WorkExecutor worker, int checkpointPeriodSec) { + public WorkProgressUpdater(@Nullable WorkExecutor worker, int checkpointPeriodSec) { this( worker, checkpointPeriodSec, @@ -138,7 +136,7 @@ public WorkProgressUpdater(WorkExecutor worker, int checkpointPeriodSec) { */ @VisibleForTesting protected WorkProgressUpdater( - WorkExecutor worker, + @Nullable WorkExecutor worker, int checkpointPeriodSec, ScheduledExecutorService executor, Clock clock) { @@ -287,10 +285,11 @@ private void checkForPeriodicCheckpoint() { */ @GuardedBy("executor") protected boolean tryCheckpointIfNeeded() { - if (checkpointState == CheckpointState.CHECKPOINT_REQUESTED && worker != null) { + WorkExecutor currentWorker = worker; + if (checkpointState == CheckpointState.CHECKPOINT_REQUESTED && currentWorker != null) { LOG.debug("Trying to checkpoint for work item {}.", workString()); try { - NativeReader.DynamicSplitResult checkpointPos = worker.requestCheckpoint(); + NativeReader.DynamicSplitResult checkpointPos = currentWorker.requestCheckpoint(); if (checkpointPos != null) { LOG.debug("Successful checkpoint for work item {} at {}.", workString(), checkpointPos); dynamicSplitResultToReport = checkpointPos; @@ -318,7 +317,7 @@ private void reportProgress() { LOG.info("Cancelling workitem execution: {}", workString(), e); // the batch worker reports progress while the executor is still being created, before // there is a worker to abort - if (Objects.nonNull(worker)) { + if (worker != null) { worker.abort(); } } catch (Throwable e) { @@ -342,7 +341,7 @@ protected long leaseRemainingTime(long leaseExpirationTimestamp) { } @VisibleForTesting - public NativeReader.DynamicSplitResult getDynamicSplitResultToReport() { + public NativeReader.@Nullable DynamicSplitResult getDynamicSplitResultToReport() { return dynamicSplitResultToReport; } From a5027e4ad672eb0137be79285987abdeb374126e Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Wed, 28 Oct 2026 21:57:35 +0200 Subject: [PATCH 6/7] Trigger legacy Dataflow worker postcommits --- .github/trigger_files/beam_PostCommit_Java_DataflowV1.json | 2 +- .../beam_PostCommit_Java_ValidatesRunner_Dataflow.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json b/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json index dc9b909a926c..44de32fa860f 100644 --- a/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json +++ b/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json @@ -4,6 +4,6 @@ "https://github.com/apache/beam/pull/34902": "Introducing OutputBuilder", "https://github.com/apache/beam/pull/35177": "Introducing WindowedValueReceiver to runners", "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 10, + "modification": 11, "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface" } diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json index f7bf5e8a08a0..4144ed0df4dc 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run!", - "modification": 14, + "modification": 15, } From 5b8f68d4f5e061c373679a35d4ff043df9d54755 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Thu, 29 Oct 2026 09:57:11 +0200 Subject: [PATCH 7/7] Re-run legacy Dataflow worker postcommits --- .github/trigger_files/beam_PostCommit_Java_DataflowV1.json | 2 +- .../beam_PostCommit_Java_ValidatesRunner_Dataflow.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json b/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json index 44de32fa860f..1e954ee5cd26 100644 --- a/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json +++ b/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json @@ -4,6 +4,6 @@ "https://github.com/apache/beam/pull/34902": "Introducing OutputBuilder", "https://github.com/apache/beam/pull/35177": "Introducing WindowedValueReceiver to runners", "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 11, + "modification": 12, "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface" } diff --git a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json index 4144ed0df4dc..d46523902fbb 100644 --- a/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Java_ValidatesRunner_Dataflow.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run!", - "modification": 15, + "modification": 16, }