Require tasks to explicitly register themselves with the BlockManager.

Author: JoshRosen

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

@@ -207,6 +207,8 @@ private[spark] class Executor(
         logDebug("Task " + taskId + "'s epoch is " + task.epoch)
         env.mapOutputTracker.updateEpoch(task.epoch)
 
+        env.blockManager.registerTask(taskId)
+
         // Run the actual task and measure its runtime.
         taskStart = System.currentTimeMillis()
         var threwException = true

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

@@ -17,13 +17,11 @@
 
 package org.apache.spark.storage
 
-import java.lang
 import javax.annotation.concurrent.GuardedBy
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 
-import com.google.common.cache.{CacheBuilder, CacheLoader, LoadingCache}
 import com.google.common.collect.ConcurrentHashMultiset
 
 import org.apache.spark.{Logging, SparkException, TaskContext}
@@ -141,18 +139,29 @@ private[storage] class BlockInfoManager extends Logging {
 
   /**
    * Tracks the set of blocks that each task has locked for reading, along with the number of times
-   * that a block has been locked (since our read locks are re-entrant). This is thread-safe.
+   * that a block has been locked (since our read locks are re-entrant).
    */
-  private[this] val readLocksByTask: LoadingCache[lang.Long, ConcurrentHashMultiset[BlockId]] = {
-    // We need to explicitly box as java.lang.Long to avoid a type mismatch error:
-    val loader = new CacheLoader[java.lang.Long, ConcurrentHashMultiset[BlockId]] {
-      override def load(t: java.lang.Long) = ConcurrentHashMultiset.create[BlockId]()
-    }
-    CacheBuilder.newBuilder().build(loader)
-  }
+  @GuardedBy("this")
+  private[this] val readLocksByTask =
+    new mutable.HashMap[TaskAttemptId, ConcurrentHashMultiset[BlockId]]
+
+  // ----------------------------------------------------------------------------------------------
+
+  // Initialization for special task attempt ids:
+  registerTask(BlockInfo.NON_TASK_WRITER)
 
   // ----------------------------------------------------------------------------------------------
 
+  /**
+   * Called at the start of a task in order to register that task with this [[BlockInfoManager]].
+   * This must be called prior to calling any other BlockInfoManager methods from that task.
+   */
+  def registerTask(taskAttemptId: TaskAttemptId): Unit = {
+    require(!readLocksByTask.contains(taskAttemptId),
+      s"Task attempt $taskAttemptId is already registered")
+    readLocksByTask(taskAttemptId) = ConcurrentHashMultiset.create()
+  }
+
   /**
    * Returns the current task's task attempt id (which uniquely identifies the task), or
    * [[BlockInfo.NON_TASK_WRITER]] if called by a non-task thread.
@@ -284,7 +293,7 @@ private[storage] class BlockInfoManager extends Logging {
     } else {
       assert(info.readerCount > 0, s"Block $blockId is not locked for reading")
       info.readerCount -= 1
-      val countsForTask = readLocksByTask.get(currentTaskAttemptId)
+      val countsForTask = readLocksByTask(currentTaskAttemptId)
       val newPinCountForTask: Int = countsForTask.remove(blockId, 1) - 1
       assert(newPinCountForTask >= 0,
         s"Task $currentTaskAttemptId release lock on block $blockId more times than it acquired it")
@@ -325,20 +334,21 @@ private[storage] class BlockInfoManager extends Logging {
    */
   def releaseAllLocksForTask(taskAttemptId: TaskAttemptId): Seq[BlockId] = {
     val blocksWithReleasedLocks = mutable.ArrayBuffer[BlockId]()
-    synchronized {
-      writeLocksByTask.remove(taskAttemptId).foreach { locks =>
-        for (blockId <- locks) {
-          infos.get(blockId).foreach { info =>
-            assert(info.writerTask == taskAttemptId)
-            info.writerTask = BlockInfo.NO_WRITER
-          }
-          blocksWithReleasedLocks += blockId
-        }
+
+    val readLocks = synchronized {
+      readLocksByTask.remove(taskAttemptId).get
+    }
+    val writeLocks = synchronized {
+      writeLocksByTask.remove(taskAttemptId).getOrElse(Seq.empty)
+    }
+
+    for (blockId <- writeLocks) {
+      infos.get(blockId).foreach { info =>
+        assert(info.writerTask == taskAttemptId)
+        info.writerTask = BlockInfo.NO_WRITER
       }
-      notifyAll()
+      blocksWithReleasedLocks += blockId
     }
-    val readLocks = readLocksByTask.get(taskAttemptId)
-    readLocksByTask.invalidate(taskAttemptId)
     readLocks.entrySet().iterator().asScala.foreach { entry =>
       val blockId = entry.getElement
       val lockCount = entry.getCount
@@ -350,6 +360,7 @@ private[storage] class BlockInfoManager extends Logging {
         }
       }
     }
+
     synchronized {
       notifyAll()
     }
@@ -369,8 +380,8 @@ private[storage] class BlockInfoManager extends Logging {
    */
   private[storage] def getNumberOfMapEntries: Long = synchronized {
     size +
-      readLocksByTask.size() +
-      readLocksByTask.asMap().asScala.map(_._2.size()).sum +
+      readLocksByTask.size +
+      readLocksByTask.map(_._2.size()).sum +
       writeLocksByTask.size +
       writeLocksByTask.map(_._2.size).sum
   }
@@ -419,7 +430,7 @@ private[storage] class BlockInfoManager extends Logging {
       blockInfo.removed = true
     }
     infos.clear()
-    readLocksByTask.invalidateAll()
+    readLocksByTask.clear()
     writeLocksByTask.clear()
     notifyAll()
   }

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

@@ -631,6 +631,13 @@ private[spark] class BlockManager(
     blockInfoManager.unlock(blockId)
   }
 
+  /**
+   * Registers a task with the BlockManager in order to initialize per-task bookkeeping structures.
+   */
+  def registerTask(taskAttemptId: Long): Unit = {
+    blockInfoManager.registerTask(taskAttemptId)
+  }
+
   /**
    * Release all locks for the given task.
    *

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

@@ -34,6 +34,9 @@ class BlockInfoManagerSuite extends SparkFunSuite with BeforeAndAfterEach {
   override protected def beforeEach(): Unit = {
     super.beforeEach()
     blockInfoManager = new BlockInfoManager()
+    for (t <- 0 to 4) {
+      blockInfoManager.registerTask(t)
+    }
   }
 
   override protected def afterEach(): Unit = {
@@ -62,7 +65,6 @@ class BlockInfoManagerSuite extends SparkFunSuite with BeforeAndAfterEach {
   }
 
   test("initial memory usage") {
-    assert(blockInfoManager.getNumberOfMapEntries === 0)
     assert(blockInfoManager.size === 0)
   }
 
@@ -72,7 +74,8 @@ class BlockInfoManagerSuite extends SparkFunSuite with BeforeAndAfterEach {
     assert(blockInfoManager.lockForWriting("non-existent-block").isEmpty)
   }
 
-  test("basic putAndLockForWritingIfAbsent") {
+  test("basic lockNewBlockForWriting") {
+    val initialNumMapEntries = blockInfoManager.getNumberOfMapEntries
     val blockInfo = newBlockInfo()
     withTaskId(1) {
       assert(blockInfoManager.lockNewBlockForWriting("block", blockInfo))
@@ -86,7 +89,7 @@ class BlockInfoManagerSuite extends SparkFunSuite with BeforeAndAfterEach {
       assert(blockInfo.writerTask === BlockInfo.NO_WRITER)
     }
     assert(blockInfoManager.size === 1)
-    assert(blockInfoManager.getNumberOfMapEntries === 1)
+    assert(blockInfoManager.getNumberOfMapEntries === initialNumMapEntries + 1)
   }
 
   test("read locks are reentrant") {
@@ -273,11 +276,12 @@ class BlockInfoManagerSuite extends SparkFunSuite with BeforeAndAfterEach {
   }
 
   test("releaseAllLocksForTask releases write locks") {
+    val initialNumMapEntries = blockInfoManager.getNumberOfMapEntries
     withTaskId(0) {
       assert(blockInfoManager.lockNewBlockForWriting("block", newBlockInfo()))
     }
-    assert(blockInfoManager.getNumberOfMapEntries === 3)
+    assert(blockInfoManager.getNumberOfMapEntries === initialNumMapEntries + 3)
     blockInfoManager.releaseAllLocksForTask(0)
-    assert(blockInfoManager.getNumberOfMapEntries === 1)
+    assert(blockInfoManager.getNumberOfMapEntries === initialNumMapEntries)
   }
 }

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

@@ -865,6 +865,7 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with BeforeAndAfterE
 
   test("updated block statuses") {
     store = makeBlockManager(12000)
+    store.registerTask(0)
     val list = List.fill(2)(new Array[Byte](2000))
     val bigList = List.fill(8)(new Array[Byte](2000))