[SPARK-52185][CORE] Automate the thread dump collection for Spark applications

roczei · roczei · commit 76cdbace3ab4 · 2025-05-17T07:06:59.000+02:00
### What changes were proposed in this pull request?

When a Java program runs for a long time without giving you any feedback/output, how do you determine what the program might be doing and whether it’s stuck? Thread dumps can help in such case. It shows the status of the threads (if it's running/waiting/blocked) and which part of the code is being executed by each thread, being important to detect deadlocks and which part of the program is running.

**The purpose of this pull request is to collect thread dumps at regular intervals. Why? Getting a single thread dump only shows a snapshot of threads, getting several allows us to see if threads are progressing by comparing states.**

Collecting thread dump samples from slow Spark executors or drivers can be challenging, especially in YARN or Kubernetes environments.

Actual solutions which are available for debugging:

1) We need to find out where the Java Virtual Machine (JVM) is running then run the jstack command manually.
2) Download the thread dumps from the Spark UI. For example: http://localhost:4040/executors/threadDump/?executorId=driver
3) Download the thread dumps via Spark API. For example:

curl "http://localhost:4040/api/v1/applications/local-1747400853731/executors/driver/threads"

### Why are the changes needed?

The purpose of this feature request is to automate the thread dump collection at regular intervals. New Spark parameters have been introduced:

- spark.driver.threadDumpCollector.enabled
- spark.executor.threadDumpCollector.enabled
- spark.threadDumpCollector.interval
- spark.threadDumpCollector.dir
- spark.threadDumpCollector.output.type
- spark.threadDumpCollector.include.regex

Example commands

1)

spark-shell  --master local-cluster[2,1,1050] --conf spark.driver.threadDumpCollector.enabled=true --conf spark.executor.threadDumpCollector.enabled=true --conf spark.threadDumpCollector.interval=15s --conf spark.threadDumpCollector.output.type=FILE --conf spark.threadDumpCollector.dir=hdfs:///user/example/jstack_test

The thread dumps will be saved into hdfs:///user/example/jstack_test, example file names: app-20250516161130-0000-driver-2025-05-16_16_12_50.txt, app-20250516161130-0000-0-2025-05-16_16_12_51.txt

2)

spark-shell  --master local-cluster[2,1,1050] --conf spark.driver.threadDumpCollector.enabled=true --conf spark.executor.threadDumpCollector.enabled=true --conf spark.threadDumpCollector.interval=15s --conf spark.threadDumpCollector.output.type=LOG

The thread dumps will be added to the log messages

3)

spark-shell  --master local-cluster[2,1,1050] --conf spark.driver.threadDumpCollector.enabled=true --conf spark.executor.threadDumpCollector.enabled=true --conf spark.threadDumpCollector.interval=15s --conf spark.threadDumpCollector.output.type=LOG --conf spark.threadDumpCollector.include.regex=something

Only those thread dumps will be captured which match the given regular expression (spark.threadDumpCollector.include.regex)

### Does this PR introduce _any_ user-facing change?

Yes, see above

### How was this patch tested?

New unit tests have been created and it has been manually tested as well.

### Was this patch authored or co-authored using generative AI tooling?

