Simplify logic + update a few comments

Author: andrewor14

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

@@ -222,23 +222,13 @@ private[spark] class MemoryStore(blockManager: BlockManager, maxMemory: Long)
     var memoryThreshold = initialMemoryThreshold
     // Memory to request as a multiple of current vector size
     val memoryGrowthFactor = 1.5
-    // Previous unroll memory held by this thread, for releasing later
+    // Previous unroll memory held by this thread, for releasing later (only at the very end)
     val previousMemoryReserved = currentUnrollMemoryForThisThread
     // Underlying vector for unrolling the block
     var vector = new SizeTrackingVector[Any]
 
-    // Request additional memory for our vector and return whether the request is granted
-    // This involves synchronizing across all threads, which is expensive if called frequently
-    def requestMemory(memoryToRequest: Long): Boolean = {
-      accountingLock.synchronized {
-        val granted = freeMemory > currentUnrollMemory + memoryToRequest
-        if (granted) { reserveUnrollMemoryForThisThread(memoryToRequest) }
-        granted
-      }
-    }
-
     // Request enough memory to begin unrolling
-    keepUnrolling = requestMemory(initialMemoryThreshold)
+    keepUnrolling = reserveUnrollMemoryForThisThread(initialMemoryThreshold)
 
     // Unroll this block safely, checking whether we have exceeded our threshold periodically
     try {
@@ -249,18 +239,18 @@ private[spark] class MemoryStore(blockManager: BlockManager, maxMemory: Long)
           val currentSize = vector.estimateSize()
           if (currentSize >= memoryThreshold) {
             val amountToRequest = (currentSize * (memoryGrowthFactor - 1)).toLong
-            // Hold the put lock, in case another thread concurrently puts a block that takes
-            // up the unrolling space we just ensured here
+            // Hold the accounting lock, in case another thread concurrently puts a block that
+            // takes up the unrolling space we just ensured here
             accountingLock.synchronized {
-              if (!requestMemory(amountToRequest)) {
+              if (!reserveUnrollMemoryForThisThread(amountToRequest)) {
                 // If the first request is not granted, try again after ensuring free space
                 // If there is still not enough space, give up and drop the partition
                 val spaceToEnsure = maxUnrollMemory - currentUnrollMemory
                 if (spaceToEnsure > 0) {
                   val result = ensureFreeSpace(blockId, spaceToEnsure)
                   droppedBlocks ++= result.droppedBlocks
                 }
-                keepUnrolling = requestMemory(amountToRequest)
+                keepUnrolling = reserveUnrollMemoryForThisThread(amountToRequest)
               }
             }
             // New threshold is currentSize * memoryGrowthFactor
@@ -283,7 +273,6 @@ private[spark] class MemoryStore(blockManager: BlockManager, maxMemory: Long)
       // we can immediately free up space for other threads. Otherwise, if we return an iterator,
       // we release the memory claimed by this thread later on when the task finishes.
       if (keepUnrolling) {
-        vector = null
         val amountToRelease = currentUnrollMemoryForThisThread - previousMemoryReserved
         releaseUnrollMemoryForThisThread(amountToRelease)
       }
@@ -433,11 +422,16 @@ private[spark] class MemoryStore(blockManager: BlockManager, maxMemory: Long)
 
   /**
    * Reserve additional memory for unrolling blocks used by this thread.
+   * Return whether the request is granted.
    */
-  private[spark] def reserveUnrollMemoryForThisThread(memory: Long): Unit = {
-    val threadId = Thread.currentThread().getId
+  private[spark] def reserveUnrollMemoryForThisThread(memory: Long): Boolean = {
     accountingLock.synchronized {
-      unrollMemoryMap(threadId) = unrollMemoryMap.getOrElse(threadId, 0L) + memory
+      val granted = freeMemory > currentUnrollMemory + memory
+      if (granted) {
+        val threadId = Thread.currentThread().getId
+        unrollMemoryMap(threadId) = unrollMemoryMap.getOrElse(threadId, 0L) + memory
+      }
+      granted
     }
   }
 

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

@@ -1005,21 +1005,24 @@ class BlockManagerSuite extends FunSuite with Matchers with BeforeAndAfter
     assert(memoryStore.currentUnrollMemoryForThisThread === 0)
 
     // Reserve
-    memoryStore.reserveUnrollMemoryForThisThread(1480)
-    assert(memoryStore.currentUnrollMemoryForThisThread === 1480)
+    memoryStore.reserveUnrollMemoryForThisThread(100)
+    assert(memoryStore.currentUnrollMemoryForThisThread === 100)
+    memoryStore.reserveUnrollMemoryForThisThread(200)
+    assert(memoryStore.currentUnrollMemoryForThisThread === 300)
+    memoryStore.reserveUnrollMemoryForThisThread(500)
+    assert(memoryStore.currentUnrollMemoryForThisThread === 800)
     memoryStore.reserveUnrollMemoryForThisThread(1000000)
-    assert(memoryStore.currentUnrollMemoryForThisThread === 1001480)
-
+    assert(memoryStore.currentUnrollMemoryForThisThread === 800) // not granted
     // Release
-    memoryStore.releaseUnrollMemoryForThisThread(1000000)
-    assert(memoryStore.currentUnrollMemoryForThisThread === 1480)
     memoryStore.releaseUnrollMemoryForThisThread(100)
-    assert(memoryStore.currentUnrollMemoryForThisThread === 1380)
-
+    assert(memoryStore.currentUnrollMemoryForThisThread === 700)
+    memoryStore.releaseUnrollMemoryForThisThread(100)
+    assert(memoryStore.currentUnrollMemoryForThisThread === 600)
     // Reserve again
-    memoryStore.reserveUnrollMemoryForThisThread(3620)
+    memoryStore.reserveUnrollMemoryForThisThread(4400)
     assert(memoryStore.currentUnrollMemoryForThisThread === 5000)
-
+    memoryStore.reserveUnrollMemoryForThisThread(20000)
+    assert(memoryStore.currentUnrollMemoryForThisThread === 5000) // not granted
     // Release again
     memoryStore.releaseUnrollMemoryForThisThread(1000)
     assert(memoryStore.currentUnrollMemoryForThisThread === 4000)