[SPARK-8981][CORE][test-hadoop3.2][test-java11] Add MDC support in Executor

core/src/main/scala/org/apache/spark/executor/Executor.scala

@@ -34,6 +34,7 @@ import scala.concurrent.duration._
 import scala.util.control.NonFatal
 
 import com.google.common.util.concurrent.ThreadFactoryBuilder
+import org.slf4j.MDC
 
 import org.apache.spark._
 import org.apache.spark.deploy.SparkHadoopUtil
@@ -107,7 +108,7 @@ private[spark] class Executor(
       .setNameFormat("Executor task launch worker-%d")
       .setThreadFactory((r: Runnable) => new UninterruptibleThread(r, "unused"))
       .build()
-    Executors.newCachedThreadPool(threadFactory).asInstanceOf[ThreadPoolExecutor]
+    ThreadUtils.newCachedThreadPool(threadFactory)
   }
   private val executorSource = new ExecutorSource(threadPool, executorId)
   // Pool used for threads that supervise task killing / cancellation
@@ -395,6 +396,8 @@ private[spark] class Executor(
     }
 
     override def run(): Unit = {
+      setMDCForTask(taskDescription)
+
       threadId = Thread.currentThread.getId
       Thread.currentThread.setName(threadName)
       val threadMXBean = ManagementFactory.getThreadMXBean
@@ -693,6 +696,17 @@ private[spark] class Executor(
     }
   }
 
+  private def setMDCForTask(taskDescription: TaskDescription): Unit = {
+    val properties = taskDescription.properties
+
+    MDC.put("taskName", taskDescription.name)
+
+    properties.asScala.filter(_._1.startsWith("mdc.")).foreach { item =>
+      val key = item._1.substring(4)
+      MDC.put(key, item._2)
+    }
+  }
+
   /**
    * Supervises the killing / cancellation of a task by sending the interrupted flag, optionally
    * sending a Thread.interrupt(), and monitoring the task until it finishes.

core/src/main/scala/org/apache/spark/util/ThreadUtils.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.util
 
+import java.util
 import java.util.concurrent._
 import java.util.concurrent.{Future => JFuture}
 import java.util.concurrent.locks.ReentrantLock
@@ -32,6 +33,99 @@ import org.apache.spark.SparkException
 
 private[spark] object ThreadUtils {
 
+  object MDCAwareThreadPoolExecutor {
+    def newCachedThreadPool(threadFactory: ThreadFactory): ThreadPoolExecutor = {
+      // The values needs to be synced with `Executors.newCachedThreadPool`
+      new MDCAwareThreadPoolExecutor(
+        0,
+        Integer.MAX_VALUE,
+        60L,
+        TimeUnit.SECONDS,
+        new SynchronousQueue[Runnable],
+        threadFactory)
+    }
+
+    def newFixedThreadPool(nThreads: Int, threadFactory: ThreadFactory): ThreadPoolExecutor = {
+      // The values needs to be synced with `Executors.newFixedThreadPool`
+      new MDCAwareThreadPoolExecutor(
+        nThreads,
+        nThreads,
+        0L,
+        TimeUnit.MILLISECONDS,
+        new LinkedBlockingQueue[Runnable],
+        threadFactory)
+    }
+
+    /**
+     * This method differ from the `java.util.concurrent.Executors#newSingleThreadExecutor` in
+     * 2 ways:
+     *   1. It use `org.apache.spark.util.ThreadUtils.MDCAwareThreadPoolExecutor`
+     *   as underline `java.util.concurrent.ExecutorService`
+     *   2. It does not use the
+     *   `java.util.concurrent.Executors#FinalizableDelegatedExecutorService` from JDK
+     */
+    def newSingleThreadExecutor(threadFactory: ThreadFactory): ExecutorService = {
+      // The values needs to be synced with `Executors.newSingleThreadExecutor`
+      Executors.unconfigurableExecutorService(
+        new MDCAwareThreadPoolExecutor(
+          1,
+          1,
+          0L,
+          TimeUnit.MILLISECONDS,
+          new LinkedBlockingQueue[Runnable],
+          threadFactory)
+        )
+    }
+
+  }
+
+  class MDCAwareRunnable(proxy: Runnable) extends Runnable {
+    val callerThreadMDC: util.Map[String, String] = getMDCMap
+
+    @inline
+    private def getMDCMap: util.Map[String, String] = {
+      org.slf4j.MDC.getCopyOfContextMap match {
+        case null => new util.HashMap[String, String]()
+        case m => m
+      }
+    }
+
+    override def run(): Unit = {
+      val threadMDC = getMDCMap
+      org.slf4j.MDC.setContextMap(callerThreadMDC)
+      try {
+        proxy.run()
+      } finally {
+        org.slf4j.MDC.setContextMap(threadMDC)
+      }
+    }
+  }
+
+  class MDCAwareScheduledThreadPoolExecutor(
+      corePoolSize: Int,
+      threadFactory: ThreadFactory)
+    extends ScheduledThreadPoolExecutor(corePoolSize, threadFactory) {
+
+    override def execute(runnable: Runnable) {
+      super.execute(new MDCAwareRunnable(runnable))
+    }
+  }
+
+  class MDCAwareThreadPoolExecutor(
+      corePoolSize: Int,
+      maximumPoolSize: Int,
+      keepAliveTime: Long,
+      unit: TimeUnit,
+      workQueue: BlockingQueue[Runnable],
+      threadFactory: ThreadFactory)
+    extends ThreadPoolExecutor(corePoolSize, maximumPoolSize,
+      keepAliveTime, unit, workQueue, threadFactory) {
+
+    override def execute(runnable: Runnable) {
+      super.execute(new MDCAwareRunnable(runnable))
+    }
+  }
+
   private val sameThreadExecutionContext =
     ExecutionContext.fromExecutorService(sameThreadExecutorService())
 
