Incorporating feedback from @sameeragarwal

Author: shubhamchopra

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

@@ -1134,11 +1134,14 @@ private[spark] class BlockManager(
    * Called for pro-active replenishment of blocks lost due to executor failures
    *
    * @param blockId blockId being replicate
-   * @param replicas existing block managers that have a replica
-   * @param maxReps maximum replicas needed
+   * @param existingReplicas existing block managers that have a replica
+   * @param maxReplicas maximum replicas needed
    * @return
    */
-  def replicateBlock(blockId: BlockId, replicas: Set[BlockManagerId], maxReps: Int): Boolean = {
+  def replicateBlock(
+    blockId: BlockId,
+    existingReplicas: Set[BlockManagerId],
+    maxReplicas: Int): Unit = {
     logInfo(s"Pro-actively replicating $blockId")
     val infoForReplication = blockInfoManager.lockForReading(blockId).map { info =>
       val data = doGetLocalBytes(blockId, info)
@@ -1147,24 +1150,17 @@ private[spark] class BlockManager(
         info.level.useMemory,
         info.level.useOffHeap,
         info.level.deserialized,
-        maxReps)
+        maxReplicas)
       (data, storageLevel, info.classTag)
     }
     infoForReplication.foreach { case (data, storageLevel, classTag) =>
-      replicate(blockId, data, storageLevel, classTag, replicas)
+      replicate(blockId, data, storageLevel, classTag, existingReplicas)
     }
-    true
   }
 
   /**
    * Replicate block to another node. Note that this is a blocking call that returns after
    * the block has been replicated.
-   *
-   * @param blockId
-   * @param data
-   * @param level
-   * @param classTag
-   * @param existingReplicas
    */
   private def replicate(
     blockId: BlockId,
@@ -1189,10 +1185,7 @@ private[spark] class BlockManager(
     var peersFailedToReplicateTo = mutable.HashSet.empty[BlockManagerId]
     var numFailures = 0
 
-    val initialPeers = {
-      val peers = getPeers(false)
-      if(existingReplicas.isEmpty) peers else peers.filter(!existingReplicas.contains(_))
-    }
+    val initialPeers = getPeers(false).filterNot(existingReplicas.contains(_))
 
     var peersForReplication = blockReplicationPolicy.prioritize(
       blockManagerId,

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

@@ -66,6 +66,8 @@ class BlockManagerMasterEndpoint(
     mapper
   }
 
+  val proactivelyReplicate = conf.get("spark.storage.replication.proactive", "false").toBoolean
+
   logInfo("BlockManagerMasterEndpoint up")
 
   override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
@@ -189,13 +191,14 @@ class BlockManagerMasterEndpoint(
   }
 
   private def removeBlockManager(blockManagerId: BlockManagerId) {
-    val proactivelyReplicate = conf.get("spark.storage.replication.proactive", "false").toBoolean
-
     val info = blockManagerInfo(blockManagerId)
 
     // Remove the block manager from blockManagerIdByExecutor.
     blockManagerIdByExecutor -= blockManagerId.executorId
 
+    // Remove it from blockManagerInfo and remove all the blocks.
+    blockManagerInfo.remove(blockManagerId)
+
     val iterator = info.blocks.keySet.iterator
     while (iterator.hasNext) {
       val blockId = iterator.next
@@ -204,8 +207,10 @@ class BlockManagerMasterEndpoint(
       if (locations.size == 0) {
         blockLocations.remove(blockId)
         logWarning(s"No more replicas available for $blockId !")
-      } else if ((blockId.isRDD || blockId.isInstanceOf[TestBlockId]) && proactivelyReplicate) {
-        // we only need to proactively replicate RDD blocks
+      } else if (proactivelyReplicate && (blockId.isRDD || blockId.isInstanceOf[TestBlockId])) {
+        // only RDD blocks store data that users explicitly cache so we only need to proactively
+        // replicate RDD blocks
+        // broadcast related blocks exist on all executors, so we don't worry about them
         // we also need to replicate this behavior for test blocks for unit tests
         // we send a message to a randomly chosen executor location to replicate block
         // assuming single executor failure, we find out how many replicas existed before failure
@@ -222,8 +227,6 @@ class BlockManagerMasterEndpoint(
         }
       }
     }
-    // Remove it from blockManagerInfo and remove all the blocks.
-    blockManagerInfo.remove(blockManagerId)
 
     listenerBus.post(SparkListenerBlockManagerRemoved(System.currentTimeMillis(), blockManagerId))
     logInfo(s"Removing block manager $blockManagerId")