[EXT][SPARK-13566][CORE] Avoid deadlock between BlockManager and Executor Thread apache#11546

Author: zzcclp

core/src/main/scala/org/apache/spark/executor/Executor.scala

@@ -218,6 +218,7 @@ private[spark] class Executor(
           threwException = false
           res
         } finally {
+          val releasedLocks = env.blockManager.releaseAllLocksForTask(taskId)
           val freedMemory = taskMemoryManager.cleanUpAllAllocatedMemory()
           if (freedMemory > 0) {
             val errMsg = s"Managed memory leak detected; size = $freedMemory bytes, TID = $taskId"
@@ -227,6 +228,17 @@ private[spark] class Executor(
               logError(errMsg)
             }
           }
+          
+          if (releasedLocks.nonEmpty) {
+            val errMsg =
+              s"${releasedLocks.size} block locks were not released by TID = $taskId:\n" +
+              releasedLocks.mkString("[", ", ", "]")
+            if (conf.getBoolean("spark.storage.exceptionOnPinLeak", false) && !threwException) {
+              throw new SparkException(errMsg)
+            } else {
+              logError(errMsg)
+            }
+          }
         }
         val taskFinish = System.currentTimeMillis()
 

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

@@ -19,12 +19,14 @@ package org.apache.spark.storage
 
 import java.io._
 import java.nio.{ByteBuffer, MappedByteBuffer}
+import java.util.concurrent.ConcurrentHashMap
 
 import scala.collection.mutable.{ArrayBuffer, HashMap}
 import scala.concurrent.duration._
 import scala.concurrent.{Await, ExecutionContext, Future}
 import scala.util.Random
 import scala.util.control.NonFatal
+import scala.collection.JavaConverters._
 
 import sun.nio.ch.DirectBuffer
 
