Keep progress renewal heartbeats active through bundle teardown

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java

@@ -288,8 +288,7 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr
 
       DataflowWorkProgressUpdater progressUpdater =
           new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker, options);
-      executeWork(worker, progressUpdater);
-      workItemStatusClient.reportSuccess();
+      executeWorkAndReportSuccess(worker, progressUpdater, workItemStatusClient);
       return true;
     } catch (OutOfMemoryError oom) {
       throw oom;
@@ -311,6 +310,7 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr
   }
 
   /** Executes the work and report progress. For testing only. */
+  @VisibleForTesting
   void executeWork(DataflowWorkExecutor worker, DataflowWorkProgressUpdater progressUpdater)
       throws Exception {
     progressUpdater.startReportingProgress();
@@ -324,6 +324,37 @@ void executeWork(DataflowWorkExecutor worker, DataflowWorkProgressUpdater progre
     }
   }
 
+  /**
+   * Executes the work, reports any unreported dynamic split, and reports final success while
+   * keeping progress reporting (and lease renewal) active.
+   *
+   * <p>By keeping the {@link DataflowWorkProgressUpdater} active until {@link
+   * WorkItemStatusClient#reportSuccess} completes, lease renewal heartbeats continue even if bundle
+   * completion or status reporting encounters delays, preventing lease expiration.
+   */
+  @VisibleForTesting
+  void executeWorkAndReportSuccess(
+      DataflowWorkExecutor worker,
+      DataflowWorkProgressUpdater progressUpdater,
+      WorkItemStatusClient workItemStatusClient)
+      throws Exception {
+    progressUpdater.startReportingProgress();
+    try {
+      // Blocks while executing the work.
+      worker.execute();
+      // Ensure any pending dynamic split is reported before sending the final success status.
+      progressUpdater.reportUnreportedSplit();
+      // Switch progress reporting to lightweight lease renewal pings now that compute is complete.
+      progressUpdater.setLeaseRenewalOnly(true);
+      // Report success while progress reporting (and lease renewal) is still active.
+      workItemStatusClient.reportSuccess();
+    } finally {
+      // Stop progress reporting only after work execution and reportSuccess have completed
+      // (or failed).
+      progressUpdater.stopReportingProgress();
+    }
+  }
+
   /** Runs the status server to report worker health on the specified port. */
   public void startStatusServer() {
     statusPages.start();

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/DataflowWorkProgressUpdater.java

@@ -60,6 +60,8 @@ public class DataflowWorkProgressUpdater extends WorkProgressUpdater {
 
   private boolean wasAskedToAbort = false;
 
+  private volatile boolean leaseRenewalOnly = false;
+
   public DataflowWorkProgressUpdater(
       WorkItemStatusClient workItemStatusClient,
       WorkItem workItem,
@@ -108,58 +110,102 @@ protected long getWorkUnitSuggestedReportingInterval() {
     return fromCloudDuration(workItem.getReportStatusInterval()).getMillis();
   }
 
+  /**
+   * Switches progress reporting to lightweight lease renewal pings without extracting counters,
+   * metrics, or progress.
+   */
+  public void setLeaseRenewalOnly(boolean leaseRenewalOnly) {
+    this.leaseRenewalOnly = leaseRenewalOnly;
+  }
+
+  @VisibleForTesting
+  public boolean isLeaseRenewalOnly() {
+    return leaseRenewalOnly;
+  }
+
   @Override
   protected void reportProgressHelper() throws Exception {
     if (wasAskedToAbort) {
       LOG.info("Service already asked to abort work item, not reporting ignored progress.");
       return;
     }
+    if (workItemStatusClient.isFinalStateSent()) {
+      LOG.debug(
+          "Final state already sent for work item {}, skipping progress report.", workString());
+      return;
+    }
+
+    if (leaseRenewalOnly && dynamicSplitResultToReport == null) {
+      reportLeasePing();
+      return;
+    }
+
     WorkItemServiceState result =
         workItemStatusClient.reportUpdate(
             dynamicSplitResultToReport, Duration.millis(requestedLeaseDurationMs));
 
     if (result != null) {
-      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();
-        return;
-      }
+      handleServiceState(result);
+    }
+  }
+
+  /**
+   * Reports a lightweight lease renewal heartbeat to the worker service without extracting counters
+   * or progress.
+   */
+  @VisibleForTesting
+  void reportLeasePing() throws Exception {
+    if (wasAskedToAbort || workItemStatusClient.isFinalStateSent()) {
+      return;
+    }
+    WorkItemServiceState result =
+        workItemStatusClient.reportLeasePing(Duration.millis(requestedLeaseDurationMs));
+    if (result != null) {
+      handleServiceState(result);
+    }
+  }
+
+  private void handleServiceState(WorkItemServiceState result) 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();
+      return;
+    }
 
-      if (result.getHotKeyDetection() != null
-          && result.getHotKeyDetection().getUserStepName() != null) {
-        HotKeyDetection hotKeyDetection = result.getHotKeyDetection();
-
-        // 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());
-        } else {
-          hotKeyLogger.logHotKeyDetection(
-              hotKeyDetection.getUserStepName(),
-              TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge()));
-        }
+    if (result.getHotKeyDetection() != null
+        && result.getHotKeyDetection().getUserStepName() != null) {
+      HotKeyDetection hotKeyDetection = result.getHotKeyDetection();
+
+      // The key set 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());
+      } else {
+        hotKeyLogger.logHotKeyDetection(
+            hotKeyDetection.getUserStepName(),
+            TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge()));
       }
