Applying review comments (from Imran)

Author: attilapiros

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

@@ -491,7 +491,7 @@ private[spark] class ExecutorAllocationManager(
     newExecutorTotal = numExistingExecutors
     if (testing || executorsRemoved.nonEmpty) {
       executorsRemoved.foreach { removedExecutorId =>
-        // If it is a cached block, it uses cachedExecutorIdleTimeoutS for timeout
+        // If it has an exclusive cached block then cachedExecutorIdleTimeoutS is used for timeout
         val idleTimeout = if (blockManagerMaster.hasExclusiveCachedBlocks(removedExecutorId)) {
           cachedExecutorIdleTimeoutS
         } else {
@@ -605,9 +605,9 @@ private[spark] class ExecutorAllocationManager(
   private def onExecutorIdle(executorId: String): Unit = synchronized {
     if (executorIds.contains(executorId)) {
       if (!removeTimes.contains(executorId) && !executorsPendingToRemove.contains(executorId)) {
-        // Note that it is not necessary to query the executors since all the cached
-        // blocks we are concerned with are reported to the driver. Note that this
-        // does not include broadcast blocks.
+        // Note that it is not necessary to query the executors since all the cached blocks we are
+        // concerned with are reported to the driver. This does not include broadcast blocks and
+        // non-exclusive blocks which are also available via the external shuffle service.
         val hasCachedBlocks = blockManagerMaster.hasExclusiveCachedBlocks(executorId)
         val now = clock.getTimeMillis()
         val timeout = {

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

@@ -223,7 +223,8 @@ class BlockManagerMaster(
   }
 
   /**
-   * Find out if the executor has cached blocks which are only available via this executor.
+   * Find out if the executor has cached blocks which are not available via the external shuffle
+   * service.
    * This method does not consider broadcast blocks, since they are not reported the master.
    */
   def hasExclusiveCachedBlocks(executorId: String): Boolean = {

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

@@ -545,8 +545,8 @@ private[spark] class BlockManagerInfo(
   private val _blocks = new JHashMap[BlockId, BlockStatus]
 
   /**
-   * Cached blocks held exclusively by this BlockManager. This does not include broadcast blocks
-   * and local disk persisted blocks when external shuffle service is enabled.
+   * Cached blocks which are not available via the external shuffle service.
+   * This does not include broadcast blocks.
    */
   private val _exclusiveCachedBlocks = new mutable.HashSet[BlockId]
 

core/src/test/scala/org/apache/spark/storage/BlockManagerInfoSuite.scala

@@ -63,21 +63,6 @@ class BlockManagerInfoSuite extends SparkFunSuite {
     }
   }
 
-  testWithShuffleServiceOnOff("RDD block with MEMORY_AND_DISK") { (svcEnabled, bmInfo) =>
-    val rddId: BlockId = RDDBlockId(0, 0)
-    bmInfo.updateBlockInfo(
-      rddId, StorageLevel.MEMORY_AND_DISK, memSize = 200, diskSize = 400)
-    assert(bmInfo.blocks.asScala ===
-      Map(rddId -> BlockStatus(StorageLevel.MEMORY_AND_DISK, 0, 400)))
-    val exclusiveCachedBlocksForOneMemoryOnly = if (svcEnabled) Set() else Set(rddId)
-    assert(bmInfo.exclusiveCachedBlocks === exclusiveCachedBlocksForOneMemoryOnly)
-    assert(bmInfo.remainingMem === 29800)
-    if (svcEnabled) {
-      assert(bmInfo.externalShuffleServiceBlockStatus.get.asScala ===
-        Map(rddId -> BlockStatus(StorageLevel.MEMORY_AND_DISK, 0, 400)))
-    }
-  }
-
   testWithShuffleServiceOnOff("RDD block with DISK_ONLY") { (svcEnabled, bmInfo) =>
     val rddId: BlockId = RDDBlockId(0, 0)
     bmInfo.updateBlockInfo(rddId, StorageLevel.DISK_ONLY, memSize = 0, diskSize = 200)