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 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..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 @@ -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; @@ -49,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. */ @@ -58,7 +68,19 @@ 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(); + + /** 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, @@ -70,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; } /** @@ -91,6 +115,21 @@ public DataflowWorkProgressUpdater( this.workItem = workItem; this.hotKeyLogger = hotKeyLogger; this.options = options; + this.createdAtMillis = clock.currentTimeMillis(); + this.nextSlowSetupWarningAtMillis = createdAtMillis + SLOW_SETUP_FIRST_WARNING_MILLIS; + } + + @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 @@ -114,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)); @@ -122,8 +164,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; } @@ -133,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()), @@ -154,7 +206,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( @@ -163,6 +216,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/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/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/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..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 @@ -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; @@ -170,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))) @@ -182,6 +201,39 @@ 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(); + } + + @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 { 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 {