+    }
 
-      // Resets state after a successful progress report.
-      dynamicSplitResultToReport = null;
+    // Resets state after a successful progress report.
+    dynamicSplitResultToReport = null;
 
-      progressReportIntervalMs =
-          nextProgressReportInterval(
-              fromCloudDuration(result.getReportStatusInterval()).getMillis(),
-              leaseRemainingTime(getLeaseExpirationTimestamp(result)));
+    progressReportIntervalMs =
+        nextProgressReportInterval(
+            fromCloudDuration(result.getReportStatusInterval()).getMillis(),
+            leaseRemainingTime(getLeaseExpirationTimestamp(result)));
 
-      ApproximateSplitRequest suggestedStopPoint = result.getSplitRequest();
-      if (suggestedStopPoint != null) {
-        LOG.info("Proposing dynamic split of work unit {} at {}", workString(), suggestedStopPoint);
-        dynamicSplitResultToReport =
-            worker.requestDynamicSplit(
-                SourceTranslationUtils.toDynamicSplitRequest(suggestedStopPoint));
-      }
+    ApproximateSplitRequest suggestedStopPoint = result.getSplitRequest();
+    if (suggestedStopPoint != null) {
+      LOG.info("Proposing dynamic split of work unit {} at {}", workString(), suggestedStopPoint);
+      dynamicSplitResultToReport =
+          worker.requestDynamicSplit(
+              SourceTranslationUtils.toDynamicSplitRequest(suggestedStopPoint));
     }
   }
 

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkItemStatusClient.java

@@ -17,6 +17,8 @@
  */
 package org.apache.beam.runners.dataflow.worker;
 
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
@@ -74,11 +76,16 @@ public class WorkItemStatusClient {
   private Long nextReportIndex;
 
   private transient String uniqueWorkId = null;
-  private boolean finalStateSent = false;
-  private boolean wasAskedToAbort = false;
+  private volatile boolean finalStateSent = false;
+  private volatile boolean wasAskedToAbort = false;
 
   private @Nullable BatchModeExecutionContext executionContext;
 
+  /** Returns whether a final completion status (success or error) has already been sent. */
+  public boolean isFinalStateSent() {
+    return finalStateSent;
+  }
+
   /**
    * Construct a partly-initialized {@link WorkItemStatusClient}. Once the {@link
    * DataflowWorkExecutor worker} has been instantiated initialization should be completed by
@@ -180,22 +187,62 @@ public synchronized void setWorker(
   public synchronized @Nullable WorkItemServiceState reportUpdate(
       @Nullable DynamicSplitResult dynamicSplitResult, Duration requestedLeaseDuration)
       throws Exception {
+    return reportUpdate(dynamicSplitResult, requestedLeaseDuration, true);
+  }
+
+  /**
+   * Return the {@link WorkItemServiceState} resulting from sending a progress update, optionally
+   * including counters and metrics.
+   */
+  public synchronized @Nullable WorkItemServiceState reportUpdate(
+      @Nullable DynamicSplitResult dynamicSplitResult,
+      Duration requestedLeaseDuration,
+      boolean includeCounters)
+      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 (finalStateSent) {
+      LOG.debug(
+          "Final state already sent for work item {}, skipping progress update.", uniqueWorkId());
+      return null;
+    }
     if (wasAskedToAbort) {
       LOG.info("Service already asked to abort work item, not reporting ignored progress.");
       return null;
     }
 
-    WorkItemStatus status = createStatusUpdate(false);
+    WorkItemStatus status = createStatusUpdate(false, includeCounters);
     status.setRequestedLeaseDuration(TimeUtil.toCloudDuration(requestedLeaseDuration));
     populateProgress(status);
     populateSplitResult(status, dynamicSplitResult);
 
     return execute(status);
   }
 