No
diff --git a/common/utils/src/main/scala/org/apache/spark/internal/LogKey.scala b/common/utils/src/main/scala/org/apache/spark/internal/LogKey.scala
@@ -835,6 +835,7 @@ private[spark] object LogKeys {
   case object TEMP_PATH extends LogKey
   case object TEST_SIZE extends LogKey
   case object THREAD extends LogKey
+  case object THREAD_DUMPS extends LogKey
   case object THREAD_ID extends LogKey
   case object THREAD_NAME extends LogKey
   case object THREAD_POOL_KEEPALIVE_TIME extends LogKey
diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -240,6 +240,7 @@ class SparkContext(config: SparkConf) extends Logging {
   private var _shutdownHookRef: AnyRef = _
   private var _statusStore: AppStatusStore = _
   private var _heartbeater: Heartbeater = _
+  private var _driverThreadDumpCollector: ThreadDumpCollector = _
   private var _resources: immutable.Map[String, ResourceInformation] = _
   private var _shuffleDriverComponents: ShuffleDriverComponents = _
   private var _plugins: Option[PluginContainer] = None
@@ -629,6 +630,15 @@ class SparkContext(config: SparkConf) extends Logging {
       _env.blockManager.blockStoreClient.setAppAttemptId(attemptId)
     }
 
+    // Create and start the thread dump collector for the Spark driver
+    if (_conf.get(DRIVER_THREAD_DUMP_COLLECTOR_ENABLED)) {
+      _driverThreadDumpCollector = new ThreadDumpCollector(
+        () => ThreadDumpCollector.saveThreadDumps(env),
+        "driver-threadDumpCollector",
+        conf.get(THREAD_DUMP_COLLECTOR_INTERVAL))
+      _driverThreadDumpCollector.start()
+    }
+
     // initialize after application id and attempt id has been initialized
     _shuffleDriverComponents = ShuffleDataIOUtils.loadShuffleDataIO(_conf).driver()
     _shuffleDriverComponents.initializeApplication().asScala.foreach { case (k, v) =>
@@ -2383,6 +2393,12 @@ class SparkContext(config: SparkConf) extends Logging {
       }
       _heartbeater = null
     }
+    if (_conf.get(DRIVER_THREAD_DUMP_COLLECTOR_ENABLED) && _driverThreadDumpCollector != null) {
+      Utils.tryLogNonFatalError {
+        _driverThreadDumpCollector.stop()
+      }
+      _driverThreadDumpCollector = null
+    }
     if (env != null && _heartbeatReceiver != null) {
       Utils.tryLogNonFatalError {
         env.rpcEnv.stop(_heartbeatReceiver)
diff --git a/core/src/main/scala/org/apache/spark/ThreadDumpCollector.scala b/core/src/main/scala/org/apache/spark/ThreadDumpCollector.scala
@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark
+
+import java.nio.charset.StandardCharsets
+import java.time.LocalDateTime
+import java.time.format.DateTimeFormatter
+import java.util.concurrent.TimeUnit
+
+import org.apache.hadoop.fs.{FileSystem, Path}
+import org.apache.hadoop.fs.permission.FsPermission
+
+import org.apache.spark.deploy.SparkHadoopUtil
+import org.apache.spark.internal.{Logging, LogKeys, MDC}
+import org.apache.spark.internal.config.{THREAD_DUMP_COLLECTOR_DIR,
+  THREAD_DUMP_COLLECTOR_OUTPUT_TYPE, THREAD_DUMP_COLLECTOR_PATTERN}
+import org.apache.spark.util.{ThreadUtils, Utils}
+
+/**
+ * Creates a thread dump collector thread which will call the specified collectThreadDumps
+ * function at intervals of intervalMs.
+ *
+ * @param collectThreadDumps the thread dump collector function to call.
+ * @param name               the thread name for the thread dump collector.
+ * @param intervalMs         the interval between stack trace collections.
+ */
+private[spark] class ThreadDumpCollector(
+                                          collectThreadDumps: () => Unit,
+                                          name: String,
+                                          intervalMs: Long) extends Logging {
+  // Executor for the thread dump collector task
+  private val threadDumpCollector = ThreadUtils.newDaemonSingleThreadScheduledExecutor(name)
+
+  /** Schedules a task to collect the thread dumps */
+  def start(): Unit = {
+    val threadDumpCollectorTask = new Runnable() {
+      override def run(): Unit = Utils.logUncaughtExceptions(collectThreadDumps())
+    }
+    threadDumpCollector.scheduleAtFixedRate(threadDumpCollectorTask, intervalMs, intervalMs,
+      TimeUnit.MILLISECONDS)
+  }
+
+  def stop(): Unit = {
+    threadDumpCollector.shutdown()
+    threadDumpCollector.awaitTermination(10, TimeUnit.SECONDS)
+  }
+
+}
+
+private[spark] object ThreadDumpCollector extends Logging {
+  def saveThreadDumps(env: SparkEnv): Unit = {
+    env.conf.get(THREAD_DUMP_COLLECTOR_OUTPUT_TYPE) match {
+      case "LOG" => writeThreadDumpsToLog(env)
+      case "FILE" => writeThreadDumpsToFile(env)
+    }
+  }
+
+  private def validateRegex(env: SparkEnv, collectedThreadDump: String): Boolean = {
+    val regexPattern = env.conf.get(THREAD_DUMP_COLLECTOR_PATTERN).r
+    regexPattern.findFirstIn(collectedThreadDump).isDefined
+  }
+
+  private def writeThreadDumpsToLog(env: SparkEnv): Unit = {
+    val collectedThreadDump = Utils.getThreadDump().map(_.toString).mkString
+    if (validateRegex(env, collectedThreadDump)) {
+      logWarning(log"Thread dumps from ${MDC(LogKeys.EXECUTOR_ID, env.executorId)}:\n" +
+        log"${MDC(LogKeys.THREAD_DUMPS, collectedThreadDump)}")
+    }
+  }
+
+  private def writeThreadDumpsToFile(env: SparkEnv): Unit = {
+    val formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd_HH_mm_ss")
+    val timestamp = LocalDateTime.now.format(formatter)
+    val threadDumpFileName = env.conf.getAppId + "-" + env.executorId + "-" + timestamp + ".txt"
+    val collectedThreadDump = Utils.getThreadDump().map(_.toString).mkString
+    if (validateRegex(env, collectedThreadDump)) {
+      val hadoopConf = SparkHadoopUtil.get.newConfiguration(env.conf)
+      val rootDir = env.conf.get(THREAD_DUMP_COLLECTOR_DIR)
+      val fileSystem: FileSystem = new Path(rootDir).getFileSystem(hadoopConf)
+      val threadDumpFilePermissions = new FsPermission(Integer.parseInt("700", 8).toShort)
+      val dfsLogFile: Path = fileSystem.makeQualified(new Path(rootDir, threadDumpFileName))
+      try {
+        val outputStream = SparkHadoopUtil.createFile(fileSystem, dfsLogFile, allowEC = true)
+        fileSystem.setPermission(dfsLogFile, threadDumpFilePermissions)
+        outputStream.write(collectedThreadDump.getBytes(StandardCharsets
+          .UTF_8))
+        outputStream.close()
+      } catch {
+        case e: Exception =>
+          logError(
+            log"Could not save thread dumps into file from executor ${
+              MDC(LogKeys.EXECUTOR_ID,
+                env.executorId)
+            }", e)
+      }
+    }
+  }
+}
diff --git a/core/src/main/scala/org/apache/spark/executor/Executor.scala b/core/src/main/scala/org/apache/spark/executor/Executor.scala
@@ -89,6 +89,8 @@ private[spark] class Executor(
     log"${LogMDC(OS_ARCH, System.getProperty("os.arch"))}")
   logInfo(log"Java version ${LogMDC(JAVA_VERSION, System.getProperty("java.version"))}")
 
+  private var executorThreadDumpCollector: ThreadDumpCollector = _
+
   private val executorShutdown = new AtomicBoolean(false)
   val stopHookReference = ShutdownHookManager.addShutdownHook(
     () => stop()
@@ -326,6 +328,15 @@ private[spark] class Executor(
 
   heartbeater.start()
 
+  // Create and start the thread dump collector for the Spark executor
+  if (conf.get(EXECUTOR_THREAD_DUMP_COLLECTOR_ENABLED)) {
+    executorThreadDumpCollector = new ThreadDumpCollector(
+      () => ThreadDumpCollector.saveThreadDumps(env),
+      "executor-ThreadDumpCollector",
+      conf.get(THREAD_DUMP_COLLECTOR_INTERVAL))
+    executorThreadDumpCollector.start()
+  }
+
   private val appStartTime = conf.getLong("spark.app.startTime", 0)
 
   // To allow users to distribute plugins and their required files
@@ -446,6 +457,15 @@ private[spark] class Executor(
         case NonFatal(e) =>
           logWarning("Unable to stop heartbeater", e)
       }
+      try {
+        if (conf.get(EXECUTOR_THREAD_DUMP_COLLECTOR_ENABLED) && executorThreadDumpCollector !=
+          null) {
+          executorThreadDumpCollector.stop()
+        }
+      } catch {
+        case NonFatal(e) =>
+          logWarning("Unable to stop the executor thread dump collector", e)
+      }
       ShuffleBlockPusher.stop()
       if (threadPool != null) {
         threadPool.shutdown()
diff --git a/core/src/main/scala/org/apache/spark/internal/config/package.scala b/core/src/main/scala/org/apache/spark/internal/config/package.scala
@@ -2838,4 +2838,50 @@ package object config {
       .checkValues(Set("connect", "classic"))
       .createWithDefault(
         if (sys.env.get("SPARK_CONNECT_MODE").contains("1")) "connect" else "classic")
+
+  private[spark] val DRIVER_THREAD_DUMP_COLLECTOR_ENABLED = ConfigBuilder("spark.driver" +
+    ".threadDumpCollector.enabled")
+    .doc("Whether to enable automatic thread dump collection for driver")
+    .version("4.1.0")
+    .booleanConf
+    .createWithDefault(false)
+
+  private[spark] val EXECUTOR_THREAD_DUMP_COLLECTOR_ENABLED = ConfigBuilder("spark.executor" +
+    ".threadDumpCollector.enabled")
+    .doc("Whether to enable automatic thread dump collection for each executor")
+    .version("4.1.0")
+    .booleanConf
+    .createWithDefault(false)
+
+  private[spark] val THREAD_DUMP_COLLECTOR_INTERVAL =
+    ConfigBuilder("spark.threadDumpCollector.interval")
+      .doc("The interval of time between two thread dump collections.")
+      .version("4.1.0")
+      .timeConf(TimeUnit.MILLISECONDS)
+      .checkValue(_ > 0, "Value should be positive")
+      .createWithDefaultString("10s")
+
+  private[spark] val THREAD_DUMP_COLLECTOR_DIR = ConfigBuilder("spark.threadDumpCollector.dir")
+    .doc("Set the default directory for saving the thread dump files.")
+    .version("4.1.0")
+    .stringConf
+    .createWithDefault("file:/tmp/spark-thread-dumps")
+
+  private[spark] val THREAD_DUMP_COLLECTOR_OUTPUT_TYPE =
+    ConfigBuilder("spark.threadDumpCollector.output.type")
+      .doc("Specifies the type of saving the thread dumps. Can be either LOG (the default) or " +
+        "FILE")
+      .version("4.1.0")
+      .stringConf
+      .transform(_.toUpperCase(Locale.ROOT))
+      .checkValues(Set("LOG", "FILE"))
+      .createWithDefault("LOG")
+
+  private[spark] val THREAD_DUMP_COLLECTOR_PATTERN =
+    ConfigBuilder("spark.threadDumpCollector.include.regex")
+      .doc("Regular expression for determining which thread dumps will be captured")
+      .version("4.1.0")
+      .stringConf
+      .createWithDefault(".*")
+
 }
diff --git a/core/src/test/scala/org/apache/spark/ThreadDumpCollectorSuite.scala b/core/src/test/scala/org/apache/spark/ThreadDumpCollectorSuite.scala
@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark
+
+import java.io.File
+import java.nio.file.Files
+
+import org.apache.commons.io.FileUtils
+
+import org.apache.spark.internal.config.{DRIVER_LOG_DFS_DIR, DRIVER_LOG_PERSISTTODFS,
+  DRIVER_THREAD_DUMP_COLLECTOR_ENABLED,
+  EXECUTOR_THREAD_DUMP_COLLECTOR_ENABLED, THREAD_DUMP_COLLECTOR_DIR,
+  THREAD_DUMP_COLLECTOR_INTERVAL, THREAD_DUMP_COLLECTOR_OUTPUT_TYPE}
+import org.apache.spark.internal.config.ConfigEntry
+import org.apache.spark.network.util.JavaUtils
+import org.apache.spark.util.Utils
+
+class ThreadDumpCollectorSuite extends SparkFunSuite {
+
+  private var tmpDir : File = _
+
+  override def beforeEach(): Unit = {
+    super.beforeEach()
+    tmpDir = Utils.createTempDir()
+  }
+
+  override def afterEach(): Unit = {
+    super.afterEach()
+    JavaUtils.deleteRecursively(tmpDir)
+  }
+
+  def writeThreadDump(confEntry: ConfigEntry[Boolean], pattern: String, outputType: String)
+  : Unit = {
+    val conf = new SparkConf()
+    conf.set(confEntry, true)
+    conf.set(THREAD_DUMP_COLLECTOR_INTERVAL, 1000L)
+    conf.set(THREAD_DUMP_COLLECTOR_OUTPUT_TYPE, outputType)
+
+    if (outputType == "LOG") {
+      conf.set(DRIVER_LOG_DFS_DIR, tmpDir.getAbsolutePath)
+      conf.set(DRIVER_LOG_PERSISTTODFS, true)
+    }
+    else {
+      conf.set(THREAD_DUMP_COLLECTOR_DIR, tmpDir.getAbsolutePath)
+    }
+
+    val sc = new SparkContext("local", "ThreadDumpWriteToFileTest", conf)
+    Thread.sleep(3000)
+    // Run a simple spark application
+    sc.parallelize(1 to 10).count()
+    sc.stop()
+    val threadDumpDir = FileUtils.getFile(tmpDir.getAbsolutePath)
+    assert(threadDumpDir.exists())
+    val files = threadDumpDir.listFiles().filter(file => !file.isHidden)
+    assert(files.length >= 1)
+    assert(files.forall { file =>
+      Files.lines(file.toPath).anyMatch(line => line.contains(pattern))
+    })
+  }
+
+  test("Spark executor thread dumps are persisted to dfs") {
+    writeThreadDump(EXECUTOR_THREAD_DUMP_COLLECTOR_ENABLED, "executor-heartbeater", "FILE")
+  }
+
+  test("Spark driver thread dumps are persisted to dfs") {
+    writeThreadDump(DRIVER_THREAD_DUMP_COLLECTOR_ENABLED, "driver-heartbeater", "FILE")
+  }
+
+  test("Spark driver thread dumps are persisted to driver log") {
+    writeThreadDump(DRIVER_THREAD_DUMP_COLLECTOR_ENABLED, "driver-heartbeater", "LOG")
+  }
+
+}
