Put fail-fast for non-fitting blocks earlier in call chain.

Author: JoshRosen

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

@@ -57,7 +57,14 @@ private[spark] class StaticMemoryManager(
       blockId: BlockId,
       numBytes: Long,
       evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = synchronized {
-    storageMemoryPool.acquireMemory(blockId, numBytes, evictedBlocks)
+    if (numBytes > storageMemoryPool.poolSize) {
+      // Fail fast if the block simply won't fit
+      logInfo(s"Will not store $blockId as the required space ($numBytes bytes) exceeds our " +
+        s"memory limit (${storageMemoryPool.poolSize} bytes)")
+      false
+    } else {
+      storageMemoryPool.acquireMemory(blockId, numBytes, evictedBlocks)
+    }
   }
 
   override def acquireUnrollMemory(

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

@@ -110,6 +110,12 @@ private[spark] class UnifiedMemoryManager private[memory] (
       evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = synchronized {
     assert(onHeapExecutionMemoryPool.poolSize + storageMemoryPool.poolSize == maxMemory)
     assert(numBytes >= 0)
+    if (numBytes > maxMemory) {
+      // Fail fast if the block simply won't fit
+      logInfo(s"Will not store $blockId as the required space ($numBytes bytes) exceeds our " +
+        s"memory limit ($maxMemory bytes)")
+      return false
+    }
     if (numBytes > storageMemoryPool.memoryFree) {
       // There is not enough free memory in the storage pool, so try to borrow free memory from
       // the execution pool.

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

@@ -454,12 +454,9 @@ private[spark] class MemoryStore(blockManager: BlockManager, memoryManager: Memo
         blockId.map { id => s"for block $id" }.getOrElse("") +
         s"(free: $freeMemory, max: $maxMemory)")
 
-      if (space > maxMemory) {
-        // Fail fast if the block simply won't fit
-        logInfo("Will not " + blockId.map { id => s"store $id" }.getOrElse("free memory") +
-          s" as the required space ($space bytes) exceeds our memory limit ($maxMemory bytes)")
-        false
-      } else if (freeMemory >= space) {
+      assert(space <= maxMemory)
+
+      if (freeMemory >= space) {
         // No need to evict anything if there is already enough free space
         true
       } else {

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

@@ -90,7 +90,7 @@ class StaticMemoryManagerSuite extends MemoryManagerSuite {
     assert(mm.storageMemoryUsed === 110L)
     // Acquire more than the max, not granted
     assert(!mm.acquireStorageMemory(dummyBlock, maxStorageMem + 1L, evictedBlocks))
-    assertEnsureFreeSpaceCalled(ms, maxStorageMem + 1L)
+    assertEnsureFreeSpaceNotCalled(ms)
     assert(mm.storageMemoryUsed === 110L)
     // Acquire up to the max, requests after this are still granted due to LRU eviction
     assert(mm.acquireStorageMemory(dummyBlock, maxStorageMem, evictedBlocks))

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

@@ -84,7 +84,7 @@ class UnifiedMemoryManagerSuite extends MemoryManagerSuite with PrivateMethodTes
     assert(evictedBlocks.isEmpty)
     // Acquire more than the max, not granted
     assert(!mm.acquireStorageMemory(dummyBlock, maxMemory + 1L, evictedBlocks))
-    assertEnsureFreeSpaceCalled(ms, maxMemory + 1L)
+    assertEnsureFreeSpaceNotCalled(ms)
     assert(mm.storageMemoryUsed === 110L)
     assert(evictedBlocks.isEmpty)
     // Acquire up to the max, requests after this are still granted due to LRU eviction