@@ -65,7 +67,7 @@ private[spark] class BlockManager(
     val master: BlockManagerMaster,
     defaultSerializer: Serializer,
     val conf: SparkConf,
-    memoryManager: MemoryManager,
+    val memoryManager: MemoryManager,
     mapOutputTracker: MapOutputTracker,
     shuffleManager: ShuffleManager,
     blockTransferService: BlockTransferService,
@@ -163,6 +165,11 @@ private[spark] class BlockManager(
    * Executor.updateDependencies. When the BlockManager is initialized, user level jars hasn't been
    * loaded yet. */
   private lazy val compressionCodec: CompressionCodec = CompressionCodec.createCodec(conf)
+  
+  // Blocks are removing by another thread
+  private val pendingToRemove = new ConcurrentHashMap[BlockId, Long]()
+
+  private val NON_TASK_WRITER = -1024L
 
   /**
    * Initializes the BlockManager with the given appId. This is not performed in the constructor as
@@ -1025,7 +1032,7 @@ private[spark] class BlockManager(
     val info = blockInfo.get(blockId).orNull
 
     // If the block has not already been dropped
-    if (info != null) {
+    if (info != null && !pendingToRemove.containsKey(blockId)) {
       info.synchronized {
         // required ? As of now, this will be invoked only for blocks which are ready
         // But in case this changes in future, adding for consistency sake.
@@ -1051,11 +1058,13 @@ private[spark] class BlockManager(
           }
           blockIsUpdated = true
         }
+        pendingToRemove.put(blockId, currentTaskAttemptId)
 
         // Actually drop from memory store
         val droppedMemorySize =
           if (memoryStore.contains(blockId)) memoryStore.getSize(blockId) else 0L
         val blockIsRemoved = memoryStore.remove(blockId)
+        pendingToRemove.remove(blockId)
         if (blockIsRemoved) {
           blockIsUpdated = true
         } else {
@@ -1080,6 +1089,7 @@ private[spark] class BlockManager(
 
   /**
    * Remove all blocks belonging to the given RDD.
+   * 
    * @return The number of blocks removed.
    */
   def removeRdd(rddId: Int): Int = {
@@ -1108,11 +1118,14 @@ private[spark] class BlockManager(
   def removeBlock(blockId: BlockId, tellMaster: Boolean = true): Unit = {
     logDebug(s"Removing block $blockId")
     val info = blockInfo.get(blockId).orNull
-    if (info != null) {
+    if (info != null && !pendingToRemove.containsKey(blockId)) {
+      pendingToRemove.put(blockId, currentTaskAttemptId)
       info.synchronized {
+        val level = info.level
         // Removals are idempotent in disk store and memory store. At worst, we get a warning.
-        val removedFromMemory = memoryStore.remove(blockId)
-        val removedFromDisk = diskStore.remove(blockId)
+        val removedFromMemory = if (level.useMemory) memoryStore.remove(blockId) else false
+        pendingToRemove.remove(blockId)
+        val removedFromDisk = if (level.useDisk) diskStore.remove(blockId) else false
         val removedFromExternalBlockStore =
           if (externalBlockStoreInitialized) externalBlockStore.remove(blockId) else false
         if (!removedFromMemory && !removedFromDisk && !removedFromExternalBlockStore) {
@@ -1147,9 +1160,11 @@ private[spark] class BlockManager(
       val entry = iterator.next()
       val (id, info, time) = (entry.getKey, entry.getValue.value, entry.getValue.timestamp)
       if (time < cleanupTime && shouldDrop(id)) {
+        pendingToRemove.put(id, currentTaskAttemptId)
         info.synchronized {
           val level = info.level
           if (level.useMemory) { memoryStore.remove(id) }
+          pendingToRemove.remove(id)
           if (level.useDisk) { diskStore.remove(id) }
           if (level.useOffHeap) { externalBlockStore.remove(id) }
           iterator.remove()
@@ -1160,6 +1175,28 @@ private[spark] class BlockManager(
       }
     }
   }
+  
+  private def currentTaskAttemptId: Long = {
+    Option(TaskContext.get()).map(_.taskAttemptId()).getOrElse(NON_TASK_WRITER)
+  }
+
+  /**
+   * Release all lock held by the given task, clearing that task's pin bookkeeping
+   * structures and updating the global pin counts. This method should be called at the
+   * end of a task (either by a task completion handler or in `TaskRunner.run()`).
+   *
+   * @return the ids of blocks whose pins were released
+   */
+  def releaseAllLocksForTask(taskAttemptId: Long): ArrayBuffer[BlockId] = {
+    var selectLocks = ArrayBuffer[BlockId]()
+    pendingToRemove.entrySet().asScala.foreach { entry =>
+      if (entry.getValue == taskAttemptId) {
+        pendingToRemove.remove(entry.getKey)
+        selectLocks += entry.getKey
+      }
+    }
+    selectLocks
+  }
 
   private def shouldCompress(blockId: BlockId): Boolean = {
     blockId match {
@@ -1234,6 +1271,7 @@ private[spark] class BlockManager(
     rpcEnv.stop(slaveEndpoint)
     blockInfo.clear()
     memoryStore.clear()
+    pendingToRemove.clear()
     diskStore.clear()
     if (externalBlockStoreInitialized) {
       externalBlockStore.clear()

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

@@ -423,6 +423,32 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with BeforeAndAfterE
       store.waitForAsyncReregister()
     }
   }
+  
+  test("deadlock between dropFromMemory and removeBlock") {
+    store = makeBlockManager(2000)
+    val a1 = new Array[Byte](400)
+    store.putSingle("a1", a1, StorageLevel.MEMORY_ONLY)
+    val t1 = new Thread {
+      override def run() = {
+        store.memoryManager.synchronized {
+          Thread.sleep(1000)
+          val status = store.dropFromMemory("a1", null: Either[Array[Any], ByteBuffer])
+          assert(status == None, "this thread can not get block a1")
+        }
+      }
+    }
+
+    val t2 = new Thread {
+      override def run() = {
+        store.removeBlock("a1", tellMaster = false)
+      }
+    }
+
+    t1.start()
+    t2.start()
+    t1.join()
+    t2.join()
+  }
 
   test("correct BlockResult returned from get() calls") {
     store = makeBlockManager(12000)