Incorporating suggestions from @JoshRosen and @sameeragarwal

Author: shubhamchopra

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

@@ -1136,25 +1136,25 @@ private[spark] class BlockManager(
    * @param blockId blockId being replicate
    * @param existingReplicas existing block managers that have a replica
    * @param maxReplicas maximum replicas needed
-   * @return
    */
   def replicateBlock(
-    blockId: BlockId,
-    existingReplicas: Set[BlockManagerId],
-    maxReplicas: Int): Unit = {
+      blockId: BlockId,
+      existingReplicas: Set[BlockManagerId],
+      maxReplicas: Int): Unit = {
     logInfo(s"Pro-actively replicating $blockId")
-    val infoForReplication = blockInfoManager.lockForReading(blockId).map { info =>
+    val blockInfo = blockInfoManager.lockForReading(blockId).foreach { info =>
       val data = doGetLocalBytes(blockId, info)
       val storageLevel = StorageLevel(
-        info.level.useDisk,
-        info.level.useMemory,
-        info.level.useOffHeap,
-        info.level.deserialized,
-        maxReplicas)
-      (data, storageLevel, info.classTag)
-    }
-    infoForReplication.foreach { case (data, storageLevel, classTag) =>
-      replicate(blockId, data, storageLevel, classTag, existingReplicas)
+        useDisk = info.level.useDisk,
+        useMemory = info.level.useMemory,
+        useOffHeap = info.level.useOffHeap,
+        deserialized = info.level.deserialized,
+        replication = maxReplicas)
+      try {
+        replicate(blockId, data, storageLevel, info.classTag, existingReplicas)
+      } finally {
+        releaseLock(blockId)
+      }
     }
   }
 
@@ -1163,11 +1163,11 @@ private[spark] class BlockManager(
    * the block has been replicated.
    */
   private def replicate(
-    blockId: BlockId,
-    data: ChunkedByteBuffer,
-    level: StorageLevel,
-    classTag: ClassTag[_],
-    existingReplicas: Set[BlockManagerId] = Set.empty): Unit = {
+      blockId: BlockId,
+      data: ChunkedByteBuffer,
+      level: StorageLevel,
+      classTag: ClassTag[_],
+      existingReplicas: Set[BlockManagerId] = Set.empty): Unit = {
 
     val maxReplicationFailures = conf.getInt("spark.storage.maxReplicationFailures", 1)
     val tLevel = StorageLevel(

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

@@ -204,26 +204,25 @@ class BlockManagerMasterEndpoint(
       val blockId = iterator.next
       val locations = blockLocations.get(blockId)
       locations -= blockManagerId
+      // De-register the block if none of the block managers have it. Otherwise, if pro-active
+      // replication is enabled, and a block is either an RDD or a test block (the latter is used
+      // for unit testing), we send a message to a randomly chosen executor location to replicate
+      // the given block. Note that we ignore other block types (such as broadcast/shuffle blocks
+      // etc.) as replication doesn't make much sense in that context.
       if (locations.size == 0) {
         blockLocations.remove(blockId)
         logWarning(s"No more replicas available for $blockId !")
       } 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
+        // As a heursitic, assume single executor failure to find out the number of replicas that
+        // existed before failure
         val maxReplicas = locations.size + 1
-
         val i = (new Random(blockId.hashCode)).nextInt(locations.size)
         val blockLocations = locations.toSeq
         val candidateBMId = blockLocations(i)
-        val blockManager = blockManagerInfo.get(candidateBMId)
-        if(blockManager.isDefined) {
+        blockManagerInfo.get(candidateBMId).foreach { bm =>
           val remainingLocations = locations.toSeq.filter(bm => bm != candidateBMId)
           val replicateMsg = ReplicateBlock(blockId, remainingLocations, maxReplicas)
-          blockManager.get.slaveEndpoint.ask[Boolean](replicateMsg)
+          bm.slaveEndpoint.ask[Boolean](replicateMsg)
         }
       }
     }

docs/configuration.md

@@ -952,6 +952,15 @@ Apart from these, the following properties are also available, and may be useful
     storage space to unroll the new block in its entirety.
   </td>
 </tr>
+<tr>
+  <td><code>spark.storage.replication.proactive<code></td>
+  <td>false</td>
+  <td>
+    Enables proactive block replication for RDD blocks. Cached RDD block replicas lost due to
+    executor failures are replenished if there are any existing available replicas. This tries
+    to get the replication level of the block to the initial number.
+  </td>
+</tr>
 </table>
 
 #### Execution Behavior