@@ -130,7 +224,15 @@ private[spark] object ThreadUtils {
    */
   def newDaemonCachedThreadPool(prefix: String): ThreadPoolExecutor = {
     val threadFactory = namedThreadFactory(prefix)
-    Executors.newCachedThreadPool(threadFactory).asInstanceOf[ThreadPoolExecutor]
+    MDCAwareThreadPoolExecutor.newCachedThreadPool(threadFactory)
+  }
+
+  /**
+   * Wrapper over newCachedThreadPool. Thread names are formatted as prefix-ID, where ID is a
+   * unique, sequentially assigned integer.
+   */
+  def newCachedThreadPool(threadFactory: ThreadFactory): ThreadPoolExecutor = {
+    MDCAwareThreadPoolExecutor.newCachedThreadPool(threadFactory)
   }
 
   /**
@@ -140,7 +242,7 @@ private[spark] object ThreadUtils {
   def newDaemonCachedThreadPool(
       prefix: String, maxThreadNumber: Int, keepAliveSeconds: Int = 60): ThreadPoolExecutor = {
     val threadFactory = namedThreadFactory(prefix)
-    val threadPool = new ThreadPoolExecutor(
+    val threadPool = new MDCAwareThreadPoolExecutor(
       maxThreadNumber, // corePoolSize: the max number of threads to create before queuing the tasks
       maxThreadNumber, // maximumPoolSize: because we use LinkedBlockingDeque, this one is not used
       keepAliveSeconds,
@@ -165,15 +267,15 @@ private[spark] object ThreadUtils {
    */
   def newDaemonSingleThreadExecutor(threadName: String): ExecutorService = {
     val threadFactory = new ThreadFactoryBuilder().setDaemon(true).setNameFormat(threadName).build()
-    Executors.newSingleThreadExecutor(threadFactory)
+    MDCAwareThreadPoolExecutor.newSingleThreadExecutor(threadFactory)
   }
 
   /**
    * Wrapper over ScheduledThreadPoolExecutor.
    */
   def newDaemonSingleThreadScheduledExecutor(threadName: String): ScheduledExecutorService = {
     val threadFactory = new ThreadFactoryBuilder().setDaemon(true).setNameFormat(threadName).build()
-    val executor = new ScheduledThreadPoolExecutor(1, threadFactory)
+    val executor = new MDCAwareScheduledThreadPoolExecutor(1, threadFactory)
     // By default, a cancelled task is not automatically removed from the work queue until its delay
     // elapses. We have to enable it manually.
     executor.setRemoveOnCancelPolicy(true)
@@ -189,7 +291,7 @@ private[spark] object ThreadUtils {
       .setDaemon(true)
       .setNameFormat(s"$threadNamePrefix-%d")
       .build()
-    val executor = new ScheduledThreadPoolExecutor(numThreads, threadFactory)
+    val executor = new MDCAwareScheduledThreadPoolExecutor(numThreads, threadFactory)
     // By default, a cancelled task is not automatically removed from the work queue until its delay
     // elapses. We have to enable it manually.
     executor.setRemoveOnCancelPolicy(true)

docs/configuration.md

@@ -2955,6 +2955,12 @@ Spark uses [log4j](http://logging.apache.org/log4j/) for logging. You can config
 `log4j.properties` file in the `conf` directory. One way to start is to copy the existing
 `log4j.properties.template` located there.
 
+By default, Spark adds 1 record to the MDC: `taskName`, which shows something
+like `task 1.0 in stage 0.0`. You can add `%X{taskName}` to your patternLayout in
+order to print it in the logs.
+Moreover, you can use `spark.sparkContext.setLocalProperty("mdc." + name, "value")` to add user specific data into MDC.
+The key in MDC will be the string after the `mdc.` prefix.
+
 # Overriding configuration directory
 
 To specify a different configuration directory other than the default "SPARK_HOME/conf",