Move pooling logic into allocators themselves.

core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java

@@ -163,7 +163,7 @@ public MemoryBlock allocatePage(long size) {
       }
       return null;
     }
-    final MemoryBlock page = memoryManager.allocateMemoryBlock(size);
+    final MemoryBlock page = memoryManager.tungstenMemoryAllocator().allocate(size);
     page.pageNumber = pageNumber;
     pageTable[pageNumber] = page;
     if (logger.isTraceEnabled()) {
@@ -187,7 +187,7 @@ public void freePage(MemoryBlock page) {
       logger.trace("Freed page number {} ({} bytes)", page.pageNumber, page.size());
     }
     long pageSize = page.size();
-    memoryManager.releaseMemoryBlock(page);
+    memoryManager.tungstenMemoryAllocator().free(page);
     releaseExecutionMemory(pageSize);
   }
 
@@ -202,7 +202,7 @@ public void freePage(MemoryBlock page) {
    */
   public MemoryBlock allocate(long size) throws OutOfMemoryError {
     assert(size > 0) : "Size must be positive, but got " + size;
-    final MemoryBlock memory = memoryManager.allocateMemoryBlock(size);
+    final MemoryBlock memory = memoryManager.tungstenMemoryAllocator().allocate(size);
     synchronized(allocatedNonPageMemory) {
       allocatedNonPageMemory.add(memory);
     }
@@ -214,7 +214,7 @@ public MemoryBlock allocate(long size) throws OutOfMemoryError {
    */
   public void free(MemoryBlock memory) {
     assert (memory.pageNumber == -1) : "Should call freePage() for pages, not free()";
-    memoryManager.releaseMemoryBlock(memory);
+    memoryManager.tungstenMemoryAllocator().free(memory);
     synchronized(allocatedNonPageMemory) {
       final boolean wasAlreadyRemoved = !allocatedNonPageMemory.remove(memory);
       assert (!wasAlreadyRemoved) : "Called free() on memory that was already freed!";
@@ -312,7 +312,7 @@ public long cleanUpAllAllocatedMemory() {
         freedBytes += memory.size();
         // We don't call free() here because that calls Set.remove, which would lead to a
         // ConcurrentModificationException here.
-        memoryManager.releaseMemoryBlock(memory);
+        memoryManager.tungstenMemoryAllocator().free(memory);
         iter.remove();
       }
     }

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

@@ -17,8 +17,6 @@
 
 package org.apache.spark.memory
 
-import java.lang.ref.WeakReference
-import java.util
 import javax.annotation.concurrent.GuardedBy
 
 import scala.collection.mutable
@@ -27,7 +25,7 @@ import scala.collection.mutable.ArrayBuffer
 import org.apache.spark.{SparkException, TaskContext, SparkConf, Logging}
 import org.apache.spark.storage.{BlockId, BlockStatus, MemoryStore}
 import org.apache.spark.unsafe.array.ByteArrayMethods
-import org.apache.spark.unsafe.memory.{MemoryAllocator, MemoryBlock}
+import org.apache.spark.unsafe.memory.MemoryAllocator
 
 /**
  * An abstract memory manager that enforces how memory is shared between execution and storage.
@@ -313,72 +311,6 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int = 1)
    * Allocates memory for use by Unsafe/Tungsten code. Exposed to enable untracked allocations of
    * temporary data structures.
    */
-  final val tungstenMemoryAllocator: MemoryAllocator =
+  private[memory] final val tungstenMemoryAllocator: MemoryAllocator =
     if (tungstenMemoryIsAllocatedInHeap) MemoryAllocator.HEAP else MemoryAllocator.UNSAFE
-
-  private val POOLING_THRESHOLD_BYTES: Int = 1024 * 1024
-
-  /**
-   * Returns true if allocations of the given size should go through the pooling mechanism and
-   * false otherwise.
-   */
-  private def shouldPool(size: Long): Boolean = {
-    // Very small allocations are less likely to benefit from pooling.
-    // At some point, we should explore supporting pooling for off-heap memory, but for now we'll
-    // ignore that case in the interest of simplicity.
-    size >= POOLING_THRESHOLD_BYTES && tungstenMemoryIsAllocatedInHeap
-  }
-
-  @GuardedBy("this")
-  private val bufferPoolsBySize: util.Map[Long, util.LinkedList[WeakReference[MemoryBlock]]] =
-    new util.HashMap[Long, util.LinkedList[WeakReference[MemoryBlock]]]
-
-  /**
-   * Allocates a contiguous block of memory. Note that the allocated memory is not guaranteed
-   * to be zeroed out (call `zero()` on the result if this is necessary). This method does
-   * not integrate with the memory bookkeeping system, so callers (i.e. TaskMemoryManager) should
-   * call those methods at appropirate times.
-   */
-  @throws(classOf[OutOfMemoryError])
-  private[memory] final def allocateMemoryBlock(size: Long): MemoryBlock = {
-    if (shouldPool(size)) {
-      this synchronized {
-        val pool: util.LinkedList[WeakReference[MemoryBlock]] = bufferPoolsBySize.get(size)
-        if (pool != null) {
-          while (!pool.isEmpty) {
-            val blockReference: WeakReference[MemoryBlock] = pool.pop
-            val memory: MemoryBlock = blockReference.get
-            if (memory != null) {
-              assert(memory.size == size)
-              return memory
-            }
-          }
-          bufferPoolsBySize.remove(size)
-        }
-      }
-      tungstenMemoryAllocator.allocate(size)
-    } else {
-      tungstenMemoryAllocator.allocate(size)
-    }
-  }
-
-  /**
-   * Releases the given memory block, either freeing it immediately or storing it in a pool for
-   * re-use by other tasks.
-   */
-  private[memory] final def releaseMemoryBlock(memory: MemoryBlock) {
-    val size: Long = memory.size
-    if (shouldPool(size)) {
-      this synchronized {
-        var pool: util.LinkedList[WeakReference[MemoryBlock]] = bufferPoolsBySize.get(size)
-        if (pool == null) {
-          pool = new util.LinkedList[WeakReference[MemoryBlock]]
-          bufferPoolsBySize.put(size, pool)
-        }
-        pool.add(new WeakReference[MemoryBlock](memory))
-      }
-    } else {
-      tungstenMemoryAllocator.free(memory)
-    }
-  }
 }

unsafe/src/main/java/org/apache/spark/unsafe/memory/HeapMemoryAllocator.java

@@ -17,22 +17,71 @@
 
 package org.apache.spark.unsafe.memory;
 
+import javax.annotation.concurrent.GuardedBy;
+import java.lang.ref.WeakReference;
+import java.util.HashMap;
+import java.util.LinkedList;
+import java.util.Map;
+
 /**
  * A simple {@link MemoryAllocator} that can allocate up to 16GB using a JVM long primitive array.
  */
 public class HeapMemoryAllocator implements MemoryAllocator {
 
+  @GuardedBy("this")
+  private final Map<Long, LinkedList<WeakReference<MemoryBlock>>> bufferPoolsBySize =
+    new HashMap<>();
+
+  private static final int POOLING_THRESHOLD_BYTES = 1024 * 1024;
+
+  /**
+   * Returns true if allocations of the given size should go through the pooling mechanism and
+   * false otherwise.
+   */
+  private boolean shouldPool(long size) {
+    // Very small allocations are less likely to benefit from pooling.
+    return size >= POOLING_THRESHOLD_BYTES;
+  }
+
   @Override
   public MemoryBlock allocate(long size) throws OutOfMemoryError {
     if (size % 8 != 0) {
       throw new IllegalArgumentException("Size " + size + " was not a multiple of 8");
     }
+    if (shouldPool(size)) {
+      synchronized (this) {
+        final LinkedList<WeakReference<MemoryBlock>> pool = bufferPoolsBySize.get(size);
+        if (pool != null) {
+          while (!pool.isEmpty()) {
+            final WeakReference<MemoryBlock> blockReference = pool.pop();
+            final MemoryBlock memory = blockReference.get();
+            if (memory != null) {
+              assert (memory.size() == size);
+              return memory;
+            }
+          }
+          bufferPoolsBySize.remove(size);
+        }
+      }
+    }
     long[] array = new long[(int) (size / 8)];
     return MemoryBlock.fromLongArray(array);
   }
 
   @Override
   public void free(MemoryBlock memory) {
-    // Do nothing
+    final long size = memory.size();
+    if (shouldPool(size)) {
+      synchronized (this) {
+        LinkedList<WeakReference<MemoryBlock>> pool = bufferPoolsBySize.get(size);
+        if (pool == null) {
+          pool = new LinkedList<>();
+          bufferPoolsBySize.put(size, pool);
+        }
+        pool.add(new WeakReference<>(memory));
+      }
+    } else {
+      // Do nothing
+    }
   }
 }