diff --git a/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json b/.github/trigger_files/beam_PostCommit_Java_DataflowV1.json index dc9b909a926c..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": 10, + "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 f7bf5e8a08a0..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": 14, + "modification": 16, } 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 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..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,6 +22,7 @@ import java.io.Closeable; import java.io.IOException; import java.util.Collection; +import java.util.Optional; import java.util.function.Function; import org.apache.beam.runners.core.metrics.ExecutionStateSampler; import org.apache.beam.runners.dataflow.DataflowRunner; @@ -41,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; @@ -55,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); @@ -133,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, @@ -215,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); @@ -233,6 +231,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 +264,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 +293,30 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr } workItemStatusClient.setWorker(worker, executionContext); - DataflowWorkProgressUpdater progressUpdater = - new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker, options); - executeWork(worker, progressUpdater); + if (progressUpdater == null) { + 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 (progressUpdater != null && !executionStarted) { + try { + progressUpdater.stopReportingProgress(); + } catch (Exception stopExn) { + e.addSuppressed(stopExn); + } + } workItemStatusClient.reportError(e); return false; } finally { @@ -314,6 +337,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..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 @@ -27,12 +27,13 @@ import com.google.api.services.dataflow.model.WorkItem; import com.google.api.services.dataflow.model.WorkItemServiceState; 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; @@ -40,15 +41,25 @@ /** * 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); + /** + * 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,18 +69,32 @@ 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, WorkItem workItem, - WorkExecutor worker, + @Nullable WorkExecutor worker, DataflowPipelineOptions options) { super(worker, Integer.MAX_VALUE); this.workItemStatusClient = workItemStatusClient; this.workItem = workItem; this.hotKeyLogger = new HotKeyLogger(); this.options = options; + this.createdAtMillis = clock.currentTimeMillis(); + this.nextSlowSetupWarningAtMillis = createdAtMillis + SLOW_SETUP_FIRST_WARNING_MILLIS; } /** @@ -81,7 +106,7 @@ public DataflowWorkProgressUpdater( DataflowWorkProgressUpdater( WorkItemStatusClient workItemStatusClient, WorkItem workItem, - WorkExecutor worker, + @Nullable WorkExecutor worker, ScheduledExecutorService executor, Clock clock, HotKeyLogger hotKeyLogger, @@ -91,6 +116,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 @@ -105,7 +145,8 @@ protected long getWorkUnitLeaseExpirationTimestamp() { @Override protected long getWorkUnitSuggestedReportingInterval() { - return fromCloudDuration(workItem.getReportStatusInterval()).getMillis(); + return Preconditions.checkStateNotNull(fromCloudDuration(workItem.getReportStatusInterval())) + .getMillis(); } @Override @@ -114,6 +155,9 @@ protected void reportProgressHelper() throws Exception { LOG.info("Service already asked to abort work item, not reporting ignored progress."); return; } + if (worker == null) { + warnIfSetupIsSlow(); + } WorkItemServiceState result = workItemStatusClient.reportUpdate( dynamicSplitResultToReport, Duration.millis(requestedLeaseDurationMs)); @@ -122,26 +166,40 @@ 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 (currentWorker != null) { + currentWorker.abort(); + } return; } 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. - if (options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging")) { - hotKeyLogger.logHotKeyDetection( - hotKeyDetection.getUserStepName(), - TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge()), - workItemStatusClient.getExecutionContext().getKey()); + // 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")) + && 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); } } @@ -150,26 +208,48 @@ protected void reportProgressHelper() throws Exception { progressReportIntervalMs = nextProgressReportInterval( - fromCloudDuration(result.getReportStatusInterval()).getMillis(), + Preconditions.checkStateNotNull(fromCloudDuration(result.getReportStatusInterval())) + .getMillis(), leaseRemainingTime(getLeaseExpirationTimestamp(result))); ApproximateSplitRequest suggestedStopPoint = result.getSplitRequest(); - if (suggestedStopPoint != null) { + WorkExecutor currentWorker = worker; + // there is nothing to split until the executor exists, the service asks again later + if (suggestedStopPoint != null && currentWorker != null) { LOG.info("Proposing dynamic split of work unit {} at {}", workString(), suggestedStopPoint); dynamicSplitResultToReport = - worker.requestDynamicSplit( + currentWorker.requestDynamicSplit( SourceTranslationUtils.toDynamicSplitRequest(suggestedStopPoint)); } } } + 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(); + 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/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..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 @@ -62,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); @@ -73,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; @@ -144,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)); @@ -176,11 +178,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) { @@ -196,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; @@ -206,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; @@ -243,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; @@ -261,6 +266,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 (worker == null) { + return; + } + Progress progress = worker.getWorkerProgress(); if (progress != null) { status.setReportedProgress(SourceTranslationUtils.readerProgressToCloudProgress(progress)); @@ -270,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"); @@ -310,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; } @@ -328,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); @@ -340,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())); } @@ -395,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 12593222c4e1..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 @@ -26,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; @@ -42,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); @@ -67,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; @@ -111,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, @@ -137,7 +136,7 @@ public WorkProgressUpdater(WorkExecutor worker, int checkpointPeriodSec) { */ @VisibleForTesting protected WorkProgressUpdater( - WorkExecutor worker, + @Nullable WorkExecutor worker, int checkpointPeriodSec, ScheduledExecutorService executor, Clock clock) { @@ -286,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; @@ -315,7 +315,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 (worker != null) { + worker.abort(); + } } catch (Throwable e) { LOG.error("Error reporting workitem progress update to Dataflow service: ", e); } @@ -337,7 +341,7 @@ protected long leaseRemainingTime(long leaseExpirationTimestamp) { } @VisibleForTesting - public NativeReader.DynamicSplitResult getDynamicSplitResultToReport() { + public NativeReader.@Nullable DynamicSplitResult getDynamicSplitResultToReport() { return dynamicSplitResultToReport; } 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 {