[mxtymoshyk]

Addresses #34206.

Opened as a draft to get feedback on the approach.

Problem

In the legacy batch worker, BatchDataflowWorker.doWork builds the executor with mapTaskExecutorFactory.create(), which runs every DoFn.setup. Progress reporting starts only after that, so nothing extends the lease while setup runs. If setup takes longer than the 3 minute lease, the service fails the work item, and after 4 failures the job fails.

Change

• BatchDataflowWorker: start the progress updater before create(), then hand it the executor. If create() throws, stop the updater before reporting the error.
• WorkItemStatusClient: an update sent before setWorker carries only the requested lease, with no progress, counters or metrics. The checkState that rejected it is removed.
• DataflowWorkProgressUpdater:
  • An abort received during setup is remembered and applied in setWorker. MapTaskExecutor keeps the abort flag, so the read loop stops before the first record and the work item is reported as failed.
  • Split requests are ignored until the executor exists.
  • While the executor is still being created, it logs a warning after 3 minutes and every 5 minutes after that: Work item ... is still creating its executor after N seconds. The existing setup warning in DoFnInstanceManagers only fires once setup returns, so without this a stuck setup would log nothing.
• WorkProgressUpdater: if the reporting thread is interrupted before the executor exists, skip worker.abort() instead of throwing a NullPointerException.
• DoFnInstanceManagers: the slow setup warning no longer says the lease will expire and fail the job. It now says setup runs again on each new worker and for each thread that needs its own instance.

Tests

• BatchDataflowWorkerTest.testProgressReportedWhileExecutorIsBeingCreated: blocks create() until a status update arrives. It times out without the fix and passes with it.
• DataflowWorkProgressUpdaterTest.workProgressAbortBeforeWorkerIsSetAbortsItOnceSet
• DataflowWorkProgressUpdaterTest.workProgressWarnsWhileExecutorCreationIsSlow
• WorkItemStatusClientTest.reportUpdateBeforeSetWorkerOnlyExtendsLease replaces the test that expected the old IllegalStateException.

:runners:google-cloud-dataflow-java:worker:check passes. The first 3 new tests passed 20 runs in a row.

Checked on Dataflow

All runs used the legacy worker (disable_runner_v2, --dataflowWorkerJar) with Create.of(1, 2, 3) into a ParDo whose @Setup sleeps.

• Batch, 5 minute setup, worker built from master: right after setup the first status update came back with Service asked worker to abort with status: {"code":1}, and the work item was retried. The job still finished, because the retry ran on the same worker and reused the DoFn that had already been set up. A job fails when the retries also have to run setup, which is what the issue reports.
• Batch, 5 minute setup, this branch: the same work item finished on its first attempt (Finished processing stage F20 with 0 errors in 303.569 seconds). The service accepted the lease-only updates sent during setup.
• Batch, three independent branches with --numberOfWorkerHarnessThreads=4: three setups ran at the same time on one worker. Each logged the new warning once at about 182 seconds and finished in about 300 seconds, with no abort.
• Batch, 2 hour setup, cancelled after 12 minutes: the new warning appeared at 184 and 486 seconds, the cancel reached the worker as an abort, and the job reached Cancelled 44 seconds later.
• Streaming (no Streaming Engine), 5 minute setup: the job finished with no abort or error. This change does not touch the streaming worker. I ran it to check whether the old warning text was accurate there.

Design choices

• No cap on how long setup may hold the lease. execute() has no cap either, and a hung processElement renews its lease forever today. A cap would also not stop the stuck thread; the service would retry the item while the original setup keeps running. A hung setup used to fail the job after about 12 minutes and now runs until someone cancels it, so the repeating warning is there to make it visible.
• An abort during setup takes effect when create() returns, not sooner. Interrupting user code in setup can leave a half-built client in the pooled DoFn that later work items reuse. MapTaskExecutor.abort() does not interrupt the execution thread either.

[codecov]

Codecov Report

❌ Patch coverage is 78.84615% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.00%. Comparing base (2213e54) to head (82637cd).
⚠️ Report is 43 commits behind head on master.

Files with missing lines	Patch %	Lines
...m/runners/dataflow/worker/BatchDataflowWorker.java	50.00%	7 Missing and 2 partials ⚠️
...s/dataflow/worker/DataflowWorkProgressUpdater.java	96.66%	0 Missing and 1 partial ⚠️
...worker/util/common/worker/WorkProgressUpdater.java	50.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@              Coverage Diff              @@
##             master   #40850       +/-   ##
=============================================
+ Coverage          0   76.00%   +76.00%     
- Complexity        0     2283     +2283     
=============================================
  Files             0      259      +259     
  Lines             0    14421    +14421     
  Branches          0     1485     +1485     
=============================================
+ Hits              0    10960    +10960     
- Misses            0     2767     +2767     
- Partials          0      694      +694     

Flag	Coverage Δ
java	76.00% <78.84%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.