Address comments

Author: Andrew Or

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

@@ -70,7 +70,7 @@ class SparkEnv (
     val httpFileServer: HttpFileServer,
     val sparkFilesDir: String,
     val metricsSystem: MetricsSystem,
-    // TODO: unify these *MemoryManager classes
+    // TODO: unify these *MemoryManager classes (SPARK-10984)
     val memoryManager: MemoryManager,
     val shuffleMemoryManager: ShuffleMemoryManager,
     val executorMemoryManager: ExecutorMemoryManager,

core/src/main/scala/org/apache/spark/memory/MemoryManager.scala

@@ -19,18 +19,35 @@ package org.apache.spark.memory
 
 import scala.collection.mutable
 
-import org.apache.spark.storage.{BlockId, BlockStatus}
+import org.apache.spark.storage.{BlockId, BlockStatus, MemoryStore}
 
 
 /**
  * An abstract memory manager that enforces how memory is shared between execution and storage.
  *
  * In this context, execution memory refers to that used for computation in shuffles, joins,
  * sorts and aggregations, while storage memory refers to that used for caching and propagating
- * internal data across the cluster.
+ * internal data across the cluster. There exists one of these per JVM.
  */
 private[spark] abstract class MemoryManager {
 
+  // The memory store used to evict cached blocks
+  private var _memoryStore: MemoryStore = _
+  protected def memoryStore: MemoryStore = {
+    if (_memoryStore == null) {
+      throw new IllegalArgumentException("memory store not initialized yet")
+    }
+    _memoryStore
+  }
+
+  /**
+   * Set the [[MemoryStore]] used by this manager to evict cached blocks.
+   * This must be set after construction due to initialization ordering constraints.
+   */
+  def setMemoryStore(store: MemoryStore): Unit = {
+    _memoryStore = store
+  }
+
   /**
    * Acquire N bytes of memory for execution.
    * @return number of bytes successfully granted (<= N).

core/src/main/scala/org/apache/spark/memory/StaticMemoryManager.scala

@@ -19,8 +19,8 @@ package org.apache.spark.memory
 
 import scala.collection.mutable
 
-import org.apache.spark.{Logging, SparkConf, SparkEnv}
-import org.apache.spark.storage.{BlockId, BlockStatus, MemoryStore}
+import org.apache.spark.{Logging, SparkConf}
+import org.apache.spark.storage.{BlockId, BlockStatus}
 
 
 /**
@@ -41,27 +41,10 @@ private[spark] class StaticMemoryManager(
     (maxStorageMemory * conf.getDouble("spark.storage.unrollFraction", 0.2)).toLong
   }
 
-  // Amount of execution memory in use. Accesses must be synchronized on `executionLock`.
+  // Amount of execution / storage memory in use
+  // Accesses must be synchronized on `this`
   private var _executionMemoryUsed: Long = 0
-  private val executionLock = new Object
-
-  // Amount of storage memory in use. Accesses must be synchronized on `storageLock`.
   private var _storageMemoryUsed: Long = 0
-  private val storageLock = new Object
-
-  // The memory store used to evict cached blocks
-  private var _memoryStore: MemoryStore = _
-  private def memoryStore: MemoryStore = {
-    if (_memoryStore == null) {
-      _memoryStore = SparkEnv.get.blockManager.memoryStore
-    }
-    _memoryStore
-  }
-
-  // For testing only
-  def setMemoryStore(store: MemoryStore): Unit = {
-    _memoryStore = store
-  }
 
   def this(conf: SparkConf) {
     this(
@@ -74,13 +57,11 @@ private[spark] class StaticMemoryManager(
    * Acquire N bytes of memory for execution.
    * @return number of bytes successfully granted (<= N).
    */
-  override def acquireExecutionMemory(numBytes: Long): Long = {
-    executionLock.synchronized {
-      assert(_executionMemoryUsed <= maxExecutionMemory)
-      val bytesToGrant = math.min(numBytes, maxExecutionMemory - _executionMemoryUsed)
-      _executionMemoryUsed += bytesToGrant
-      bytesToGrant
-    }
+  override def acquireExecutionMemory(numBytes: Long): Long = synchronized {
+    assert(_executionMemoryUsed <= maxExecutionMemory)
+    val bytesToGrant = math.min(numBytes, maxExecutionMemory - _executionMemoryUsed)
+    _executionMemoryUsed += bytesToGrant
+    bytesToGrant
   }
 
   /**
@@ -130,7 +111,7 @@ private[spark] class StaticMemoryManager(
       evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = {
     // Note: Keep this outside synchronized block to avoid potential deadlocks!
     memoryStore.ensureFreeSpace(blockId, numBytesToFree, evictedBlocks)
-    storageLock.synchronized {
+    synchronized {
       assert(_storageMemoryUsed <= maxStorageMemory)
       val enoughMemory = _storageMemoryUsed + numBytesToAcquire <= maxStorageMemory
       if (enoughMemory) {
@@ -143,40 +124,34 @@ private[spark] class StaticMemoryManager(
   /**
    * Release N bytes of execution memory.
    */
-  override def releaseExecutionMemory(numBytes: Long): Unit = {
-    executionLock.synchronized {
-      if (numBytes > _executionMemoryUsed) {
-        logWarning(s"Attempted to release $numBytes bytes of execution " +
-          s"memory when we only have ${_executionMemoryUsed} bytes")
-        _executionMemoryUsed = 0
-      } else {
-        _executionMemoryUsed -= numBytes
-      }
+  override def releaseExecutionMemory(numBytes: Long): Unit = synchronized {
+    if (numBytes > _executionMemoryUsed) {
+      logWarning(s"Attempted to release $numBytes bytes of execution " +
+        s"memory when we only have ${_executionMemoryUsed} bytes")
+      _executionMemoryUsed = 0
+    } else {
+      _executionMemoryUsed -= numBytes
     }
   }
 
   /**
    * Release N bytes of storage memory.
    */
-  override def releaseStorageMemory(numBytes: Long): Unit = {
-    storageLock.synchronized {
-      if (numBytes > _storageMemoryUsed) {
-        logWarning(s"Attempted to release $numBytes bytes of storage " +
-          s"memory when we only have ${_storageMemoryUsed} bytes")
-        _storageMemoryUsed = 0
-      } else {
-        _storageMemoryUsed -= numBytes
-      }
+  override def releaseStorageMemory(numBytes: Long): Unit = synchronized {
+    if (numBytes > _storageMemoryUsed) {
+      logWarning(s"Attempted to release $numBytes bytes of storage " +
+        s"memory when we only have ${_storageMemoryUsed} bytes")
+      _storageMemoryUsed = 0
+    } else {
+      _storageMemoryUsed -= numBytes
     }
   }
 
   /**
    * Release all storage memory acquired.
    */
-  override def releaseStorageMemory(): Unit = {
-    storageLock.synchronized {
-      _storageMemoryUsed = 0
-    }
+  override def releaseStorageMemory(): Unit = synchronized {
+    _storageMemoryUsed = 0
   }
 
   /**
@@ -189,19 +164,15 @@ private[spark] class StaticMemoryManager(
   /**
    * Amount of execution memory currently in use, in bytes.
    */
-  override def executionMemoryUsed: Long = {
-    executionLock.synchronized {
-      _executionMemoryUsed
-    }
+  override def executionMemoryUsed: Long = synchronized {
+    _executionMemoryUsed
   }
 
   /**
    * Amount of storage memory currently in use, in bytes.
    */
-  override def storageMemoryUsed: Long = {
-    storageLock.synchronized {
-      _storageMemoryUsed
-    }
+  override def storageMemoryUsed: Long = synchronized {
+    _storageMemoryUsed
   }
 
 }

core/src/main/scala/org/apache/spark/shuffle/ShuffleMemoryManager.scala

@@ -165,7 +165,8 @@ private[spark] object ShuffleMemoryManager {
    */
   def create(maxMemory: Long, pageSizeBytes: Long): ShuffleMemoryManager = {
     val conf = new SparkConf
-    val memoryManager = new StaticMemoryManager(conf, maxMemory, Long.MaxValue)
+    val memoryManager = new StaticMemoryManager(
+      conf, maxExecutionMemory = maxMemory, maxStorageMemory = Long.MaxValue)
     new ShuffleMemoryManager(memoryManager, pageSizeBytes)
   }
 

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

@@ -89,6 +89,7 @@ private[spark] class BlockManager(
     externalBlockStoreInitialized = true
     new ExternalBlockStore(this, executorId)
   }
+  memoryManager.setMemoryStore(memoryStore)
 
   private val maxMemory = memoryManager.maxStorageMemory
 

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

@@ -496,6 +496,7 @@ private[spark] class MemoryStore(blockManager: BlockManager, memoryManager: Memo
       memory: Long,
       droppedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = {
     accountingLock.synchronized {
+      // Note: all acquisitions of unroll memory must be synchronized on `accountingLock`
       val success = memoryManager.acquireUnrollMemory(blockId, memory, droppedBlocks)
       if (success) {
         val taskAttemptId = currentTaskAttemptId()

core/src/test/scala/org/apache/spark/memory/StaticMemoryManagerSuite.scala

@@ -38,9 +38,9 @@ class StaticMemoryManagerSuite extends SparkFunSuite {
     assert(mm.acquireExecutionMemory(100L) === 100L)
     // Acquire up to the max
     assert(mm.acquireExecutionMemory(1000L) === 890L)
-    assert(mm.executionMemoryUsed === 1000L)
+    assert(mm.executionMemoryUsed === maxExecutionMem)
     assert(mm.acquireExecutionMemory(1L) === 0L)
-    assert(mm.executionMemoryUsed === 1000L)
+    assert(mm.executionMemoryUsed === maxExecutionMem)
     mm.releaseExecutionMemory(800L)
     assert(mm.executionMemoryUsed === 200L)
     // Acquire after release
@@ -54,35 +54,36 @@ class StaticMemoryManagerSuite extends SparkFunSuite {
   test("basic storage memory") {
     val maxStorageMem = 1000L
     val dummyBlock = TestBlockId("you can see the world you brought to live")
-    val dummyBlocks = new ArrayBuffer[(BlockId, BlockStatus)]
+    val evictedBlocks = new ArrayBuffer[(BlockId, BlockStatus)]
     val (mm, ms) = makeThings(Long.MaxValue, maxStorageMem)
     assert(mm.storageMemoryUsed === 0L)
-    assert(mm.acquireStorageMemory(dummyBlock, 10L, dummyBlocks))
+    assert(mm.acquireStorageMemory(dummyBlock, 10L, evictedBlocks))
     // `ensureFreeSpace` should be called with the number of bytes requested
     assertEnsureFreeSpaceCalled(ms, dummyBlock, 10L)
     assert(mm.storageMemoryUsed === 10L)
-    assert(dummyBlocks.isEmpty)
-    assert(mm.acquireStorageMemory(dummyBlock, 100L, dummyBlocks))
+    assert(evictedBlocks.isEmpty)
+    assert(mm.acquireStorageMemory(dummyBlock, 100L, evictedBlocks))
     assertEnsureFreeSpaceCalled(ms, dummyBlock, 100L)
+    assert(mm.storageMemoryUsed === 110L)
     // Acquire up to the max, not granted
-    assert(!mm.acquireStorageMemory(dummyBlock, 1000L, dummyBlocks))
+    assert(!mm.acquireStorageMemory(dummyBlock, 1000L, evictedBlocks))
     assertEnsureFreeSpaceCalled(ms, dummyBlock, 1000L)
     assert(mm.storageMemoryUsed === 110L)
-    assert(mm.acquireStorageMemory(dummyBlock, 890L, dummyBlocks))
+    assert(mm.acquireStorageMemory(dummyBlock, 890L, evictedBlocks))
     assertEnsureFreeSpaceCalled(ms, dummyBlock, 890L)
     assert(mm.storageMemoryUsed === 1000L)
-    assert(!mm.acquireStorageMemory(dummyBlock, 1L, dummyBlocks))
+    assert(!mm.acquireStorageMemory(dummyBlock, 1L, evictedBlocks))
     assertEnsureFreeSpaceCalled(ms, dummyBlock, 1L)
     assert(mm.storageMemoryUsed === 1000L)
     mm.releaseStorageMemory(800L)
     assert(mm.storageMemoryUsed === 200L)
     // Acquire after release
-    assert(mm.acquireStorageMemory(dummyBlock, 1L, dummyBlocks))
+    assert(mm.acquireStorageMemory(dummyBlock, 1L, evictedBlocks))
     assertEnsureFreeSpaceCalled(ms, dummyBlock, 1L)
     assert(mm.storageMemoryUsed === 201L)
     mm.releaseStorageMemory()
     assert(mm.storageMemoryUsed === 0L)
-    assert(mm.acquireStorageMemory(dummyBlock, 1L, dummyBlocks))
+    assert(mm.acquireStorageMemory(dummyBlock, 1L, evictedBlocks))
     assertEnsureFreeSpaceCalled(ms, dummyBlock, 1L)
     assert(mm.storageMemoryUsed === 1L)
     // Release beyond what was acquired
@@ -150,7 +151,8 @@ class StaticMemoryManagerSuite extends SparkFunSuite {
   private def makeThings(
       maxExecutionMem: Long,
       maxStorageMem: Long): (StaticMemoryManager, MemoryStore) = {
-    val mm = new StaticMemoryManager(conf, maxExecutionMem, maxStorageMem)
+    val mm = new StaticMemoryManager(
+      conf, maxExecutionMemory = maxExecutionMem, maxStorageMemory = maxStorageMem)
     val ms = mock(classOf[MemoryStore])
     mm.setMemoryStore(ms)
     (mm, ms)

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

@@ -1045,13 +1045,12 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with BeforeAndAfterE
   test("reserve/release unroll memory") {
     store = makeBlockManager(12000)
     val memoryStore = store.memoryStore
-    val dummyBlock = TestBlockId("")
-    val dummyEvictedBlocks = new ArrayBuffer[(BlockId, BlockStatus)]
     assert(memoryStore.currentUnrollMemory === 0)
     assert(memoryStore.currentUnrollMemoryForThisTask === 0)
 
     def reserveUnrollMemoryForThisTask(memory: Long): Boolean = {
-      memoryStore.reserveUnrollMemoryForThisTask(dummyBlock, memory, dummyEvictedBlocks)
+      memoryStore.reserveUnrollMemoryForThisTask(
+        TestBlockId(""), memory, new ArrayBuffer[(BlockId, BlockStatus)])
     }
 
     // Reserve