+  /**
+   * Sends a lightweight lease renewal heartbeat without counter extraction or progress sampling.
+   *
+   * <p>This can be used during bundle finalization, teardown, or when the worker is otherwise busy,
+   * to keep the lease active on the Dataflow service without incurring serialization overhead.
+   */
+  public synchronized @Nullable WorkItemServiceState reportLeasePing(
+      Duration requestedLeaseDuration) throws Exception {
+    checkStateNotNull(worker, "setWorker should be called before reportLeasePing");
+    checkArgumentNotNull(requestedLeaseDuration, "requestLeaseDuration must be non-null");
+    if (finalStateSent) {
+      LOG.debug("Final state already sent for work item {}, skipping lease ping.", uniqueWorkId());
+      return null;
+    }
+    if (wasAskedToAbort) {
+      LOG.info("Service already asked to abort work item, not reporting ignored progress.");
+      return null;
+    }
+
+    WorkItemStatus status = createStatusUpdate(false, false);
+    status.setRequestedLeaseDuration(TimeUtil.toCloudDuration(requestedLeaseDuration));
+    return execute(status);
+  }
+
   private static boolean isOutOfMemoryError(Throwable t) {
     while (t != null) {
       if (t instanceof OutOfMemoryError) {
@@ -218,6 +265,15 @@ private static boolean isReadLoopAbortedError(Throwable t) {
 
   private synchronized @Nullable WorkItemServiceState execute(WorkItemStatus status)
       throws IOException {
+    if (finalStateSent && !status.getCompleted()) {
+      LOG.debug(
+          "Final state already sent for work item {}, skipping non-final status update.",
+          uniqueWorkId());
+      return null;
+    }
+    status.setReportIndex(
+        checkNotNull(nextReportIndex, "nextReportIndex should be non-null when sending an update"));
+
     WorkItemServiceState result = workUnitClient.reportWorkItemStatus(status);
     if (result != null) {
       if (result.getCompleteWorkStatus() != null
@@ -230,7 +286,9 @@ private static boolean isReadLoopAbortedError(Throwable t) {
       if (nextReportIndex == null && !status.getCompleted()) {
         LOG.error("Missing next work index in {} when reporting {}.", result, status);
       }
-      commitMetrics();
+      if (status.getMetricUpdates() != null || status.getCounterUpdates() != null) {
+        commitMetrics();
+      }
     }
 
     if (status.getCompleted()) {
@@ -292,13 +350,15 @@ synchronized void populateMetricUpdates(WorkItemStatus status) {
   }
 
   private synchronized WorkItemStatus createStatusUpdate(boolean isFinal) {
+    return createStatusUpdate(isFinal, true);
+  }
+
+  private synchronized WorkItemStatus createStatusUpdate(boolean isFinal, boolean includeCounters) {
     WorkItemStatus status = new WorkItemStatus();
     status.setWorkItemId(Long.toString(workItem.getId()));
     status.setCompleted(isFinal);
-    status.setReportIndex(
-        checkNotNull(nextReportIndex, "nextReportIndex should be non-null when sending an update"));
 
-    if (worker != null) {
+    if (worker != null && includeCounters) {
       populateMetricUpdates(status);
       populateCounterUpdates(status);
     }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/WorkProgressUpdater.java

@@ -184,6 +184,27 @@ public void requestCheckpoint() {
     }
   }
 
+  /**
+   * Reports any pending dynamic split if one has not yet been reported.
+   *
+   * <p>This should be invoked when the worker finishes processing elements, before reporting final
+   * success, so that any dynamic split is sent to the backend while progress reporting is still
+   * active.
+   *
+   * @throws Exception if reporting the dynamic split fails
+   */
+  public void reportUnreportedSplit() throws Exception {
+    synchronized (executor) {
+      if (dynamicSplitResultToReport != null) {
+        LOG.debug(
+            "Sending progress update with unreported split: {} for work item: {}",
+            dynamicSplitResultToReport,
+            workString());
+        reportProgressHelper();
+      }
+    }
+  }
+
   /**
    * Stops sending work progress updates to the worker service. It may throw an exception if the
    * final progress report fails to be sent for some reason.
@@ -196,7 +217,7 @@ public void stopReportingProgress() throws Exception {
       // We send a final progress report in case there was an unreported dynamic split.
       if (dynamicSplitResultToReport != null) {
         LOG.debug(
-            "Sending final progress update with unreported split: {} " + "for work item: {}",
+            "Sending final progress update with unreported split: {} for work item: {}",
             dynamicSplitResultToReport,
             workString());
         reportProgressHelper(); // This call can fail with an exception