(spark) branch branch-4.1 updated: [SPARK-54312][CORE] Avoid repeatedly scheduling tasks for SendHeartbeat/WorkDirClean in standalone worker

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.1 by this push:
     new 7fcfdf30a18d [SPARK-54312][CORE] Avoid repeatedly scheduling tasks for SendHeartbeat/WorkDirClean in standalone worker
7fcfdf30a18d is described below

commit 7fcfdf30a18d1628d09c24e585d36fb0648253fa
Author: Tengfei Huang <te...@gmail.com>
AuthorDate: Wed Nov 19 21:47:02 2025 -0800

    [SPARK-54312][CORE] Avoid repeatedly scheduling tasks for SendHeartbeat/WorkDirClean in standalone worker
    
    ### What changes were proposed in this pull request?
    Currently, [worker](https://github.com/apache/spark/blob/87b3b94232436528f88c9a7aa7ee70758b85a33a/core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala#L495) will schedule tasks forwarding `SendHeartbeat` and `WorkDirCleanup` while `handleRegisterResponse`.
    
    While worker registration could happen multiple times in case of heartbeat timeout/disconnected from master, in these cases the tasks would be scheduled multiple times.
    
    To fix the issue:
    - Adding `heartbeatTask` and `workDirCleanupTask` in worker to tell whether these tasks have been scheduled
    - `heartbeatTask` and `workDirCleanupTask` will be initialized after the 1st registration, and then skipped scheduling these tasks in later registration.
    - Cancel the task and reset `heartbeatTask` and `workDirCleanupTask` when worker stops.
    
    ### Why are the changes needed?
    Fix the issue repeatedly scheduling SendHeartbeat/WorkDirClean tasks after worker registration.
    
    ### Does this PR introduce _any_ user-facing change?
    No
    
    ### How was this patch tested?
    UT added
    
    ### Was this patch authored or co-authored using generative AI tooling?
    No
    
    Closes #53054 from ivoson/duplicate-worker-heartbeat.
    
    Authored-by: Tengfei Huang <te...@gmail.com>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
    (cherry picked from commit d51b4331e419d56822990846ba1df0efdbf0b6b2)
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 .../org/apache/spark/deploy/worker/Worker.scala    | 27 ++++++++++----
 .../apache/spark/deploy/worker/WorkerSuite.scala   | 42 ++++++++++++++++++++--
 2 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala b/core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala
index 98da33a429ec..93fb64f485f6 100755
--- a/core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala
@@ -211,6 +211,8 @@ private[deploy] class Worker(
 
   private var registerMasterFutures: Array[JFuture[_]] = null
   private var registrationRetryTimer: Option[JScheduledFuture[_]] = None
+  private var heartbeatTask: Option[JScheduledFuture[_]] = None
+  private var workDirCleanupTask: Option[JScheduledFuture[_]] = None
 
   // A thread pool for registering with masters. Because registering with a master is a blocking
   // action, this thread pool must be able to create "masterRpcAddresses.size" threads at the same
@@ -492,16 +494,25 @@ private[deploy] class Worker(
         logInfo(log"Successfully registered with master ${MDC(MASTER_URL, preferredMasterAddress)}")
         registered = true
         changeMaster(masterRef, masterWebUiUrl, masterAddress)
-        forwardMessageScheduler.scheduleAtFixedRate(
-          () => Utils.tryLogNonFatalError { self.send(SendHeartbeat) },
-          0, HEARTBEAT_MILLIS, TimeUnit.MILLISECONDS)
-        if (CLEANUP_ENABLED) {
+
+        // Only schedule heartbeat task if not already scheduled. The existing task will
+        // continue running through reconnections, and the SendHeartbeat handler already
+        // checks the 'connected' flag before sending heartbeats to master.
+        if (heartbeatTask.isEmpty) {
+          heartbeatTask = Some(forwardMessageScheduler.scheduleAtFixedRate(
+            () => Utils.tryLogNonFatalError {
+              self.send(SendHeartbeat)
+            },
+            0, HEARTBEAT_MILLIS, TimeUnit.MILLISECONDS))
+        }
+        // Only schedule work directory cleanup task if not already scheduled
+        if (CLEANUP_ENABLED && workDirCleanupTask.isEmpty) {
           logInfo(
             log"Worker cleanup enabled; old application directories will be deleted in: " +
             log"${MDC(PATH, workDir)}")
-          forwardMessageScheduler.scheduleAtFixedRate(
+          workDirCleanupTask = Some(forwardMessageScheduler.scheduleAtFixedRate(
             () => Utils.tryLogNonFatalError { self.send(WorkDirCleanup) },
-            CLEANUP_INTERVAL_MILLIS, CLEANUP_INTERVAL_MILLIS, TimeUnit.MILLISECONDS)
+            CLEANUP_INTERVAL_MILLIS, CLEANUP_INTERVAL_MILLIS, TimeUnit.MILLISECONDS))
         }
 
         val execs = executors.values.map { e =>
@@ -852,6 +863,10 @@ private[deploy] class Worker(
     cleanupThreadExecutor.shutdownNow()
     metricsSystem.report()
     cancelLastRegistrationRetry()
+    heartbeatTask.foreach(_.cancel(true))
+    heartbeatTask = None
+    workDirCleanupTask.foreach(_.cancel(true))
+    workDirCleanupTask = None
     forwardMessageScheduler.shutdownNow()
     registerMasterThreadPool.shutdownNow()
     executors.values.foreach(_.kill())
diff --git a/core/src/test/scala/org/apache/spark/deploy/worker/WorkerSuite.scala b/core/src/test/scala/org/apache/spark/deploy/worker/WorkerSuite.scala
index ff5d314d1688..f9a0efce8870 100644
--- a/core/src/test/scala/org/apache/spark/deploy/worker/WorkerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/worker/WorkerSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.deploy.worker
 
 import java.io.{File, IOException}
+import java.util.concurrent.{ScheduledFuture => JScheduledFuture}
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.function.Supplier
 
@@ -37,7 +38,7 @@ import org.scalatest.matchers.should.Matchers._
 import org.apache.spark.{SecurityManager, SparkConf, SparkFunSuite}
 import org.apache.spark.TestUtils.{createTempJsonFile, createTempScriptWithExpectedOutput}
 import org.apache.spark.deploy.{Command, ExecutorState, ExternalShuffleService}
-import org.apache.spark.deploy.DeployMessages.{DriverStateChanged, ExecutorStateChanged, WorkDirCleanup}
+import org.apache.spark.deploy.DeployMessages.{DriverStateChanged, ExecutorStateChanged, RegisteredWorker, WorkDirCleanup}
 import org.apache.spark.deploy.master.DriverState
 import org.apache.spark.internal.config
 import org.apache.spark.internal.config.SHUFFLE_SERVICE_DB_BACKEND
@@ -46,7 +47,7 @@ import org.apache.spark.network.shuffledb.DBBackend
 import org.apache.spark.resource.{ResourceAllocation, ResourceInformation}
 import org.apache.spark.resource.ResourceUtils._
 import org.apache.spark.resource.TestResourceIDs.{WORKER_FPGA_ID, WORKER_GPU_ID}
-import org.apache.spark.rpc.{RpcAddress, RpcEnv}
+import org.apache.spark.rpc.{RpcAddress, RpcEndpointRef, RpcEnv}
 import org.apache.spark.util.Utils
 
 class WorkerSuite extends SparkFunSuite with Matchers with BeforeAndAfter with PrivateMethodTester {
@@ -405,4 +406,41 @@ class WorkerSuite extends SparkFunSuite with Matchers with BeforeAndAfter with P
     }.getMessage
     assert(m.contains("Whitespace is not allowed"))
   }
+
+  test("SPARK-54312: heartbeat task and workdir cleanup task should only be scheduled once " +
+    "across multiple registrations") {
+    val worker = spy(makeWorker())
+    val masterWebUiUrl = "https://1.2.3.4:8080"
+    val masterAddress = RpcAddress("1.2.3.4", 1234)
+    val masterRef = mock(classOf[RpcEndpointRef])
+    when(masterRef.address).thenReturn(masterAddress)
+
+    def getHeartbeatTask(worker: Worker): Option[JScheduledFuture[_]] = {
+      val _heartbeatTask =
+        PrivateMethod[Option[JScheduledFuture[_]]](Symbol("heartbeatTask"))
+      worker.invokePrivate(_heartbeatTask())
+    }
+
+    def getWorkDirCleanupTask(worker: Worker): Option[JScheduledFuture[_]] = {
+      val _workDirCleanupTask =
+        PrivateMethod[Option[JScheduledFuture[_]]](Symbol("workDirCleanupTask"))
+      worker.invokePrivate(_workDirCleanupTask())
+    }
+
+    // Tasks should not be scheduled yet before registration
+    assert(getHeartbeatTask(worker).isEmpty && getWorkDirCleanupTask(worker).isEmpty)
+
+    val msg = RegisteredWorker(masterRef, masterWebUiUrl, masterAddress, duplicate = false)
+    // Simulate first registration - this should schedule both tasks
+    worker.receive(msg)
+    val heartbeatTask = getHeartbeatTask(worker)
+    val workDirCleanupTask = getWorkDirCleanupTask(worker)
+    assert(heartbeatTask.isDefined && workDirCleanupTask.isDefined)
+
+    // Simulate disconnection and re-registration
+    worker.receive(msg)
+    // After re-registration, the task references should be the same (not rescheduled)
+    assert(getHeartbeatTask(worker) == heartbeatTask)
+    assert(getWorkDirCleanupTask(worker) == workDirCleanupTask)
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org