Address a few minor comments

Author: Andrew Or

core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala

@@ -115,7 +115,6 @@ private[spark] class ExecutorAllocationManager(sc: SparkContext) extends Logging
    * If not, throw an appropriate exception.
    */
   private def validateSettings(): Unit = {
-    // Verify that bounds are valid
     if (minNumExecutors < 0 || maxNumExecutors < 0) {
       throw new SparkException("spark.dynamicAllocation.{min/max}Executors must be set!")
     }
@@ -126,21 +125,21 @@ private[spark] class ExecutorAllocationManager(sc: SparkContext) extends Logging
       throw new SparkException(s"spark.dynamicAllocation.minExecutors ($minNumExecutors) must " +
         s"be less than or equal to spark.dynamicAllocation.maxExecutors ($maxNumExecutors)!")
     }
-    // Verify that timeouts are positive
     if (schedulerBacklogTimeout <= 0) {
-      throw new SparkException(s"spark.dynamicAllocation.schedulerBacklogTimeout must be > 0!")
+      throw new SparkException("spark.dynamicAllocation.schedulerBacklogTimeout must be > 0!")
     }
     if (sustainedSchedulerBacklogTimeout <= 0) {
       throw new SparkException(
-        s"spark.dynamicAllocation.sustainedSchedulerBacklogTimeout must be > 0!")
+        "spark.dynamicAllocation.sustainedSchedulerBacklogTimeout must be > 0!")
     }
     if (executorIdleTimeout <= 0) {
-      throw new SparkException(s"spark.dynamicAllocation.executorIdleTimeout must be > 0!")
+      throw new SparkException("spark.dynamicAllocation.executorIdleTimeout must be > 0!")
     }
-    // Verify that external shuffle service is enabled
+    // Require external shuffle service for dynamic allocation
+    // Otherwise, we may lose shuffle files when killing executors
     if (!conf.getBoolean("spark.shuffle.service.enabled", false)) {
-      throw new SparkException(s"Dynamic allocation of executors requires the external " +
-        s"shuffle service. You may enable this through spark.shuffle.service.enabled.")
+      throw new SparkException("Dynamic allocation of executors requires the external " +
+        "shuffle service. You may enable this through spark.shuffle.service.enabled.")
     }
   }
 

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -100,7 +100,7 @@ private[spark] class BlockManager(
     val sparkPort = conf.getInt(shuffleServicePortKey, 7337)
     if (SparkHadoopUtil.get.isYarnMode) {
       val hadoopConf = SparkHadoopUtil.get.newConfiguration(conf)
-      Option(hadoopConf.get(shuffleServicePortKey)).map(_.toInt).getOrElse(sparkPort)
+      hadoopConf.getInt(shuffleServicePortKey, sparkPort)
     } else {
       sparkPort
     }

network/yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java

@@ -22,12 +22,12 @@
 
 import org.apache.spark.network.TransportContext;
 import org.apache.spark.network.server.RpcHandler;
+import org.apache.spark.network.server.TransportServer;
 import org.apache.spark.network.shuffle.ExternalShuffleBlockHandler;
 import org.apache.spark.network.util.TransportConf;
 import org.apache.spark.network.util.SystemPropertyConfigProvider;
 
 import org.apache.hadoop.conf.Configuration;
-import org.apache.hadoop.mapreduce.security.token.JobTokenSecretManager;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.hadoop.yarn.server.api.AuxiliaryService;
@@ -43,11 +43,13 @@
  */
 public class YarnShuffleService extends AuxiliaryService {
   private final Logger logger = LoggerFactory.getLogger(YarnShuffleService.class);
-  private static final JobTokenSecretManager secretManager = new JobTokenSecretManager();
 
   private static final String SPARK_SHUFFLE_SERVICE_PORT_KEY = "spark.shuffle.service.port";
   private static final int DEFAULT_SPARK_SHUFFLE_SERVICE_PORT = 7337;
 
+  // Actual server that serves the shuffle files
+  private TransportServer shuffleServer = null;
+
   public YarnShuffleService() {
     super("spark_shuffle");
     logger.info("Initializing Yarn shuffle service for Spark");
@@ -64,7 +66,7 @@ protected void serviceInit(Configuration conf) {
       TransportConf transportConf = new TransportConf(new SystemPropertyConfigProvider());
       RpcHandler rpcHandler = new ExternalShuffleBlockHandler();
       TransportContext transportContext = new TransportContext(transportConf, rpcHandler);
-      transportContext.createServer(port);
+      shuffleServer = transportContext.createServer(port);
       logger.info("Started Yarn shuffle service for Spark on port " + port);
     } catch (Exception e) {
       logger.error("Exception in starting Yarn shuffle service for Spark", e);
@@ -100,4 +102,11 @@ public void stopContainer(ContainerTerminationContext context) {
     ContainerId containerId = context.getContainerId();
     logger.debug("Stopping container " + containerId + "!");
   }
+
+  @Override
+  protected void serviceStop() {
+    if (shuffleServer != null) {
+      shuffleServer.close();
+    }
+  }
 }