[SPARK-26632][Core] Separate Thread Configurations of Driver and Executor

core/src/main/scala/org/apache/spark/network/netty/SparkTransportConf.scala

@@ -19,7 +19,7 @@ package org.apache.spark.network.netty
 
 import scala.collection.JavaConverters._
 
-import org.apache.spark.SparkConf
+import org.apache.spark.{SparkConf, SparkContext}
 import org.apache.spark.network.util.{ConfigProvider, NettyUtils, TransportConf}
 
 /**
@@ -43,9 +43,23 @@ object SparkTransportConf {
     // Specify thread configuration based on our JVM's allocation of cores (rather than necessarily
     // assuming we have all the machine's cores).
     // NB: Only set if serverThreads/clientThreads not already set.
+    val serverSpecificThreads = getSpecificNumOfThreads(conf, module, true)
+    val clientSpecificThreads = getSpecificNumOfThreads(conf, module, false)
     val numThreads = NettyUtils.defaultNumThreads(numUsableCores)
-    conf.setIfMissing(s"spark.$module.io.serverThreads", numThreads.toString)
-    conf.setIfMissing(s"spark.$module.io.clientThreads", numThreads.toString)
+
+    // override threads configurations with side (driver and executor) specific values
+    val serverKey = s"spark.$module.io.serverThreads"
+    if (serverSpecificThreads > 0) {
+      conf.set(serverKey, serverSpecificThreads.toString)
+    } else {
+      conf.setIfMissing(serverKey, numThreads.toString)
+    }
+    val clientKey = s"spark.$module.io.clientThreads"
+    if (clientSpecificThreads > 0) {
+      conf.set(clientKey, clientSpecificThreads.toString)
+    } else {
+      conf.setIfMissing(clientKey, numThreads.toString)
+    }
 
     new TransportConf(module, new ConfigProvider {
       override def get(name: String): String = conf.get(name)
@@ -55,4 +69,27 @@ object SparkTransportConf {
       }
     })
   }
+
+  /**
+   * Separate threads configuration of driver and executor
+   * @param conf the [[SparkConf]]
+   * @param module the module name
+   * @param server if true, it's for the serverThreads. Otherwise, it's for the clientThreads.
+   * @return the configured number of threads, or -1 if not configured.
+   */
+  def getSpecificNumOfThreads(
+                               conf: SparkConf,
+                               module: String,
+                               server: Boolean): Int = {
+    val executorId = conf.get("spark.executor.id", "")
+    val isDriver = executorId == SparkContext.DRIVER_IDENTIFIER ||
+      executorId == SparkContext.LEGACY_DRIVER_IDENTIFIER
+    val side = if (isDriver) "driver" else "executor"
+
+    if (server) {
+      conf.getInt(s"spark.$side.$module.io.serverThreads", -1)
+    } else {
+      conf.getInt(s"spark.$side.$module.io.clientThreads", -1)
+    }
+  }
 }

core/src/main/scala/org/apache/spark/rpc/netty/Dispatcher.scala

@@ -24,7 +24,7 @@ import scala.collection.JavaConverters._
 import scala.concurrent.Promise
 import scala.util.control.NonFatal
 
-import org.apache.spark.SparkException
+import org.apache.spark.{SparkConf, SparkContext, SparkException}
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config.Network.RPC_NETTY_DISPATCHER_NUM_THREADS
 import org.apache.spark.network.client.RpcResponseCallback
@@ -40,9 +40,9 @@ import org.apache.spark.util.ThreadUtils
 private[netty] class Dispatcher(nettyEnv: NettyRpcEnv, numUsableCores: Int) extends Logging {
 
   private class EndpointData(
-      val name: String,
-      val endpoint: RpcEndpoint,
-      val ref: NettyRpcEndpointRef) {
+                              val name: String,
+                              val endpoint: RpcEndpoint,
+                              val ref: NettyRpcEndpointRef) {
     val inbox = new Inbox(ref, endpoint)
   }
 
@@ -113,10 +113,10 @@ private[netty] class Dispatcher(nettyEnv: NettyRpcEnv, numUsableCores: Int) exte
     val iter = endpoints.keySet().iterator()
     while (iter.hasNext) {
       val name = iter.next
-        postMessage(name, message, (e) => { e match {
-          case e: RpcEnvStoppedException => logDebug (s"Message $message dropped. ${e.getMessage}")
-          case e: Throwable => logWarning(s"Message $message dropped. ${e.getMessage}")
-        }}
+      postMessage(name, message, (e) => { e match {
+        case e: RpcEnvStoppedException => logDebug (s"Message $message dropped. ${e.getMessage}")
+        case e: Throwable => logWarning(s"Message $message dropped. ${e.getMessage}")
+      }}
       )}
   }
 
@@ -150,9 +150,9 @@ private[netty] class Dispatcher(nettyEnv: NettyRpcEnv, numUsableCores: Int) exte
    * @param callbackIfStopped callback function if the endpoint is stopped.
    */
   private def postMessage(
-      endpointName: String,
-      message: InboxMessage,
-      callbackIfStopped: (Exception) => Unit): Unit = {
+                           endpointName: String,
+                           message: InboxMessage,
+                           callbackIfStopped: (Exception) => Unit): Unit = {
     val error = synchronized {
       val data = endpoints.get(endpointName)
       if (stopped) {
@@ -194,12 +194,26 @@ private[netty] class Dispatcher(nettyEnv: NettyRpcEnv, numUsableCores: Int) exte
     endpoints.containsKey(name)
   }
 
+  def getNumOfThreads(conf: SparkConf): Int = {
+    val executorId = conf.get("spark.executor.id", "")
+    val isDriver = executorId == SparkContext.DRIVER_IDENTIFIER ||
+      executorId == SparkContext.LEGACY_DRIVER_IDENTIFIER
+    val side = if (isDriver) "driver" else "executor"
+
+    val num = conf.getInt(s"spark.$side.rpc.netty.dispatcher.numThreads", -1)
+    if (num > 0) {
+      num
+    } else {
+      val availableCores =
+        if (numUsableCores > 0) numUsableCores else Runtime.getRuntime.availableProcessors()
+      nettyEnv.conf.get(RPC_NETTY_DISPATCHER_NUM_THREADS)
+        .getOrElse(math.max(2, availableCores))
+    }
+  }
+
   /** Thread pool used for dispatching messages. */
   private val threadpool: ThreadPoolExecutor = {
-    val availableCores =
-      if (numUsableCores > 0) numUsableCores else Runtime.getRuntime.availableProcessors()
-    val numThreads = nettyEnv.conf.get(RPC_NETTY_DISPATCHER_NUM_THREADS)
-      .getOrElse(math.max(2, availableCores))
+    val numThreads = getNumOfThreads(nettyEnv.conf)
     val pool = ThreadUtils.newDaemonFixedThreadPool(numThreads, "dispatcher-event-loop")
     for (i <- 0 until numThreads) {
       pool.execute(new MessageLoop)