[Dataflow Java] Extend work item lease while DoFn.setup runs in batch worker

CHANGES.md

@@ -143,6 +143,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

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();

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;
       }
 
@@ -154,7 +203,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 +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();

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);
       }

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));

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);
     }

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<ParallelInstruction>())
+                    .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();