From d0ada7b5c6f8be9a91624e4eb5a6ca2242edf16b Mon Sep 17 00:00:00 2001 From: Davies Liu Date: Mon, 26 Oct 2015 15:29:49 -0700 Subject: [PATCH 01/17] force spilling --- .../apache/spark/memory/MemoryConsumer.java | 100 +++++ .../spark/memory/TaskMemoryManager.java | 106 ++++- .../shuffle/sort/ShuffleExternalSorter.java | 171 ++------ .../shuffle/sort/ShuffleInMemorySorter.java | 6 +- .../shuffle/sort/UnsafeShuffleWriter.java | 2 +- .../spark/unsafe/map/BytesToBytesMap.java | 371 +++++++++-------- .../unsafe/sort/UnsafeExternalSorter.java | 391 +++++++++--------- .../unsafe/sort/UnsafeInMemorySorter.java | 14 +- .../unsafe/sort/UnsafeSorterSpillReader.java | 6 +- .../unsafe/sort/UnsafeSorterSpillWriter.java | 2 +- .../sort/PackedRecordPointerSuite.java | 16 +- .../sort/UnsafeShuffleWriterSuite.java | 11 +- .../map/AbstractBytesToBytesMapSuite.java | 144 ++++++- .../sort/UnsafeExternalSorterSuite.java | 49 +-- .../memory/GrantEverythingMemoryManager.scala | 21 +- .../execution/UnsafeExternalRowSorter.java | 2 +- .../UnsafeFixedWidthAggregationMap.java | 2 +- .../sql/execution/UnsafeKVExternalSorter.java | 9 +- .../UnsafeFixedWidthAggregationMapSuite.scala | 10 +- .../UnsafeKVExternalSorterSuite.scala | 2 +- .../TungstenAggregationIteratorSuite.scala | 54 --- .../unsafe/memory/HeapMemoryAllocator.java | 9 +- .../unsafe/memory/UnsafeMemoryAllocator.java | 3 - 23 files changed, 861 insertions(+), 640 deletions(-) create mode 100644 core/src/main/java/org/apache/spark/memory/MemoryConsumer.java delete mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIteratorSuite.scala diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java new file mode 100644 index 0000000000000..e308c4b74a579 --- /dev/null +++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java @@ -0,0 +1,100 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.memory; + + +import java.io.IOException; + +import org.apache.spark.unsafe.memory.MemoryBlock; + + +/** + * An memory consumer of TaskMemoryManager, which support spilling. + */ +public class MemoryConsumer { + + private TaskMemoryManager memoryManager; + private long pageSize; + + protected MemoryConsumer(TaskMemoryManager memoryManager, long pageSize) { + this.memoryManager = memoryManager; + this.pageSize = pageSize; + } + + protected MemoryConsumer(TaskMemoryManager memoryManager) { + this(memoryManager, memoryManager.pageSizeBytes()); + } + + /** + * Spill some data to disk to release memory, which will be called by TaskMemoryManager + * when there is not enough memory for the task. + * + * @param size the amount of memory should be released + * @return the amount of released memory in bytes + * @throws IOException + */ + public long spill(long size) throws IOException { + return 0L; + } + + /** + * Acquire `size` bytes memory. + * + * If there is not enough memory, throws IOException. + * + * @throws IOException + */ + protected void acquireMemory(long size) throws IOException { + long got = memoryManager.acquireExecutionMemory(size, this); + if (got < size) { + throw new IOException("Could not acquire " + size + " bytes of memory " + got); + } + } + + /** + * Release amount of memory. + */ + protected void releaseMemory(long size) { + memoryManager.releaseExecutionMemory(size, this); + } + + /** + * Allocate a memory block with at least `required` bytes. + * + * Throws IOException if there is not enough memory. + * + * @throws IOException + */ + protected MemoryBlock allocatePage(long required) throws IOException { + MemoryBlock page = memoryManager.allocatePage(Math.max(pageSize, required), this); + if (page == null || page.size() < required) { + if (page != null) { + freePage(page); + } + throw new IOException("Unable to acquire " + required + " bytes of memory"); + } + return page; + } + + /** + * Free a memory block. + */ + protected void freePage(MemoryBlock page) { + memoryManager.freePage(page, this); + } +} diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java index 7b31c90dac666..4f7e98ace60d1 100644 --- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java +++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java @@ -17,7 +17,9 @@ package org.apache.spark.memory; -import java.util.*; +import java.io.IOException; +import java.util.BitSet; +import java.util.HashMap; import com.google.common.annotations.VisibleForTesting; import org.slf4j.Logger; @@ -100,6 +102,11 @@ public class TaskMemoryManager { */ private final boolean inHeap; + /** + * The size of memory granted to each consumer. + */ + private HashMap consumers; + /** * Construct a new TaskMemoryManager. */ @@ -107,6 +114,7 @@ public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) { this.inHeap = memoryManager.tungstenMemoryIsAllocatedInHeap(); this.memoryManager = memoryManager; this.taskAttemptId = taskAttemptId; + this.consumers = new HashMap<>(); } /** @@ -117,6 +125,46 @@ public long acquireExecutionMemory(long size) { return memoryManager.acquireExecutionMemory(size, taskAttemptId); } + /** + * Acquire N bytes of memory for a consumer. If there is no enough memory, it will call + * spill() of consumers to release more memory. + * + * @return number of bytes successfully granted (<= N). + */ + public long acquireExecutionMemory(long size, MemoryConsumer consumer) throws IOException { + synchronized (this) { + long got = acquireExecutionMemory(size); + + if (got < size && consumer != null) { + // call spill() on itself to release some memory + consumer.spill(size - got); + got += acquireExecutionMemory(size - got); + + if (got < size) { + long needed = size - got; + // call spill() on other consumers to release memory + for (MemoryConsumer c: consumers.keySet()) { + if (c != consumer) { + needed -= c.spill(size - got); + if (needed < 0) { + break; + } + } + } + got += acquireExecutionMemory(size - got); + } + } + + long old = 0L; + if (consumers.containsKey(consumer)) { + old = consumers.get(consumer); + } + consumers.put(consumer, got + old); + + return got; + } + } + /** * Release N bytes of execution memory. */ @@ -124,6 +172,28 @@ public void releaseExecutionMemory(long size) { memoryManager.releaseExecutionMemory(size, taskAttemptId); } + /** + * Release N bytes of execution memory for a MemoryConsumer. + */ + public void releaseExecutionMemory(long size, MemoryConsumer consumer) { + synchronized (this) { + if (consumer != null && consumers.containsKey(consumer)) { + long old = consumers.get(consumer); + if (old > size) { + consumers.put(consumer, old - size); + } else { + if (old < size) { + // TODO + } + consumers.remove(consumer); + } + } else { + // TODO + } + memoryManager.releaseExecutionMemory(size, taskAttemptId); + } + } + public long pageSizeBytes() { return memoryManager.pageSizeBytes(); } @@ -134,12 +204,27 @@ public long pageSizeBytes() { * * Returns `null` if there was not enough memory to allocate the page. */ - public MemoryBlock allocatePage(long size) { + public MemoryBlock allocatePage(long size) throws IOException { + return allocatePage(size, null); + } + + /** + * Allocate a block of memory that will be tracked in the MemoryManager's page table; this is + * intended for allocating large blocks of Tungsten memory that will be shared between operators. + * + * Returns `null` if there was not enough memory to allocate the page. + */ + public MemoryBlock allocatePage(long size, MemoryConsumer consumer) throws IOException { if (size > MAXIMUM_PAGE_SIZE_BYTES) { throw new IllegalArgumentException( "Cannot allocate a page with more than " + MAXIMUM_PAGE_SIZE_BYTES + " bytes"); } + long acquired = acquireExecutionMemory(size, consumer); + if (acquired <= 0) { + return null; + } + final int pageNumber; synchronized (this) { pageNumber = allocatedPages.nextClearBit(0); @@ -149,14 +234,6 @@ public MemoryBlock allocatePage(long size) { } allocatedPages.set(pageNumber); } - final long acquiredExecutionMemory = acquireExecutionMemory(size); - if (acquiredExecutionMemory != size) { - releaseExecutionMemory(acquiredExecutionMemory); - synchronized (this) { - allocatedPages.clear(pageNumber); - } - return null; - } final MemoryBlock page = memoryManager.tungstenMemoryAllocator().allocate(size); page.pageNumber = pageNumber; pageTable[pageNumber] = page; @@ -170,6 +247,13 @@ public MemoryBlock allocatePage(long size) { * Free a block of memory allocated via {@link TaskMemoryManager#allocatePage(long)}. */ public void freePage(MemoryBlock page) { + freePage(page, null); + } + + /** + * Free a block of memory allocated via {@link TaskMemoryManager#allocatePage(long)}. + */ + public void freePage(MemoryBlock page, MemoryConsumer consumer) { assert (page.pageNumber != -1) : "Called freePage() on memory that wasn't allocated with allocatePage()"; assert(allocatedPages.get(page.pageNumber)); @@ -182,7 +266,7 @@ public void freePage(MemoryBlock page) { } long pageSize = page.size(); memoryManager.tungstenMemoryAllocator().free(page); - releaseExecutionMemory(pageSize); + releaseExecutionMemory(pageSize, consumer); } /** diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java index f43236f41ae7b..69ac697a03511 100644 --- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java +++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java @@ -31,15 +31,15 @@ import org.apache.spark.SparkConf; import org.apache.spark.TaskContext; import org.apache.spark.executor.ShuffleWriteMetrics; +import org.apache.spark.memory.MemoryConsumer; +import org.apache.spark.memory.TaskMemoryManager; import org.apache.spark.serializer.DummySerializerInstance; import org.apache.spark.serializer.SerializerInstance; import org.apache.spark.storage.BlockManager; import org.apache.spark.storage.DiskBlockObjectWriter; import org.apache.spark.storage.TempShuffleBlockId; import org.apache.spark.unsafe.Platform; -import org.apache.spark.unsafe.array.ByteArrayMethods; import org.apache.spark.unsafe.memory.MemoryBlock; -import org.apache.spark.memory.TaskMemoryManager; import org.apache.spark.util.Utils; /** @@ -58,14 +58,13 @@ * spill files. Instead, this merging is performed in {@link UnsafeShuffleWriter}, which uses a * specialized merge procedure that avoids extra serialization/deserialization. */ -final class ShuffleExternalSorter { +final class ShuffleExternalSorter extends MemoryConsumer { private final Logger logger = LoggerFactory.getLogger(ShuffleExternalSorter.class); @VisibleForTesting static final int DISK_WRITE_BUFFER_SIZE = 1024 * 1024; - private final int initialSize; private final int numPartitions; private final int pageSizeBytes; @VisibleForTesting @@ -98,8 +97,7 @@ final class ShuffleExternalSorter { // These variables are reset after spilling: @Nullable private ShuffleInMemorySorter inMemSorter; @Nullable private MemoryBlock currentPage = null; - private long currentPagePosition = -1; - private long freeSpaceInCurrentPage = 0; + private long pageCursor = -1; public ShuffleExternalSorter( TaskMemoryManager memoryManager, @@ -109,10 +107,10 @@ public ShuffleExternalSorter( int numPartitions, SparkConf conf, ShuffleWriteMetrics writeMetrics) throws IOException { + super(memoryManager); this.taskMemoryManager = memoryManager; this.blockManager = blockManager; this.taskContext = taskContext; - this.initialSize = initialSize; this.peakMemoryUsedBytes = initialSize; this.numPartitions = numPartitions; // Use getSizeAsKb (not bytes) to maintain backwards compatibility if no units are provided @@ -123,25 +121,7 @@ public ShuffleExternalSorter( PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES, taskMemoryManager.pageSizeBytes()); this.maxRecordSizeBytes = pageSizeBytes - 4; this.writeMetrics = writeMetrics; - initializeForWriting(); - - // preserve first page to ensure that we have at least one page to work with. Otherwise, - // other operators in the same task may starve this sorter (SPARK-9709). - acquireNewPageIfNecessary(pageSizeBytes); - } - - /** - * Allocates new sort data structures. Called when creating the sorter and after each spill. - */ - private void initializeForWriting() throws IOException { - // TODO: move this sizing calculation logic into a static method of sorter: - final long memoryRequested = initialSize * 8L; - final long memoryAcquired = taskMemoryManager.acquireExecutionMemory(memoryRequested); - if (memoryAcquired != memoryRequested) { - taskMemoryManager.releaseExecutionMemory(memoryAcquired); - throw new IOException("Could not acquire " + memoryRequested + " bytes of memory"); - } - + acquireMemory(initialSize * 8L); this.inMemSorter = new ShuffleInMemorySorter(initialSize); numRecordsInsertedSinceLastSpill = 0; } @@ -242,6 +222,8 @@ private void writeSortedFile(boolean isLastFile) throws IOException { } } + inMemSorter.reset(); + if (!isLastFile) { // i.e. this is a spill file // The current semantics of `shuffleRecordsWritten` seem to be that it's updated when records // are written to disk, not when they enter the shuffle sorting code. DiskBlockObjectWriter @@ -267,8 +249,8 @@ private void writeSortedFile(boolean isLastFile) throws IOException { * Sort and spill the current records in response to memory pressure. */ @VisibleForTesting - void spill() throws IOException { - assert(inMemSorter != null); + @Override + public long spill(long size) throws IOException { logger.info("Thread {} spilling sort data of {} to disk ({} {} so far)", Thread.currentThread().getId(), Utils.bytesToString(getMemoryUsage()), @@ -276,13 +258,9 @@ void spill() throws IOException { spills.size() > 1 ? " times" : " time"); writeSortedFile(false); - final long inMemSorterMemoryUsage = inMemSorter.getMemoryUsage(); - inMemSorter = null; - taskMemoryManager.releaseExecutionMemory(inMemSorterMemoryUsage); final long spillSize = freeMemory(); taskContext.taskMetrics().incMemoryBytesSpilled(spillSize); - - initializeForWriting(); + return spillSize; } private long getMemoryUsage() { @@ -312,18 +290,12 @@ private long freeMemory() { updatePeakMemoryUsed(); long memoryFreed = 0; for (MemoryBlock block : allocatedPages) { - taskMemoryManager.freePage(block); + freePage(block); memoryFreed += block.size(); } - if (inMemSorter != null) { - long sorterMemoryUsage = inMemSorter.getMemoryUsage(); - inMemSorter = null; - taskMemoryManager.releaseExecutionMemory(sorterMemoryUsage); - } allocatedPages.clear(); currentPage = null; - currentPagePosition = -1; - freeSpaceInCurrentPage = 0; + pageCursor = 0; return memoryFreed; } @@ -332,16 +304,16 @@ private long freeMemory() { */ public void cleanupResources() { freeMemory(); + if (inMemSorter != null) { + long sorterMemoryUsage = inMemSorter.getMemoryUsage(); + inMemSorter = null; + releaseMemory(sorterMemoryUsage); + } for (SpillInfo spill : spills) { if (spill.file.exists() && !spill.file.delete()) { logger.error("Unable to delete spill file {}", spill.file.getPath()); } } - if (inMemSorter != null) { - long sorterMemoryUsage = inMemSorter.getMemoryUsage(); - inMemSorter = null; - taskMemoryManager.releaseExecutionMemory(sorterMemoryUsage); - } } /** @@ -354,15 +326,8 @@ private void growPointerArrayIfNecessary() throws IOException { if (!inMemSorter.hasSpaceForAnotherRecord()) { logger.debug("Attempting to expand sort pointer array"); final long oldPointerArrayMemoryUsage = inMemSorter.getMemoryUsage(); - final long memoryToGrowPointerArray = oldPointerArrayMemoryUsage * 2; - final long memoryAcquired = taskMemoryManager.acquireExecutionMemory(memoryToGrowPointerArray); - if (memoryAcquired < memoryToGrowPointerArray) { - taskMemoryManager.releaseExecutionMemory(memoryAcquired); - spill(); - } else { - inMemSorter.expandPointerArray(); - taskMemoryManager.releaseExecutionMemory(oldPointerArrayMemoryUsage); - } + acquireMemory(oldPointerArrayMemoryUsage); //TODO actual memory may be less + inMemSorter.expandPointerArray(); } } @@ -370,93 +335,42 @@ private void growPointerArrayIfNecessary() throws IOException { * Allocates more memory in order to insert an additional record. This will request additional * memory from the memory manager and spill if the requested memory can not be obtained. * - * @param requiredSpace the required space in the data page, in bytes, including space for storing + * @param required the required space in the data page, in bytes, including space for storing * the record size. This must be less than or equal to the page size (records * that exceed the page size are handled via a different code path which uses * special overflow pages). */ - private void acquireNewPageIfNecessary(int requiredSpace) throws IOException { + private void acquireNewPageIfNecessary(int required) throws IOException { growPointerArrayIfNecessary(); - if (requiredSpace > freeSpaceInCurrentPage) { - logger.trace("Required space {} is less than free space in current page ({})", requiredSpace, - freeSpaceInCurrentPage); - // TODO: we should track metrics on the amount of space wasted when we roll over to a new page - // without using the free space at the end of the current page. We should also do this for - // BytesToBytesMap. - if (requiredSpace > pageSizeBytes) { - throw new IOException("Required space " + requiredSpace + " is greater than page size (" + - pageSizeBytes + ")"); - } else { - currentPage = taskMemoryManager.allocatePage(pageSizeBytes); - if (currentPage == null) { - spill(); - currentPage = taskMemoryManager.allocatePage(pageSizeBytes); - if (currentPage == null) { - throw new IOException("Unable to acquire " + pageSizeBytes + " bytes of memory"); - } - } - currentPagePosition = currentPage.getBaseOffset(); - freeSpaceInCurrentPage = pageSizeBytes; - allocatedPages.add(currentPage); - } + if (currentPage == null || + pageCursor + required > currentPage.getBaseOffset() + currentPage.size() ) { + // TODO: try to find space in previous pages + currentPage = allocatePage(pageSizeBytes); + pageCursor = currentPage.getBaseOffset(); + allocatedPages.add(currentPage); } } /** * Write a record to the shuffle sorter. */ - public void insertRecord( - Object recordBaseObject, - long recordBaseOffset, - int lengthInBytes, - int partitionId) throws IOException { + public void insertRecord(Object recordBase, long recordOffset, int length, int partitionId) + throws IOException { + // for tests if (numRecordsInsertedSinceLastSpill > numElementsForSpillThreshold) { - spill(); + spill(0); } - growPointerArrayIfNecessary(); - // Need 4 bytes to store the record length. - final int totalSpaceRequired = lengthInBytes + 4; - - // --- Figure out where to insert the new record ---------------------------------------------- - - final MemoryBlock dataPage; - long dataPagePosition; - boolean useOverflowPage = totalSpaceRequired > pageSizeBytes; - if (useOverflowPage) { - long overflowPageSize = ByteArrayMethods.roundNumberOfBytesToNearestWord(totalSpaceRequired); - // The record is larger than the page size, so allocate a special overflow page just to hold - // that record. - MemoryBlock overflowPage = taskMemoryManager.allocatePage(overflowPageSize); - if (overflowPage == null) { - spill(); - overflowPage = taskMemoryManager.allocatePage(overflowPageSize); - if (overflowPage == null) { - throw new IOException("Unable to acquire " + overflowPageSize + " bytes of memory"); - } - } - allocatedPages.add(overflowPage); - dataPage = overflowPage; - dataPagePosition = overflowPage.getBaseOffset(); - } else { - // The record is small enough to fit in a regular data page, but the current page might not - // have enough space to hold it (or no pages have been allocated yet). - acquireNewPageIfNecessary(totalSpaceRequired); - dataPage = currentPage; - dataPagePosition = currentPagePosition; - // Update bookkeeping information - freeSpaceInCurrentPage -= totalSpaceRequired; - currentPagePosition += totalSpaceRequired; - } - final Object dataPageBaseObject = dataPage.getBaseObject(); - - final long recordAddress = - taskMemoryManager.encodePageNumberAndOffset(dataPage, dataPagePosition); - Platform.putInt(dataPageBaseObject, dataPagePosition, lengthInBytes); - dataPagePosition += 4; - Platform.copyMemory( - recordBaseObject, recordBaseOffset, dataPageBaseObject, dataPagePosition, lengthInBytes); + final int required = length + 4; + acquireNewPageIfNecessary(required); + + final Object base = currentPage.getBaseObject(); + final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor); + Platform.putInt(base, pageCursor, length); + pageCursor += 4; + Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length); + pageCursor += length; assert(inMemSorter != null); inMemSorter.insertRecord(recordAddress, partitionId); numRecordsInsertedSinceLastSpill += 1; @@ -475,6 +389,9 @@ public SpillInfo[] closeAndGetSpills() throws IOException { // Do not count the final file towards the spill count. writeSortedFile(true); freeMemory(); + long sorterMemoryUsage = inMemSorter.getMemoryUsage(); + inMemSorter = null; + releaseMemory(sorterMemoryUsage); } return spills.toArray(new SpillInfo[spills.size()]); } catch (IOException e) { diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java index a8dee6c6101c1..ca18e510e8c62 100644 --- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java +++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java @@ -50,6 +50,10 @@ public ShuffleInMemorySorter(int initialSize) { this.sorter = new Sorter(ShuffleSortDataFormat.INSTANCE); } + public void reset() { + pointerArrayInsertPosition = 0; + } + public void expandPointerArray() { final long[] oldArray = pointerArray; // Guard against overflow: @@ -59,7 +63,7 @@ public void expandPointerArray() { } public boolean hasSpaceForAnotherRecord() { - return pointerArrayInsertPosition + 1 < pointerArray.length; + return pointerArrayInsertPosition < pointerArray.length; } public long getMemoryUsage() { diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java index f6c5c944bd77b..fcb3f3064794c 100644 --- a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java +++ b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java @@ -245,7 +245,7 @@ void insertRecordIntoSorter(Product2 record) throws IOException { @VisibleForTesting void forceSorterToSpill() throws IOException { assert (sorter != null); - sorter.spill(); + sorter.spill(Long.MAX_VALUE); } /** diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java index f035bdac810bd..42c6ac34dc22c 100644 --- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java +++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java @@ -18,14 +18,19 @@ package org.apache.spark.unsafe.map; import javax.annotation.Nullable; +import java.io.IOException; import java.util.Iterator; import java.util.LinkedList; -import java.util.List; import com.google.common.annotations.VisibleForTesting; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.apache.spark.SparkEnv; +import org.apache.spark.executor.ShuffleWriteMetrics; +import org.apache.spark.memory.MemoryConsumer; +import org.apache.spark.memory.TaskMemoryManager; +import org.apache.spark.storage.BlockManager; import org.apache.spark.unsafe.Platform; import org.apache.spark.unsafe.array.ByteArrayMethods; import org.apache.spark.unsafe.array.LongArray; @@ -33,7 +38,8 @@ import org.apache.spark.unsafe.hash.Murmur3_x86_32; import org.apache.spark.unsafe.memory.MemoryBlock; import org.apache.spark.unsafe.memory.MemoryLocation; -import org.apache.spark.memory.TaskMemoryManager; +import org.apache.spark.util.collection.unsafe.sort.UnsafeSorterSpillReader; +import org.apache.spark.util.collection.unsafe.sort.UnsafeSorterSpillWriter; /** * An append-only hash map where keys and values are contiguous regions of bytes. @@ -54,7 +60,7 @@ * is consistent with {@link org.apache.spark.util.collection.unsafe.sort.UnsafeExternalSorter}, * so we can pass records from this map directly into the sorter to sort records in place. */ -public final class BytesToBytesMap { +public final class BytesToBytesMap extends MemoryConsumer { private final Logger logger = LoggerFactory.getLogger(BytesToBytesMap.class); @@ -62,27 +68,22 @@ public final class BytesToBytesMap { private static final HashMapGrowthStrategy growthStrategy = HashMapGrowthStrategy.DOUBLING; - /** - * Special record length that is placed after the last record in a data page. - */ - private static final int END_OF_PAGE_MARKER = -1; - private final TaskMemoryManager taskMemoryManager; /** * A linked list for tracking all allocated data pages so that we can free all of our memory. */ - private final List dataPages = new LinkedList(); + private final LinkedList dataPages = new LinkedList<>(); /** * The data page that will be used to store keys and values for new hashtable entries. When this * page becomes full, a new page will be allocated and this pointer will change to point to that * new page. */ - private MemoryBlock currentDataPage = null; + private MemoryBlock currentPage = null; /** - * Offset into `currentDataPage` that points to the location where new data can be inserted into + * Offset into `currentPage` that points to the location where new data can be inserted into * the page. This does not incorporate the page's base offset. */ private long pageCursor = 0; @@ -164,13 +165,19 @@ public final class BytesToBytesMap { private long peakMemoryUsedBytes = 0L; + private final BlockManager blockManager; + private MapIterator destructiveIterator = null; + public BytesToBytesMap( TaskMemoryManager taskMemoryManager, + BlockManager blockManager, int initialCapacity, double loadFactor, long pageSizeBytes, boolean enablePerfMetrics) { + super(taskMemoryManager, pageSizeBytes); this.taskMemoryManager = taskMemoryManager; + this.blockManager = blockManager; this.loadFactor = loadFactor; this.loc = new Location(); this.pageSizeBytes = pageSizeBytes; @@ -187,18 +194,13 @@ public BytesToBytesMap( TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES); } allocate(initialCapacity); - - // Acquire a new page as soon as we construct the map to ensure that we have at least - // one page to work with. Otherwise, other operators in the same task may starve this - // map (SPARK-9747). - acquireNewPage(); } public BytesToBytesMap( TaskMemoryManager taskMemoryManager, int initialCapacity, long pageSizeBytes) { - this(taskMemoryManager, initialCapacity, 0.70, pageSizeBytes, false); + this(taskMemoryManager, initialCapacity, pageSizeBytes, false); } public BytesToBytesMap( @@ -208,6 +210,7 @@ public BytesToBytesMap( boolean enablePerfMetrics) { this( taskMemoryManager, + SparkEnv.get() != null ? SparkEnv.get().blockManager() : null, initialCapacity, 0.70, pageSizeBytes, @@ -219,61 +222,146 @@ public BytesToBytesMap( */ public int numElements() { return numElements; } - public static final class BytesToBytesMapIterator implements Iterator { + public final class MapIterator implements Iterator { - private final int numRecords; - private final Iterator dataPagesIterator; + private int numRecords; private final Location loc; private MemoryBlock currentPage = null; - private int currentRecordNumber = 0; + private int recordsInPage = 0; private Object pageBaseObject; private long offsetInPage; // If this iterator destructive or not. When it is true, it frees each page as it moves onto // next one. private boolean destructive = false; - private BytesToBytesMap bmap; + private LinkedList spillWriters = + new LinkedList(); + private UnsafeSorterSpillReader reader = null; - private BytesToBytesMapIterator( - int numRecords, Iterator dataPagesIterator, Location loc, - boolean destructive, BytesToBytesMap bmap) { + private MapIterator(int numRecords, Location loc, boolean destructive) { this.numRecords = numRecords; - this.dataPagesIterator = dataPagesIterator; this.loc = loc; this.destructive = destructive; - this.bmap = bmap; - if (dataPagesIterator.hasNext()) { - advanceToNextPage(); + if (destructive) { + destructiveIterator = this; } } private void advanceToNextPage() { - if (destructive && currentPage != null) { - dataPagesIterator.remove(); - this.bmap.taskMemoryManager.freePage(currentPage); + synchronized (this) { + int nextIdx = dataPages.indexOf(currentPage) + 1; + if (destructive && currentPage != null) { + dataPages.remove(currentPage); + freePage(currentPage); + nextIdx --; + } + if (dataPages.size() > nextIdx) { + currentPage = dataPages.get(nextIdx); + pageBaseObject = currentPage.getBaseObject(); + offsetInPage = currentPage.getBaseOffset(); + recordsInPage = Platform.getInt(pageBaseObject, offsetInPage); + offsetInPage += 4; + } else { + currentPage = null; + try { + reader = spillWriters.removeFirst().getReader(blockManager); + recordsInPage = -1; + } catch (IOException e) { + // Scala iterator does not handle exception + Platform.throwException(e); + } + } } - currentPage = dataPagesIterator.next(); - pageBaseObject = currentPage.getBaseObject(); - offsetInPage = currentPage.getBaseOffset(); } @Override public boolean hasNext() { - return currentRecordNumber != numRecords; + return numRecords > 0; } @Override public Location next() { - int totalLength = Platform.getInt(pageBaseObject, offsetInPage); - if (totalLength == END_OF_PAGE_MARKER) { + if (recordsInPage == 0) { advanceToNextPage(); - totalLength = Platform.getInt(pageBaseObject, offsetInPage); } - loc.with(currentPage, offsetInPage); - offsetInPage += 4 + totalLength; - currentRecordNumber++; - return loc; + numRecords--; + if (currentPage != null) { + int totalLength = Platform.getInt(pageBaseObject, offsetInPage); + loc.with(currentPage, offsetInPage); + offsetInPage += 4 + totalLength; + recordsInPage --; + return loc; + } else { + assert(reader != null); + if (!reader.hasNext()) { + advanceToNextPage(); + } + try { + reader.loadNext(); + } catch (IOException e) { + // Scala iterator does not handle exception + Platform.throwException(e); + } + loc.with(reader.getBaseObject(), reader.getBaseOffset(), reader.getRecordLength()); + return loc; + } + } + + public long spill(long numBytes) throws IOException { + synchronized (this) { + if (!destructive || dataPages.size() == 1) { + return 0L; + } + + // TODO: use existing ShuffleWriteMetrics + ShuffleWriteMetrics writeMetrics = new ShuffleWriteMetrics(); + + long released = 0L; + while (dataPages.size() > 0) { + MemoryBlock block = dataPages.getLast(); + // The currentPage is used, cannot be released + if (block == currentPage) { + break; + } + + Object base = block.getBaseObject(); + long offset = block.getBaseOffset(); + int numRecords = Platform.getInt(base, offset); + offset += 4; + final UnsafeSorterSpillWriter writer = + new UnsafeSorterSpillWriter(blockManager, 32 * 1024, writeMetrics, numRecords); + while (numRecords-- > 0) { + int length = Platform.getInt(base, offset); + writer.write(base, offset + 4, length, 0); + offset += 4 + length; + } + writer.close(); + spillWriters.add(writer); +// taskContext.addOnCompleteCallback(new AbstractFunction0() { +// @Override +// public BoxedUnit apply() { +// File file = writer.getFile(); +// if (file != null && file.exists()) { +// if (!file.delete()) { +// logger.error("Was unable to delete spill file {}", file.getAbsolutePath()); +// } +// } +// return null; +// } +// }); + + dataPages.removeLast(); + freePage(block); + released += block.size(); + + if (released > numBytes) { + break; + } + } + + return released; + } } @Override @@ -290,8 +378,8 @@ public void remove() { * If any other lookups or operations are performed on this map while iterating over it, including * `lookup()`, the behavior of the returned iterator is undefined. */ - public BytesToBytesMapIterator iterator() { - return new BytesToBytesMapIterator(numElements, dataPages.iterator(), loc, false, this); + public MapIterator iterator() { + return new MapIterator(numElements, loc, false); } /** @@ -304,8 +392,8 @@ public BytesToBytesMapIterator iterator() { * If any other lookups or operations are performed on this map while iterating over it, including * `lookup()`, the behavior of the returned iterator is undefined. */ - public BytesToBytesMapIterator destructiveIterator() { - return new BytesToBytesMapIterator(numElements, dataPages.iterator(), loc, true, this); + public MapIterator destructiveIterator() { + return new MapIterator(numElements, loc, true); } /** @@ -314,11 +402,8 @@ public BytesToBytesMapIterator destructiveIterator() { * * This function always return the same {@link Location} instance to avoid object allocation. */ - public Location lookup( - Object keyBaseObject, - long keyBaseOffset, - int keyRowLengthBytes) { - safeLookup(keyBaseObject, keyBaseOffset, keyRowLengthBytes, loc); + public Location lookup(Object keyBase, long keyOffset, int keyLength) { + safeLookup(keyBase, keyOffset, keyLength, loc); return loc; } @@ -327,18 +412,14 @@ public Location lookup( * * This is a thread-safe version of `lookup`, could be used by multiple threads. */ - public void safeLookup( - Object keyBaseObject, - long keyBaseOffset, - int keyRowLengthBytes, - Location loc) { + public void safeLookup(Object keyBase, long keyOffset, int keyLength, Location loc) { assert(bitset != null); assert(longArray != null); if (enablePerfMetrics) { numKeyLookups++; } - final int hashcode = HASHER.hashUnsafeWords(keyBaseObject, keyBaseOffset, keyRowLengthBytes); + final int hashcode = HASHER.hashUnsafeWords(keyBase, keyOffset, keyLength); int pos = hashcode & mask; int step = 1; while (true) { @@ -354,16 +435,16 @@ public void safeLookup( if ((int) (stored) == hashcode) { // Full hash code matches. Let's compare the keys for equality. loc.with(pos, hashcode, true); - if (loc.getKeyLength() == keyRowLengthBytes) { + if (loc.getKeyLength() == keyLength) { final MemoryLocation keyAddress = loc.getKeyAddress(); - final Object storedKeyBaseObject = keyAddress.getBaseObject(); - final long storedKeyBaseOffset = keyAddress.getBaseOffset(); + final Object storedkeyBase = keyAddress.getBaseObject(); + final long storedkeyOffset = keyAddress.getBaseOffset(); final boolean areEqual = ByteArrayMethods.arrayEquals( - keyBaseObject, - keyBaseOffset, - storedKeyBaseObject, - storedKeyBaseOffset, - keyRowLengthBytes + keyBase, + keyOffset, + storedkeyBase, + storedkeyOffset, + keyLength ); if (areEqual) { return; @@ -410,18 +491,18 @@ private void updateAddressesAndSizes(long fullKeyAddress) { taskMemoryManager.getOffsetInPage(fullKeyAddress)); } - private void updateAddressesAndSizes(final Object page, final long offsetInPage) { - long position = offsetInPage; - final int totalLength = Platform.getInt(page, position); + private void updateAddressesAndSizes(final Object base, final long offset) { + long position = offset; + final int totalLength = Platform.getInt(base, position); position += 4; - keyLength = Platform.getInt(page, position); + keyLength = Platform.getInt(base, position); position += 4; valueLength = totalLength - keyLength - 4; - keyMemoryLocation.setObjAndOffset(page, position); + keyMemoryLocation.setObjAndOffset(base, position); position += keyLength; - valueMemoryLocation.setObjAndOffset(page, position); + valueMemoryLocation.setObjAndOffset(base, position); } private Location with(int pos, int keyHashcode, boolean isDefined) { @@ -443,6 +524,19 @@ private Location with(MemoryBlock page, long offsetInPage) { return this; } + /** + + * This is only used for spilling + + */ + private Location with(Object base, long offset, int length) { + this.isDefined = true; + this.memoryPage = null; + keyLength = Platform.getInt(base, offset); + valueLength = length - 4 - keyLength; + keyMemoryLocation.setObjAndOffset(base, offset + 4); + valueMemoryLocation.setObjAndOffset(base, offset + 4 + keyLength); + return this; + } + /** * Returns the memory page that contains the current record. * This is only valid if this is returned by {@link BytesToBytesMap#iterator()}. @@ -517,9 +611,9 @@ public int getValueLength() { * As an example usage, here's the proper way to store a new key: *

* 
-     *   Location loc = map.lookup(keyBaseObject, keyBaseOffset, keyLengthInBytes);
+     *   Location loc = map.lookup(keyBase, keyOffset, keyLength);
      *   if (!loc.isDefined()) {
-     *     if (!loc.putNewKey(keyBaseObject, keyBaseOffset, keyLengthInBytes, ...)) {
+     *     if (!loc.putNewKey(keyBase, keyOffset, keyLength, ...)) {
      *       // handle failure to grow map (by spilling, for example)
      *     }
      *   }
@@ -531,16 +625,11 @@ public int getValueLength() {
      * @return true if the put() was successful and false if the put() failed because memory could
      *         not be acquired.
      */
-    public boolean putNewKey(
-        Object keyBaseObject,
-        long keyBaseOffset,
-        int keyLengthBytes,
-        Object valueBaseObject,
-        long valueBaseOffset,
-        int valueLengthBytes) {
+    public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
+        Object valueBase, long valueOffset, int valueLength) {
       assert (!isDefined) : "Can only set value once for a key";
-      assert (keyLengthBytes % 8 == 0);
-      assert (valueLengthBytes % 8 == 0);
+      assert (keyLength % 8 == 0);
+      assert (valueLength % 8 == 0);
       assert(bitset != null);
       assert(longArray != null);
 
@@ -552,86 +641,32 @@ public boolean putNewKey(
       // the key address instead of storing the absolute address of the value, the key and value
       // must be stored in the same memory page.
       // (8 byte key length) (key) (value)
-      final long requiredSize = 8 + keyLengthBytes + valueLengthBytes;
-
-      // --- Figure out where to insert the new record ---------------------------------------------
-
-      final MemoryBlock dataPage;
-      final Object dataPageBaseObject;
-      final long dataPageInsertOffset;
-      boolean useOverflowPage = requiredSize > pageSizeBytes - 8;
-      if (useOverflowPage) {
-        // The record is larger than the page size, so allocate a special overflow page just to hold
-        // that record.
-        final long overflowPageSize = requiredSize + 8;
-        MemoryBlock overflowPage = taskMemoryManager.allocatePage(overflowPageSize);
-        if (overflowPage == null) {
-          logger.debug("Failed to acquire {} bytes of memory", overflowPageSize);
-          return false;
-        }
-        dataPages.add(overflowPage);
-        dataPage = overflowPage;
-        dataPageBaseObject = overflowPage.getBaseObject();
-        dataPageInsertOffset = overflowPage.getBaseOffset();
-      } else if (currentDataPage == null || pageSizeBytes - 8 - pageCursor < requiredSize) {
-        // The record can fit in a data page, but either we have not allocated any pages yet or
-        // the current page does not have enough space.
-        if (currentDataPage != null) {
-          // There wasn't enough space in the current page, so write an end-of-page marker:
-          final Object pageBaseObject = currentDataPage.getBaseObject();
-          final long lengthOffsetInPage = currentDataPage.getBaseOffset() + pageCursor;
-          Platform.putInt(pageBaseObject, lengthOffsetInPage, END_OF_PAGE_MARKER);
-        }
-        if (!acquireNewPage()) {
+      final long recordLength = 8 + keyLength + valueLength;
+      if (currentPage == null || currentPage.size() - pageCursor < recordLength) {
+        if (!acquireNewPage(recordLength + 4L)) {
           return false;
         }
-        dataPage = currentDataPage;
-        dataPageBaseObject = currentDataPage.getBaseObject();
-        dataPageInsertOffset = currentDataPage.getBaseOffset();
-      } else {
-        // There is enough space in the current data page.
-        dataPage = currentDataPage;
-        dataPageBaseObject = currentDataPage.getBaseObject();
-        dataPageInsertOffset = currentDataPage.getBaseOffset() + pageCursor;
       }
 
       // --- Append the key and value data to the current data page --------------------------------
-
-      long insertCursor = dataPageInsertOffset;
-
-      // Compute all of our offsets up-front:
-      final long recordOffset = insertCursor;
-      insertCursor += 4;
-      final long keyLengthOffset = insertCursor;
-      insertCursor += 4;
-      final long keyDataOffsetInPage = insertCursor;
-      insertCursor += keyLengthBytes;
-      final long valueDataOffsetInPage = insertCursor;
-      insertCursor += valueLengthBytes; // word used to store the value size
-
-      Platform.putInt(dataPageBaseObject, recordOffset,
-        keyLengthBytes + valueLengthBytes + 4);
-      Platform.putInt(dataPageBaseObject, keyLengthOffset, keyLengthBytes);
-      // Copy the key
-      Platform.copyMemory(
-        keyBaseObject, keyBaseOffset, dataPageBaseObject, keyDataOffsetInPage, keyLengthBytes);
-      // Copy the value
-      Platform.copyMemory(valueBaseObject, valueBaseOffset, dataPageBaseObject,
-        valueDataOffsetInPage, valueLengthBytes);
-
-      // --- Update bookeeping data structures -----------------------------------------------------
-
-      if (useOverflowPage) {
-        // Store the end-of-page marker at the end of the data page
-        Platform.putInt(dataPageBaseObject, insertCursor, END_OF_PAGE_MARKER);
-      } else {
-        pageCursor += requiredSize;
-      }
-
+      final Object base = currentPage.getBaseObject();
+      long offset = currentPage.getBaseOffset() + pageCursor;
+      final long recordOffset = offset;
+      Platform.putInt(base, offset, keyLength + valueLength + 4);
+      Platform.putInt(base, offset + 4, keyLength);
+      offset += 8;
+      Platform.copyMemory(keyBase, keyOffset, base, offset, keyLength);
+      offset += keyLength;
+      Platform.copyMemory(valueBase, valueOffset, base, offset, valueLength);
+
+      // --- Update bookkeeping data structures -----------------------------------------------------
+      offset = currentPage.getBaseOffset();
+      Platform.putInt(base, offset, Platform.getInt(base, offset) + 1);
+      pageCursor += recordLength;
       numElements++;
       bitset.set(pos);
       final long storedKeyAddress = taskMemoryManager.encodePageNumberAndOffset(
-        dataPage, recordOffset);
+        currentPage, recordOffset);
       longArray.set(pos * 2, storedKeyAddress);
       longArray.set(pos * 2 + 1, keyHashcode);
       updateAddressesAndSizes(storedKeyAddress);
@@ -647,16 +682,24 @@ public boolean putNewKey(
    * Acquire a new page from the memory manager.
    * @return whether there is enough space to allocate the new page.
    */
-  private boolean acquireNewPage() {
-    MemoryBlock newPage = taskMemoryManager.allocatePage(pageSizeBytes);
-    if (newPage == null) {
-      logger.debug("Failed to acquire {} bytes of memory", pageSizeBytes);
+  private boolean acquireNewPage(long required) {
+    try {
+      currentPage = allocatePage(required);
+      dataPages.add(currentPage);
+      Platform.putInt(currentPage.getBaseObject(), currentPage.getBaseOffset(), 0);
+      pageCursor = 4;
+      return true;
+    } catch (IOException e) {
       return false;
     }
-    dataPages.add(newPage);
-    pageCursor = 0;
-    currentDataPage = newPage;
-    return true;
+  }
+
+  @Override
+  public long spill(long size) throws IOException {
+    if (destructiveIterator != null) {
+      return destructiveIterator.spill(size);
+    }
+    return 0L;
   }
 
   /**
@@ -691,7 +734,7 @@ public void free() {
     while (dataPagesIterator.hasNext()) {
       MemoryBlock dataPage = dataPagesIterator.next();
       dataPagesIterator.remove();
-      taskMemoryManager.freePage(dataPage);
+      freePage(dataPage);
     }
     assert(dataPages.isEmpty());
   }
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index e317ea391c556..e554695f4e3e5 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -17,12 +17,11 @@
 
 package org.apache.spark.util.collection.unsafe.sort;
 
+import javax.annotation.Nullable;
 import java.io.File;
 import java.io.IOException;
 import java.util.LinkedList;
 
-import javax.annotation.Nullable;
-
 import scala.runtime.AbstractFunction0;
 import scala.runtime.BoxedUnit;
 
@@ -32,24 +31,22 @@
 
 import org.apache.spark.TaskContext;
 import org.apache.spark.executor.ShuffleWriteMetrics;
+import org.apache.spark.memory.MemoryConsumer;
+import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.storage.BlockManager;
-import org.apache.spark.unsafe.array.ByteArrayMethods;
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.memory.MemoryBlock;
-import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.util.Utils;
 
 /**
  * External sorter based on {@link UnsafeInMemorySorter}.
  */
-public final class UnsafeExternalSorter {
+public final class UnsafeExternalSorter extends MemoryConsumer {
 
   private final Logger logger = LoggerFactory.getLogger(UnsafeExternalSorter.class);
 
-  private final long pageSizeBytes;
   private final PrefixComparator prefixComparator;
   private final RecordComparator recordComparator;
-  private final int initialSize;
   private final TaskMemoryManager taskMemoryManager;
   private final BlockManager blockManager;
   private final TaskContext taskContext;
@@ -70,13 +67,13 @@ public final class UnsafeExternalSorter {
 
   // These variables are reset after spilling:
   @Nullable private UnsafeInMemorySorter inMemSorter;
-  // Whether the in-mem sorter is created internally, or passed in from outside.
-  // If it is passed in from outside, we shouldn't release the in-mem sorter's memory.
-  private boolean isInMemSorterExternal = false;
+  // The acquired memory for in-memory sorter
+  private long acquiredMem = 0L;
+
   private MemoryBlock currentPage = null;
-  private long currentPagePosition = -1;
-  private long freeSpaceInCurrentPage = 0;
+  private long pageCursor = -1;
   private long peakMemoryUsedBytes = 0;
+  private SpillableIterator readingIterator = null;
 
   public static UnsafeExternalSorter createWithExistingInMemorySorter(
       TaskMemoryManager taskMemoryManager,
@@ -112,26 +109,26 @@ private UnsafeExternalSorter(
       int initialSize,
       long pageSizeBytes,
       @Nullable UnsafeInMemorySorter existingInMemorySorter) throws IOException {
+    super(taskMemoryManager, pageSizeBytes);
     this.taskMemoryManager = taskMemoryManager;
     this.blockManager = blockManager;
     this.taskContext = taskContext;
     this.recordComparator = recordComparator;
     this.prefixComparator = prefixComparator;
-    this.initialSize = initialSize;
     // Use getSizeAsKb (not bytes) to maintain backwards compatibility for units
     // this.fileBufferSizeBytes = (int) conf.getSizeAsKb("spark.shuffle.file.buffer", "32k") * 1024;
     this.fileBufferSizeBytes = 32 * 1024;
-    this.pageSizeBytes = pageSizeBytes;
+    // TODO: metrics tracking + integration with shuffle write metrics
+    // need to connect the write metrics to task metrics so we count the spill IO somewhere.
     this.writeMetrics = new ShuffleWriteMetrics();
 
     if (existingInMemorySorter == null) {
-      initializeForWriting();
-      // Acquire a new page as soon as we construct the sorter to ensure that we have at
-      // least one page to work with. Otherwise, other operators in the same task may starve
-      // this sorter (SPARK-9709). We don't need to do this if we already have an existing sorter.
-      acquireNewPage();
+      this.inMemSorter =
+        new UnsafeInMemorySorter(taskMemoryManager, recordComparator, prefixComparator, initialSize);
+      acquireMemory(inMemSorter.getMemoryUsage());
+      acquiredMem = inMemSorter.getMemoryUsage();
     } else {
-      this.isInMemSorterExternal = true;
+      acquiredMem = 0;
       this.inMemSorter = existingInMemorySorter;
     }
 
@@ -147,25 +144,7 @@ public BoxedUnit apply() {
     });
   }
 
-  // TODO: metrics tracking + integration with shuffle write metrics
-  // need to connect the write metrics to task metrics so we count the spill IO somewhere.
 
-  /**
-   * Allocates new sort data structures. Called when creating the sorter and after each spill.
-   */
-  private void initializeForWriting() throws IOException {
-    // Note: Do not track memory for the pointer array for now because of SPARK-10474.
-    // In more detail, in TungstenAggregate we only reserve a page, but when we fall back to
-    // sort-based aggregation we try to acquire a page AND a pointer array, which inevitably
-    // fails if all other memory is already occupied. It should be safe to not track the array
-    // because its memory footprint is frequently much smaller than that of a page. This is a
-    // temporary hack that we should address in 1.6.0.
-    // TODO: track the pointer array memory!
-    this.writeMetrics = new ShuffleWriteMetrics();
-    this.inMemSorter =
-      new UnsafeInMemorySorter(taskMemoryManager, recordComparator, prefixComparator, initialSize);
-    this.isInMemSorterExternal = false;
-  }
 
   /**
    * Marks the current page as no-more-space-available, and as a result, either allocate a
@@ -173,13 +152,16 @@ private void initializeForWriting() throws IOException {
    */
   @VisibleForTesting
   public void closeCurrentPage() {
-    freeSpaceInCurrentPage = 0;
+    if (currentPage != null) {
+      pageCursor = currentPage.getBaseOffset() + currentPage.size();
+    }
   }
 
   /**
    * Sort and spill the current records in response to memory pressure.
    */
-  public void spill() throws IOException {
+  @Override
+  public long spill(long size) throws IOException {
     assert(inMemSorter != null);
     logger.info("Thread {} spilling sort data of {} to disk ({} {} so far)",
       Thread.currentThread().getId(),
@@ -187,6 +169,10 @@ public void spill() throws IOException {
       spillWriters.size(),
       spillWriters.size() > 1 ? " times" : " time");
 
+    if (readingIterator != null) {
+      return readingIterator.spill();
+    }
+
     // We only write out contents of the inMemSorter if it is not empty.
     if (inMemSorter.numRecords() > 0) {
       final UnsafeSorterSpillWriter spillWriter =
@@ -202,6 +188,8 @@ public void spill() throws IOException {
         spillWriter.write(baseObject, baseOffset, recordLength, sortedRecords.getKeyPrefix());
       }
       spillWriter.close();
+
+      inMemSorter.reset();
     }
 
     final long spillSize = freeMemory();
@@ -210,7 +198,7 @@ public void spill() throws IOException {
     // written to disk. This also counts the space needed to store the sorter's pointer array.
     taskContext.taskMetrics().incMemoryBytesSpilled(spillSize);
 
-    initializeForWriting();
+    return spillSize;
   }
 
   /**
@@ -246,7 +234,7 @@ public int getNumberOfAllocatedPages() {
   }
 
   /**
-   * Free this sorter's in-memory data structures, including its data pages and pointer array.
+   * Free this sorter's data pages.
    *
    * @return the number of bytes freed.
    */
@@ -254,14 +242,12 @@ private long freeMemory() {
     updatePeakMemoryUsed();
     long memoryFreed = 0;
     for (MemoryBlock block : allocatedPages) {
-      taskMemoryManager.freePage(block);
+      freePage(block);
       memoryFreed += block.size();
     }
-    // TODO: track in-memory sorter memory usage (SPARK-10474)
     allocatedPages.clear();
     currentPage = null;
-    currentPagePosition = -1;
-    freeSpaceInCurrentPage = 0;
+    pageCursor = 0;
     return memoryFreed;
   }
 
@@ -285,6 +271,11 @@ private void deleteSpillFiles() {
   public void cleanupResources() {
     deleteSpillFiles();
     freeMemory();
+    if (inMemSorter != null) {
+      inMemSorter = null;
+      releaseMemory(acquiredMem);
+      acquiredMem = 0;
+    }
   }
 
   /**
@@ -295,8 +286,15 @@ public void cleanupResources() {
   private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
-      // TODO: track the pointer array memory! (SPARK-10474)
-      inMemSorter.expandPointerArray();
+      // assume that the memory of array will be doubled
+      long needed = inMemSorter.getMemoryUsage();
+      acquireMemory(needed);  // could trigger spilling
+      if (inMemSorter.hasSpaceForAnotherRecord()) {
+        releaseMemory(needed);
+      } else {
+        acquiredMem += needed;
+        inMemSorter.expandPointerArray();
+      }
     }
   }
 
@@ -304,101 +302,38 @@ private void growPointerArrayIfNecessary() throws IOException {
    * Allocates more memory in order to insert an additional record. This will request additional
    * memory from the memory manager and spill if the requested memory can not be obtained.
    *
-   * @param requiredSpace the required space in the data page, in bytes, including space for storing
+   * @param required the required space in the data page, in bytes, including space for storing
    *                      the record size. This must be less than or equal to the page size (records
    *                      that exceed the page size are handled via a different code path which uses
    *                      special overflow pages).
    */
-  private void acquireNewPageIfNecessary(int requiredSpace) throws IOException {
-    assert (requiredSpace <= pageSizeBytes);
-    if (requiredSpace > freeSpaceInCurrentPage) {
-      logger.trace("Required space {} is less than free space in current page ({})", requiredSpace,
-        freeSpaceInCurrentPage);
-      // TODO: we should track metrics on the amount of space wasted when we roll over to a new page
-      // without using the free space at the end of the current page. We should also do this for
-      // BytesToBytesMap.
-      if (requiredSpace > pageSizeBytes) {
-        throw new IOException("Required space " + requiredSpace + " is greater than page size (" +
-          pageSizeBytes + ")");
-      } else {
-        acquireNewPage();
-      }
+  private void acquireNewPageIfNecessary(int required) throws IOException {
+    if (currentPage == null ||
+      pageCursor + required > currentPage.getBaseOffset() + currentPage.size()) {
+      // TODO: try to find space on previous pages
+      currentPage = allocatePage(required);
+      pageCursor = currentPage.getBaseOffset();
+      allocatedPages.add(currentPage);
     }
   }
 
-  /**
-   * Acquire a new page from the memory manager.
-   *
-   * If there is not enough space to allocate the new page, spill all existing ones
-   * and try again. If there is still not enough space, report error to the caller.
-   */
-  private void acquireNewPage() throws IOException {
-    currentPage = taskMemoryManager.allocatePage(pageSizeBytes);
-    if (currentPage == null) {
-      spill();
-      currentPage = taskMemoryManager.allocatePage(pageSizeBytes);
-      if (currentPage == null) {
-        throw new IOException("Unable to acquire " + pageSizeBytes + " bytes of memory");
-      }
-    }
-    currentPagePosition = currentPage.getBaseOffset();
-    freeSpaceInCurrentPage = pageSizeBytes;
-    allocatedPages.add(currentPage);
-  }
-
   /**
    * Write a record to the sorter.
    */
-  public void insertRecord(
-      Object recordBaseObject,
-      long recordBaseOffset,
-      int lengthInBytes,
-      long prefix) throws IOException {
+  public void insertRecord(Object recordBase, long recordOffset, int length, long prefix)
+    throws IOException {
 
     growPointerArrayIfNecessary();
     // Need 4 bytes to store the record length.
-    final int totalSpaceRequired = lengthInBytes + 4;
-
-    // --- Figure out where to insert the new record ----------------------------------------------
-
-    final MemoryBlock dataPage;
-    long dataPagePosition;
-    boolean useOverflowPage = totalSpaceRequired > pageSizeBytes;
-    if (useOverflowPage) {
-      long overflowPageSize = ByteArrayMethods.roundNumberOfBytesToNearestWord(totalSpaceRequired);
-      // The record is larger than the page size, so allocate a special overflow page just to hold
-      // that record.
-      MemoryBlock overflowPage = taskMemoryManager.allocatePage(overflowPageSize);
-      if (overflowPage == null) {
-        spill();
-        overflowPage = taskMemoryManager.allocatePage(overflowPageSize);
-        if (overflowPage == null) {
-          throw new IOException("Unable to acquire " + overflowPageSize + " bytes of memory");
-        }
-      }
-      allocatedPages.add(overflowPage);
-      dataPage = overflowPage;
-      dataPagePosition = overflowPage.getBaseOffset();
-    } else {
-      // The record is small enough to fit in a regular data page, but the current page might not
-      // have enough space to hold it (or no pages have been allocated yet).
-      acquireNewPageIfNecessary(totalSpaceRequired);
-      dataPage = currentPage;
-      dataPagePosition = currentPagePosition;
-      // Update bookkeeping information
-      freeSpaceInCurrentPage -= totalSpaceRequired;
-      currentPagePosition += totalSpaceRequired;
-    }
-    final Object dataPageBaseObject = dataPage.getBaseObject();
-
-    // --- Insert the record ----------------------------------------------------------------------
-
-    final long recordAddress =
-      taskMemoryManager.encodePageNumberAndOffset(dataPage, dataPagePosition);
-    Platform.putInt(dataPageBaseObject, dataPagePosition, lengthInBytes);
-    dataPagePosition += 4;
-    Platform.copyMemory(
-      recordBaseObject, recordBaseOffset, dataPageBaseObject, dataPagePosition, lengthInBytes);
+    final int required = length + 4;
+    acquireNewPageIfNecessary(required);
+
+    final Object base = currentPage.getBaseObject();
+    final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
+    Platform.putInt(base, pageCursor, length);
+    pageCursor += 4;
+    Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
+    pageCursor += length;
     assert(inMemSorter != null);
     inMemSorter.insertRecord(recordAddress, prefix);
   }
@@ -411,59 +346,23 @@ public void insertRecord(
    *
    * record length = key length + value length + 4
    */
-  public void insertKVRecord(
-      Object keyBaseObj, long keyOffset, int keyLen,
-      Object valueBaseObj, long valueOffset, int valueLen, long prefix) throws IOException {
+  public void insertKVRecord(Object keyBase, long keyOffset, int keyLen,
+      Object valueBase, long valueOffset, int valueLen, long prefix) throws IOException {
 
     growPointerArrayIfNecessary();
-    final int totalSpaceRequired = keyLen + valueLen + 4 + 4;
-
-    // --- Figure out where to insert the new record ----------------------------------------------
-
-    final MemoryBlock dataPage;
-    long dataPagePosition;
-    boolean useOverflowPage = totalSpaceRequired > pageSizeBytes;
-    if (useOverflowPage) {
-      long overflowPageSize = ByteArrayMethods.roundNumberOfBytesToNearestWord(totalSpaceRequired);
-      // The record is larger than the page size, so allocate a special overflow page just to hold
-      // that record.
-      MemoryBlock overflowPage = taskMemoryManager.allocatePage(overflowPageSize);
-      if (overflowPage == null) {
-        spill();
-        overflowPage = taskMemoryManager.allocatePage(overflowPageSize);
-        if (overflowPage == null) {
-          throw new IOException("Unable to acquire " + overflowPageSize + " bytes of memory");
-        }
-      }
-      allocatedPages.add(overflowPage);
-      dataPage = overflowPage;
-      dataPagePosition = overflowPage.getBaseOffset();
-    } else {
-      // The record is small enough to fit in a regular data page, but the current page might not
-      // have enough space to hold it (or no pages have been allocated yet).
-      acquireNewPageIfNecessary(totalSpaceRequired);
-      dataPage = currentPage;
-      dataPagePosition = currentPagePosition;
-      // Update bookkeeping information
-      freeSpaceInCurrentPage -= totalSpaceRequired;
-      currentPagePosition += totalSpaceRequired;
-    }
-    final Object dataPageBaseObject = dataPage.getBaseObject();
-
-    // --- Insert the record ----------------------------------------------------------------------
-
-    final long recordAddress =
-      taskMemoryManager.encodePageNumberAndOffset(dataPage, dataPagePosition);
-    Platform.putInt(dataPageBaseObject, dataPagePosition, keyLen + valueLen + 4);
-    dataPagePosition += 4;
-
-    Platform.putInt(dataPageBaseObject, dataPagePosition, keyLen);
-    dataPagePosition += 4;
-
-    Platform.copyMemory(keyBaseObj, keyOffset, dataPageBaseObject, dataPagePosition, keyLen);
-    dataPagePosition += keyLen;
-
-    Platform.copyMemory(valueBaseObj, valueOffset, dataPageBaseObject, dataPagePosition, valueLen);
+    final int required = keyLen + valueLen + 4 + 4;
+    acquireNewPageIfNecessary(required);
+
+    final Object base = currentPage.getBaseObject();
+    final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
+    Platform.putInt(base, pageCursor, keyLen + valueLen + 4);
+    pageCursor += 4;
+    Platform.putInt(base, pageCursor, keyLen);
+    pageCursor += 4;
+    Platform.copyMemory(keyBase, keyOffset, base, pageCursor, keyLen);
+    pageCursor += keyLen;
+    Platform.copyMemory(valueBase, valueOffset, base, pageCursor, valueLen);
+    pageCursor += valueLen;
 
     assert(inMemSorter != null);
     inMemSorter.insertRecord(recordAddress, prefix);
@@ -475,10 +374,10 @@ public void insertKVRecord(
    */
   public UnsafeSorterIterator getSortedIterator() throws IOException {
     assert(inMemSorter != null);
-    final UnsafeInMemorySorter.SortedIterator inMemoryIterator = inMemSorter.getSortedIterator();
-    int numIteratorsToMerge = spillWriters.size() + (inMemoryIterator.hasNext() ? 1 : 0);
+    readingIterator = new SpillableIterator(inMemSorter.getSortedIterator());
+    int numIteratorsToMerge = spillWriters.size() + (readingIterator.hasNext() ? 1 : 0);
     if (spillWriters.isEmpty()) {
-      return inMemoryIterator;
+      return readingIterator;
     } else {
       final UnsafeSorterSpillMerger spillMerger =
         new UnsafeSorterSpillMerger(recordComparator, prefixComparator, numIteratorsToMerge);
@@ -486,9 +385,131 @@ public UnsafeSorterIterator getSortedIterator() throws IOException {
         spillMerger.addSpillIfNotEmpty(spillWriter.getReader(blockManager));
       }
       spillWriters.clear();
-      spillMerger.addSpillIfNotEmpty(inMemoryIterator);
+      spillMerger.addSpillIfNotEmpty(readingIterator);
 
       return spillMerger.getSortedIterator();
     }
   }
+
+  class SpillableIterator extends UnsafeSorterIterator {
+    private Object baseObject = null;
+    private long baseOffset = 0;
+    private long keyPrefix = 0;
+    private int recordLength = 0;
+
+    private boolean cached = false;
+    private MemoryBlock lastPage = null;
+
+    private UnsafeSorterIterator upstream;
+
+    public SpillableIterator(UnsafeInMemorySorter.SortedIterator inMemIterator) {
+      this.upstream = inMemIterator;
+    }
+
+    public long spill() throws IOException {
+      synchronized (this) {
+        if (!(upstream instanceof UnsafeInMemorySorter.SortedIterator && upstream.hasNext())) {
+          return 0L;
+        }
+
+        UnsafeInMemorySorter.SortedIterator inMemIterator =
+          (UnsafeInMemorySorter.SortedIterator) upstream;
+
+        // backup the current record
+        baseObject = inMemIterator.getBaseObject();
+        baseOffset = inMemIterator.getBaseOffset();
+        keyPrefix = inMemIterator.getKeyPrefix();
+        recordLength = inMemIterator.getRecordLength();
+
+        final UnsafeSorterSpillWriter spillWriter =
+          new UnsafeSorterSpillWriter(blockManager, fileBufferSizeBytes, writeMetrics,
+            inMemIterator.numRecordsLeft());
+        while (inMemIterator.hasNext()) {
+          inMemIterator.loadNext();
+          final Object baseObject = inMemIterator.getBaseObject();
+          final long baseOffset = inMemIterator.getBaseOffset();
+          final int recordLength = inMemIterator.getRecordLength();
+          spillWriter.write(baseObject, baseOffset, recordLength, inMemIterator.getKeyPrefix());
+        }
+        spillWriter.close();
+        upstream = spillWriter.getReader(blockManager);
+
+        cached = true;
+        // release the pages except the one that is used
+        long released = 0L;
+        for (MemoryBlock block : allocatedPages) {
+          if (block.getBaseObject() != baseObject) {
+            freePage(block);
+            released += block.size();
+          } else {
+            lastPage = block;
+          }
+        }
+        allocatedPages.clear();
+        return released;
+      }
+    }
+
+    @Override
+    public boolean hasNext() {
+      synchronized (this) {
+        return upstream.hasNext();
+      }
+    }
+
+    @Override
+    public void loadNext() throws IOException {
+      synchronized (this) {
+        if (cached) {
+          // Just consumed the last record from in memory iterator
+          if (lastPage != null) {
+            freePage(lastPage);
+            lastPage = null;
+          }
+          cached = false;
+        }
+        upstream.loadNext();
+      }
+    }
+
+    @Override
+    public Object getBaseObject() {
+      synchronized (this) {
+        if (cached) {
+          return baseObject;
+        }
+        return upstream.getBaseObject();
+      }
+    }
+
+    @Override
+    public long getBaseOffset() {
+      synchronized (this) {
+        if (cached) {
+          return baseOffset;
+        }
+        return upstream.getBaseOffset();
+      }
+    }
+
+    @Override
+    public int getRecordLength() {
+      synchronized (this) {
+        if (cached) {
+          return recordLength;
+        }
+        return upstream.getRecordLength();
+      }
+    }
+
+    @Override
+    public long getKeyPrefix() {
+      synchronized (this) {
+        if (cached) {
+          return keyPrefix;
+        }
+        return upstream.getKeyPrefix();
+      }
+    }
+  }
 }
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
index 5aad72c374c37..45596fba8b2fe 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
@@ -89,6 +89,10 @@ public UnsafeInMemorySorter(
     this.sortComparator = new SortComparator(recordComparator, prefixComparator, memoryManager);
   }
 
+  public void reset() {
+    pointerArrayInsertPosition = 0;
+  }
+
   /**
    * @return the number of records that have been inserted into this sorter.
    */
@@ -100,12 +104,8 @@ public long getMemoryUsage() {
     return pointerArray.length * 8L;
   }
 
-  static long getMemoryRequirementsForPointerArray(long numEntries) {
-    return numEntries * 2L * 8L;
-  }
-
   public boolean hasSpaceForAnotherRecord() {
-    return pointerArrayInsertPosition + 2 < pointerArray.length;
+    return pointerArrayInsertPosition + 2 <= pointerArray.length;
   }
 
   public void expandPointerArray() {
@@ -158,6 +158,10 @@ public boolean hasNext() {
       return position < sortBufferInsertPosition;
     }
 
+    public int numRecordsLeft() {
+      return sortBufferInsertPosition - position;
+    }
+
     @Override
     public void loadNext() {
       // This pointer points to a 4-byte record length, followed by the record's bytes
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillReader.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillReader.java
index 501dfe77d13cb..039e940a357ea 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillReader.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillReader.java
@@ -20,18 +20,18 @@
 import java.io.*;
 
 import com.google.common.io.ByteStreams;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import org.apache.spark.storage.BlockId;
 import org.apache.spark.storage.BlockManager;
 import org.apache.spark.unsafe.Platform;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /**
  * Reads spill files written by {@link UnsafeSorterSpillWriter} (see that class for a description
  * of the file format).
  */
-final class UnsafeSorterSpillReader extends UnsafeSorterIterator {
+public final class UnsafeSorterSpillReader extends UnsafeSorterIterator {
   private static final Logger logger = LoggerFactory.getLogger(UnsafeSorterSpillReader.class);
 
   private final File file;
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillWriter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillWriter.java
index e59a84ff8d118..234e21140a1dd 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillWriter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillWriter.java
@@ -35,7 +35,7 @@
  *
  *   [# of records (int)] [[len (int)][prefix (long)][data (bytes)]...]
  */
-final class UnsafeSorterSpillWriter {
+public final class UnsafeSorterSpillWriter {
 
   static final int DISK_WRITE_BUFFER_SIZE = 1024 * 1024;
 
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java
index 7fb2f92ca80e8..4ba4d7ab9641b 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java
@@ -17,20 +17,24 @@
 
 package org.apache.spark.shuffle.sort;
 
-import org.apache.spark.shuffle.sort.PackedRecordPointer;
+import java.io.IOException;
+
 import org.junit.Test;
-import static org.junit.Assert.*;
 
 import org.apache.spark.SparkConf;
 import org.apache.spark.memory.GrantEverythingMemoryManager;
-import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.memory.TaskMemoryManager;
-import static org.apache.spark.shuffle.sort.PackedRecordPointer.*;
+import org.apache.spark.unsafe.memory.MemoryBlock;
+
+import static org.apache.spark.shuffle.sort.PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES;
+import static org.apache.spark.shuffle.sort.PackedRecordPointer.MAXIMUM_PARTITION_ID;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 
 public class PackedRecordPointerSuite {
 
   @Test
-  public void heap() {
+  public void heap() throws IOException {
     final SparkConf conf = new SparkConf().set("spark.unsafe.offHeap", "false");
     final TaskMemoryManager memoryManager =
       new TaskMemoryManager(new GrantEverythingMemoryManager(conf), 0);
@@ -49,7 +53,7 @@ public void heap() {
   }
 
   @Test
-  public void offHeap() {
+  public void offHeap() throws IOException {
     final SparkConf conf = new SparkConf().set("spark.unsafe.offHeap", "true");
     final TaskMemoryManager memoryManager =
       new TaskMemoryManager(new GrantEverythingMemoryManager(conf), 0);
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
index d65926949c036..00f3c1eea7545 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
@@ -49,6 +49,7 @@
 import org.apache.spark.io.SnappyCompressionCodec;
 import org.apache.spark.executor.ShuffleWriteMetrics;
 import org.apache.spark.executor.TaskMetrics;
+import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.network.util.LimitedInputStream;
 import org.apache.spark.serializer.*;
 import org.apache.spark.scheduler.MapStatus;
@@ -344,9 +345,7 @@ private void testMergingSpills(
     }
     assertEquals(sumOfPartitionSizes, mergedOutputFile.length());
 
-    assertEquals(
-      HashMultiset.create(dataToWrite),
-      HashMultiset.create(readRecordsFromFile()));
+    assertEquals(HashMultiset.create(dataToWrite), HashMultiset.create(readRecordsFromFile()));
     assertSpillFilesWereCleanedUp();
     ShuffleWriteMetrics shuffleWriteMetrics = taskMetrics.shuffleWriteMetrics().get();
     assertEquals(dataToWrite.size(), shuffleWriteMetrics.shuffleRecordsWritten());
@@ -411,7 +410,7 @@ public void writeEnoughDataToTriggerSpill() throws Exception {
       dataToWrite.add(new Tuple2(i, bigByteArray));
     }
     writer.write(dataToWrite.iterator());
-    verify(taskMemoryManager, times(5)).acquireExecutionMemory(anyLong());
+    verify(taskMemoryManager, times(3)).acquireExecutionMemory(anyLong(), any(MemoryConsumer.class));
     assertEquals(2, spillFilesCreated.size());
     writer.stop(true);
     readRecordsFromFile();
@@ -438,7 +437,7 @@ public void writeEnoughRecordsToTriggerSortBufferExpansionAndSpill() throws Exce
       dataToWrite.add(new Tuple2(i, i));
     }
     writer.write(dataToWrite.iterator());
-    verify(taskMemoryManager, times(5)).acquireExecutionMemory(anyLong());
+    verify(taskMemoryManager, times(2)).acquireExecutionMemory(anyLong(), any(MemoryConsumer.class));
     assertEquals(2, spillFilesCreated.size());
     writer.stop(true);
     readRecordsFromFile();
@@ -524,7 +523,7 @@ public void testPeakMemoryUsed() throws Exception {
       for (int i = 0; i < numRecordsPerPage * 10; i++) {
         writer.insertRecordIntoSorter(new Tuple2(1, 1));
         newPeakMemory = writer.getPeakMemoryUsedBytes();
-        if (i % numRecordsPerPage == 0 && i != 0) {
+        if (i % numRecordsPerPage == 0) {
           // The first page is allocated in constructor, another page will be allocated after
           // every numRecordsPerPage records (peak memory should change).
           assertEquals(previousPeakMemory + pageSizeBytes, newPeakMemory);
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index 6e52496cf933b..3c70bdce666f2 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -17,20 +17,45 @@
 
 package org.apache.spark.unsafe.map;
 
-import java.lang.Exception;
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.util.*;
 
-import org.apache.spark.memory.TaskMemoryManager;
-import org.junit.*;
-import static org.hamcrest.Matchers.greaterThan;
-import static org.junit.Assert.*;
+import scala.Tuple2;
+import scala.Tuple2$;
+import scala.runtime.AbstractFunction1;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 
 import org.apache.spark.SparkConf;
+import org.apache.spark.executor.ShuffleWriteMetrics;
 import org.apache.spark.memory.GrantEverythingMemoryManager;
-import org.apache.spark.unsafe.array.ByteArrayMethods;
-import org.apache.spark.unsafe.memory.*;
+import org.apache.spark.memory.TaskMemoryManager;
+import org.apache.spark.serializer.SerializerInstance;
+import org.apache.spark.storage.*;
 import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.array.ByteArrayMethods;
+import org.apache.spark.unsafe.memory.MemoryLocation;
+import org.apache.spark.util.Utils;
+
+import static org.hamcrest.Matchers.greaterThan;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.mockito.AdditionalAnswers.returnsSecondArg;
+import static org.mockito.Answers.RETURNS_SMART_NULLS;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Mockito.when;
 
 
 public abstract class AbstractBytesToBytesMapSuite {
@@ -41,16 +66,68 @@ public abstract class AbstractBytesToBytesMapSuite {
   private TaskMemoryManager taskMemoryManager;
   private final long PAGE_SIZE_BYTES = 1L << 26; // 64 megabytes
 
+  final LinkedList spillFilesCreated = new LinkedList();
+  File tempDir;
+
+  @Mock(answer = RETURNS_SMART_NULLS) BlockManager blockManager;
+  @Mock(answer = RETURNS_SMART_NULLS) DiskBlockManager diskBlockManager;
+
+  private static final class CompressStream extends AbstractFunction1 {
+    @Override
+    public OutputStream apply(OutputStream stream) {
+      return stream;
+    }
+  }
+
   @Before
   public void setup() {
     memoryManager =
       new GrantEverythingMemoryManager(
         new SparkConf().set("spark.unsafe.offHeap", "" + useOffHeapMemoryAllocator()));
     taskMemoryManager = new TaskMemoryManager(memoryManager, 0);
+
+    tempDir = Utils.createTempDir(System.getProperty("java.io.tmpdir"), "unsafe-test");
+    spillFilesCreated.clear();
+    MockitoAnnotations.initMocks(this);
+    when(blockManager.diskBlockManager()).thenReturn(diskBlockManager);
+    when(diskBlockManager.createTempLocalBlock()).thenAnswer(new Answer>() {
+      @Override
+      public Tuple2 answer(InvocationOnMock invocationOnMock) throws Throwable {
+        TempLocalBlockId blockId = new TempLocalBlockId(UUID.randomUUID());
+        File file = File.createTempFile("spillFile", ".spill", tempDir);
+        spillFilesCreated.add(file);
+        return Tuple2$.MODULE$.apply(blockId, file);
+      }
+    });
+    when(blockManager.getDiskWriter(
+      any(BlockId.class),
+      any(File.class),
+      any(SerializerInstance.class),
+      anyInt(),
+      any(ShuffleWriteMetrics.class))).thenAnswer(new Answer() {
+      @Override
+      public DiskBlockObjectWriter answer(InvocationOnMock invocationOnMock) throws Throwable {
+        Object[] args = invocationOnMock.getArguments();
+
+        return new DiskBlockObjectWriter(
+          (File) args[1],
+          (SerializerInstance) args[2],
+          (Integer) args[3],
+          new CompressStream(),
+          false,
+          (ShuffleWriteMetrics) args[4]
+        );
+      }
+    });
+    when(blockManager.wrapForCompression(any(BlockId.class), any(InputStream.class)))
+      .then(returnsSecondArg());
   }
 
   @After
   public void tearDown() {
+    Utils.deleteRecursively(tempDir);
+    tempDir = null;
+
     Assert.assertEquals(0L, taskMemoryManager.cleanUpAllAllocatedMemory());
     if (taskMemoryManager != null) {
       long leakedMemory = taskMemoryManager.getMemoryConsumptionForThisTask();
@@ -415,9 +492,8 @@ public void randomizedTestWithRecordsLargerThanPageSize() {
 
   @Test
   public void failureToAllocateFirstPage() {
-    memoryManager.markExecutionAsOutOfMemory();
+    memoryManager.markExecutionAsOutOfMemoryAlways();
     BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, 1, PAGE_SIZE_BYTES);
-    memoryManager.markExecutionAsOutOfMemory();
     try {
       final long[] emptyArray = new long[0];
       final BytesToBytesMap.Location loc =
@@ -427,6 +503,7 @@ public void failureToAllocateFirstPage() {
         emptyArray, Platform.LONG_ARRAY_OFFSET, 0, emptyArray, Platform.LONG_ARRAY_OFFSET, 0));
     } finally {
       map.free();
+      memoryManager.resetOutOfMemory();
     }
   }
 
@@ -439,7 +516,7 @@ public void failureToGrow() {
       int i;
       for (i = 0; i < 127; i++) {
         if (i > 0) {
-          memoryManager.markExecutionAsOutOfMemory();
+          memoryManager.markExecutionAsOutOfMemoryAlways();
         }
         final long[] arr = new long[]{i};
         final BytesToBytesMap.Location loc = map.lookup(arr, Platform.LONG_ARRAY_OFFSET, 8);
@@ -453,6 +530,44 @@ public void failureToGrow() {
       Assert.assertFalse(success);
     } finally {
       map.free();
+      memoryManager.resetOutOfMemory();
+    }
+  }
+
+  @Test
+  public void spillInIterator() throws IOException {
+    BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, blockManager, 1, 0.75, 1024, false);
+    try {
+      int i;
+      for (i = 0; i < 1024; i++) {
+        final long[] arr = new long[]{i};
+        final BytesToBytesMap.Location loc = map.lookup(arr, Platform.LONG_ARRAY_OFFSET, 8);
+        loc.putNewKey(arr, Platform.LONG_ARRAY_OFFSET, 8, arr, Platform.LONG_ARRAY_OFFSET, 8);
+      }
+      BytesToBytesMap.MapIterator iter = map.iterator();
+      for (i = 0; i < 100; i++) {
+        iter.next();
+      }
+      // Non-destructive iterator is not spillable
+      Assert.assertEquals(0, iter.spill(1024L * 10));
+      for (i = 100; i < 1024; i++) {
+        iter.next();
+      }
+      BytesToBytesMap.MapIterator iter2 = map.destructiveIterator();
+      for (i = 0; i < 100; i++) {
+        iter2.next();
+      }
+      Assert.assertTrue(iter2.spill(1024) > 1024);
+      for (i = 100; i < 1024; i++) {
+        iter2.next();
+      }
+      assertFalse(iter2.hasNext());
+    } finally {
+      map.free();
+      for (File spillFile : spillFilesCreated) {
+        assertFalse("Spill file " + spillFile.getPath() + " was not cleaned up",
+          spillFile.exists());
+      }
     }
   }
 
@@ -500,7 +615,7 @@ public void testPeakMemoryUsed() {
           Platform.LONG_ARRAY_OFFSET,
           8);
         newPeakMemory = map.getPeakMemoryUsedBytes();
-        if (i % numRecordsPerPage == 0 && i > 0) {
+        if (i % numRecordsPerPage == 0) {
           // We allocated a new page for this record, so peak memory should change
           assertEquals(previousPeakMemory + pageSizeBytes, newPeakMemory);
         } else {
@@ -519,11 +634,4 @@ public void testPeakMemoryUsed() {
     }
   }
 
-  @Test
-  public void testAcquirePageInConstructor() {
-    final BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, 1, PAGE_SIZE_BYTES);
-    assertEquals(1, map.getNumDataPages());
-    map.free();
-  }
-
 }
diff --git a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
index 94d50b94fde3f..3c7d0791ac09a 100644
--- a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
+++ b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
@@ -36,23 +36,24 @@
 import org.mockito.MockitoAnnotations;
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
-import static org.hamcrest.Matchers.greaterThanOrEqualTo;
-import static org.junit.Assert.*;
-import static org.mockito.AdditionalAnswers.returnsSecondArg;
-import static org.mockito.Answers.RETURNS_SMART_NULLS;
-import static org.mockito.Mockito.*;
 
 import org.apache.spark.SparkConf;
 import org.apache.spark.TaskContext;
 import org.apache.spark.executor.ShuffleWriteMetrics;
 import org.apache.spark.executor.TaskMetrics;
 import org.apache.spark.memory.GrantEverythingMemoryManager;
+import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.serializer.SerializerInstance;
 import org.apache.spark.storage.*;
 import org.apache.spark.unsafe.Platform;
-import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.util.Utils;
 
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.junit.Assert.*;
+import static org.mockito.AdditionalAnswers.returnsSecondArg;
+import static org.mockito.Answers.RETURNS_SMART_NULLS;
+import static org.mockito.Mockito.*;
+
 public class UnsafeExternalSorterSuite {
 
   final LinkedList spillFilesCreated = new LinkedList();
@@ -86,7 +87,7 @@ public int compare(
   @Mock(answer = RETURNS_SMART_NULLS) TaskContext taskContext;
 
 
-  private final long pageSizeBytes = new SparkConf().getSizeAsBytes("spark.buffer.pageSize", "64m");
+  private final long pageSizeBytes = new SparkConf().getSizeAsBytes("spark.buffer.pageSize", "4m");
 
   private static final class CompressStream extends AbstractFunction1 {
     @Override
@@ -183,9 +184,9 @@ public void testSortingOnlyByPrefix() throws Exception {
     insertNumber(sorter, 5);
     insertNumber(sorter, 1);
     insertNumber(sorter, 3);
-    sorter.spill();
+    sorter.spill(0);
     insertNumber(sorter, 4);
-    sorter.spill();
+    sorter.spill(0);
     insertNumber(sorter, 2);
 
     UnsafeSorterIterator iter = sorter.getSortedIterator();
@@ -206,9 +207,9 @@ public void testSortingEmptyArrays() throws Exception {
     final UnsafeExternalSorter sorter = newSorter();
     sorter.insertRecord(null, 0, 0, 0);
     sorter.insertRecord(null, 0, 0, 0);
-    sorter.spill();
+    sorter.spill(0L);
     sorter.insertRecord(null, 0, 0, 0);
-    sorter.spill();
+    sorter.spill(0L);
     sorter.insertRecord(null, 0, 0, 0);
     sorter.insertRecord(null, 0, 0, 0);
 
@@ -233,7 +234,7 @@ public void spillingOccursInResponseToMemoryPressure() throws Exception {
       insertNumber(sorter, numRecords - i);
     }
     assertEquals(1, sorter.getNumberOfAllocatedPages());
-    memoryManager.markExecutionAsOutOfMemory();
+    memoryManager.markExecutionAsOutOfMemoryOnce();
     // The insertion of this record should trigger a spill:
     insertNumber(sorter, 0);
     // Ensure that spill files were created
@@ -274,9 +275,9 @@ public void sortingRecordsThatExceedPageSize() throws Exception {
     Arrays.fill(smallRecord, 123);
 
     insertRecord(sorter, largeRecord, 456);
-    sorter.spill();
+    sorter.spill(0L);
     insertRecord(sorter, smallRecord, 123);
-    sorter.spill();
+    sorter.spill(0L);
     insertRecord(sorter, smallRecord, 123);
     insertRecord(sorter, largeRecord, 456);
 
@@ -334,7 +335,7 @@ public void testPeakMemoryUsed() throws Exception {
         insertNumber(sorter, i);
         newPeakMemory = sorter.getPeakMemoryUsedBytes();
         // The first page is pre-allocated on instantiation
-        if (i % numRecordsPerPage == 0 && i > 0) {
+        if (i % numRecordsPerPage == 0) {
           // We allocated a new page for this record, so peak memory should change
           assertEquals(previousPeakMemory + pageSizeBytes, newPeakMemory);
         } else {
@@ -344,7 +345,7 @@ public void testPeakMemoryUsed() throws Exception {
       }
 
       // Spilling should not change peak memory
-      sorter.spill();
+      sorter.spill(0L);
       newPeakMemory = sorter.getPeakMemoryUsedBytes();
       assertEquals(previousPeakMemory, newPeakMemory);
       for (int i = 0; i < numRecordsPerPage; i++) {
@@ -358,21 +359,5 @@ public void testPeakMemoryUsed() throws Exception {
     }
   }
 
-  @Test
-  public void testReservePageOnInstantiation() throws Exception {
-    final UnsafeExternalSorter sorter = newSorter();
-    try {
-      assertEquals(1, sorter.getNumberOfAllocatedPages());
-      // Inserting a new record doesn't allocate more memory since we already have a page
-      long peakMemory = sorter.getPeakMemoryUsedBytes();
-      insertNumber(sorter, 100);
-      assertEquals(peakMemory, sorter.getPeakMemoryUsedBytes());
-      assertEquals(1, sorter.getNumberOfAllocatedPages());
-    } finally {
-      sorter.cleanupResources();
-      assertSpillFilesWereCleanedUp();
-    }
-  }
-
 }
 
diff --git a/core/src/test/scala/org/apache/spark/memory/GrantEverythingMemoryManager.scala b/core/src/test/scala/org/apache/spark/memory/GrantEverythingMemoryManager.scala
index fe102d8aeb2a5..cde7b48b461dc 100644
--- a/core/src/test/scala/org/apache/spark/memory/GrantEverythingMemoryManager.scala
+++ b/core/src/test/scala/org/apache/spark/memory/GrantEverythingMemoryManager.scala
@@ -26,8 +26,10 @@ class GrantEverythingMemoryManager(conf: SparkConf) extends MemoryManager(conf,
   private[memory] override def doAcquireExecutionMemory(
       numBytes: Long,
       evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Long = synchronized {
-    if (oom) {
-      oom = false
+    if (oomAlways) {
+      0
+    } else if (oomOnce) {
+      oomOnce = false
       0
     } else {
       _executionMemoryUsed += numBytes // To suppress warnings when freeing unallocated memory
@@ -46,9 +48,18 @@ class GrantEverythingMemoryManager(conf: SparkConf) extends MemoryManager(conf,
   override def maxExecutionMemory: Long = Long.MaxValue
   override def maxStorageMemory: Long = Long.MaxValue
 
-  private var oom = false
+  private var oomOnce = false
+  private var oomAlways = false
+
+  def markExecutionAsOutOfMemoryOnce(): Unit = {
+    oomOnce = true
+  }
+
+  def markExecutionAsOutOfMemoryAlways(): Unit = {
+    oomAlways = true
+  }
 
-  def markExecutionAsOutOfMemory(): Unit = {
-    oom = true
+  def resetOutOfMemory(): Unit = {
+    oomAlways = false
   }
 }
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java b/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
index 810c74fd2fb96..a56d31c2ef369 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
@@ -102,7 +102,7 @@ void insertRow(UnsafeRow row) throws IOException {
 
   @VisibleForTesting
   void spill() throws IOException {
-    sorter.spill();
+    sorter.spill(Long.MAX_VALUE);
   }
 
   /**
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java
index 82c645df284de..889f97003450c 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java
@@ -165,7 +165,7 @@ public UnsafeRow getAggregationBufferFromUnsafeRow(UnsafeRow unsafeGroupingKeyRo
   public KVIterator iterator() {
     return new KVIterator() {
 
-      private final BytesToBytesMap.BytesToBytesMapIterator mapLocationIterator =
+      private final BytesToBytesMap.MapIterator mapLocationIterator =
         map.destructiveIterator();
       private final UnsafeRow key = new UnsafeRow();
       private final UnsafeRow value = new UnsafeRow();
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
index 46301f0042954..bd5f91f36726b 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
@@ -17,13 +17,13 @@
 
 package org.apache.spark.sql.execution;
 
-import java.io.IOException;
-
 import javax.annotation.Nullable;
+import java.io.IOException;
 
 import com.google.common.annotations.VisibleForTesting;
 
 import org.apache.spark.TaskContext;
+import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
 import org.apache.spark.sql.catalyst.expressions.codegen.BaseOrdering;
 import org.apache.spark.sql.catalyst.expressions.codegen.GenerateOrdering;
@@ -33,7 +33,6 @@
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.map.BytesToBytesMap;
 import org.apache.spark.unsafe.memory.MemoryBlock;
-import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.util.collection.unsafe.sort.*;
 
 /**
@@ -95,7 +94,7 @@ public UnsafeKVExternalSorter(
       // We cannot use the destructive iterator here because we are reusing the existing memory
       // pages in BytesToBytesMap to hold records during sorting.
       // The only new memory we are allocating is the pointer/prefix array.
-      BytesToBytesMap.BytesToBytesMapIterator iter = map.iterator();
+      BytesToBytesMap.MapIterator iter = map.iterator();
       final int numKeyFields = keySchema.size();
       UnsafeRow row = new UnsafeRow();
       while (iter.hasNext()) {
@@ -126,7 +125,7 @@ public UnsafeKVExternalSorter(
         pageSizeBytes,
         inMemSorter);
 
-      sorter.spill();
+      sorter.spill(Long.MAX_VALUE);
       map.free();
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
index dbf4863b767bf..23e66adcaffc5 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
@@ -208,7 +208,7 @@ class UnsafeFixedWidthAggregationMapSuite
       sorter.insertKV(keyConverter.apply(k), valueConverter.apply(v))
 
       if ((i % 100) == 0) {
-        memoryManager.markExecutionAsOutOfMemory()
+        memoryManager.markExecutionAsOutOfMemoryOnce()
         sorter.closeCurrentPage()
       }
     }
@@ -251,7 +251,7 @@ class UnsafeFixedWidthAggregationMapSuite
       sorter.insertKV(keyConverter.apply(k), valueConverter.apply(v))
 
       if ((i % 100) == 0) {
-        memoryManager.markExecutionAsOutOfMemory()
+        memoryManager.markExecutionAsOutOfMemoryOnce()
         sorter.closeCurrentPage()
       }
     }
@@ -303,7 +303,7 @@ class UnsafeFixedWidthAggregationMapSuite
       sorter.insertKV(UnsafeRow.createFromByteArray(0, 0), UnsafeRow.createFromByteArray(0, 0))
 
       if ((i % 100) == 0) {
-        memoryManager.markExecutionAsOutOfMemory()
+        memoryManager.markExecutionAsOutOfMemoryOnce()
         sorter.closeCurrentPage()
       }
     }
@@ -342,8 +342,8 @@ class UnsafeFixedWidthAggregationMapSuite
       buf.setInt(0, str.length)
     }
     // Simulate running out of space
-    memoryManager.markExecutionAsOutOfMemory()
-    val str = rand.nextString(1024)
+    memoryManager.markExecutionAsOutOfMemoryAlways()
+    val str = rand.nextString(4096)
     val buf = map.getAggregationBuffer(InternalRow(UTF8String.fromString(str)))
     assert(buf == null)
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeKVExternalSorterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeKVExternalSorterSuite.scala
index 13dc1754c9ff0..261861fbf81a5 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeKVExternalSorterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeKVExternalSorterSuite.scala
@@ -128,7 +128,7 @@ class UnsafeKVExternalSorterSuite extends SparkFunSuite with SharedSQLContext {
       sorter.insertKV(k.asInstanceOf[UnsafeRow], v.asInstanceOf[UnsafeRow])
       // 1% chance we will spill
       if (rand.nextDouble() < 0.01 && spill) {
-        memoryManager.markExecutionAsOutOfMemory()
+        memoryManager.markExecutionAsOutOfMemoryOnce()
         sorter.closeCurrentPage()
       }
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIteratorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIteratorSuite.scala
deleted file mode 100644
index 475037bd45379..0000000000000
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIteratorSuite.scala
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.execution.aggregate
-
-import org.apache.spark._
-import org.apache.spark.memory.TaskMemoryManager
-import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.InterpretedMutableProjection
-import org.apache.spark.sql.execution.metric.SQLMetrics
-import org.apache.spark.sql.test.SharedSQLContext
-
-class TungstenAggregationIteratorSuite extends SparkFunSuite with SharedSQLContext {
-
-  test("memory acquired on construction") {
-    val taskMemoryManager = new TaskMemoryManager(SparkEnv.get.memoryManager, 0)
-    val taskContext = new TaskContextImpl(0, 0, 0, 0, taskMemoryManager, null, Seq.empty)
-    TaskContext.setTaskContext(taskContext)
-
-    // Assert that a page is allocated before processing starts
-    var iter: TungstenAggregationIterator = null
-    try {
-      val newMutableProjection = (expr: Seq[Expression], schema: Seq[Attribute]) => {
-        () => new InterpretedMutableProjection(expr, schema)
-      }
-      val dummyAccum = SQLMetrics.createLongMetric(sparkContext, "dummy")
-      iter = new TungstenAggregationIterator(Seq.empty, Seq.empty, Seq.empty, Seq.empty, Seq.empty,
-        0, Seq.empty, newMutableProjection, Seq.empty, None,
-        dummyAccum, dummyAccum, dummyAccum, dummyAccum)
-      val numPages = iter.getHashMap.getNumDataPages
-      assert(numPages === 1)
-    } finally {
-      // Clean up
-      if (iter != null) {
-        iter.free()
-      }
-      TaskContext.unset()
-    }
-  }
-}
diff --git a/unsafe/src/main/java/org/apache/spark/unsafe/memory/HeapMemoryAllocator.java b/unsafe/src/main/java/org/apache/spark/unsafe/memory/HeapMemoryAllocator.java
index ebe90d9e63d83..09847cec9c4ca 100644
--- a/unsafe/src/main/java/org/apache/spark/unsafe/memory/HeapMemoryAllocator.java
+++ b/unsafe/src/main/java/org/apache/spark/unsafe/memory/HeapMemoryAllocator.java
@@ -23,6 +23,8 @@
 import java.util.LinkedList;
 import java.util.Map;
 
+import org.apache.spark.unsafe.Platform;
+
 /**
  * A simple {@link MemoryAllocator} that can allocate up to 16GB using a JVM long primitive array.
  */
@@ -45,9 +47,6 @@ private boolean shouldPool(long size) {
 
   @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> pool = bufferPoolsBySize.get(size);
@@ -64,8 +63,8 @@ public MemoryBlock allocate(long size) throws OutOfMemoryError {
         }
       }
     }
-    long[] array = new long[(int) (size / 8)];
-    return MemoryBlock.fromLongArray(array);
+    long[] array = new long[(int) ((size + 7) / 8)];
+    return new MemoryBlock(array, Platform.LONG_ARRAY_OFFSET, size);
   }
 
   @Override
diff --git a/unsafe/src/main/java/org/apache/spark/unsafe/memory/UnsafeMemoryAllocator.java b/unsafe/src/main/java/org/apache/spark/unsafe/memory/UnsafeMemoryAllocator.java
index cda7826c8c99b..98ce711176e43 100644
--- a/unsafe/src/main/java/org/apache/spark/unsafe/memory/UnsafeMemoryAllocator.java
+++ b/unsafe/src/main/java/org/apache/spark/unsafe/memory/UnsafeMemoryAllocator.java
@@ -26,9 +26,6 @@ public class UnsafeMemoryAllocator implements MemoryAllocator {
 
   @Override
   public MemoryBlock allocate(long size) throws OutOfMemoryError {
-    if (size % 8 != 0) {
-      throw new IllegalArgumentException("Size " + size + " was not a multiple of 8");
-    }
     long address = Platform.allocateMemory(size);
     return new MemoryBlock(null, address, size);
   }

From ee6b9a40d01b0ae9d43e34511c43c77f8a080bf0 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Tue, 27 Oct 2015 10:35:39 -0700
Subject: [PATCH 02/17] fix build

---
 .../spark/memory/TaskMemoryManager.java       | 87 ++++++++-----------
 .../spark/util/collection/Spillable.scala     |  4 +-
 .../spark/memory/TaskMemoryManagerSuite.java  | 14 +--
 .../sort/PackedRecordPointerSuite.java        |  8 +-
 .../sort/ShuffleInMemorySorterSuite.java      |  2 +-
 .../sort/UnsafeShuffleWriterSuite.java        |  4 +-
 .../sort/UnsafeInMemorySorterSuite.java       | 16 ++--
 .../scala/org/apache/spark/FailureSuite.scala |  4 +-
 .../spark/memory/MemoryManagerSuite.scala     | 60 ++++++-------
 9 files changed, 93 insertions(+), 106 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 4f7e98ace60d1..2ec75c9424f0c 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -25,7 +25,10 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.spark.SparkException;
+import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.memory.MemoryBlock;
+import org.apache.spark.util.Utils;
 
 /**
  * Manages the memory allocated by an individual task.
@@ -117,14 +120,6 @@ public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) {
     this.consumers = new HashMap<>();
   }
 
-  /**
-   * Acquire N bytes of memory for execution, evicting cached blocks if necessary.
-   * @return number of bytes successfully granted (<= N).
-   */
-  public long acquireExecutionMemory(long size) {
-    return memoryManager.acquireExecutionMemory(size, taskAttemptId);
-  }
-
   /**
    * Acquire N bytes of memory for a consumer. If there is no enough memory, it will call
    * spill() of consumers to release more memory.
@@ -133,26 +128,26 @@ public long acquireExecutionMemory(long size) {
    */
   public long acquireExecutionMemory(long size, MemoryConsumer consumer) throws IOException {
     synchronized (this) {
-      long got = acquireExecutionMemory(size);
+      long got = memoryManager.acquireExecutionMemory(size, taskAttemptId);
 
+      // call spill() on itself to release some memory
       if (got < size && consumer != null) {
-        // call spill() on itself to release some memory
         consumer.spill(size - got);
-        got += acquireExecutionMemory(size - got);
-
-        if (got < size) {
-          long needed = size - got;
-          // call spill() on other consumers to release memory
-          for (MemoryConsumer c: consumers.keySet()) {
-            if (c != consumer) {
-              needed -= c.spill(size - got);
-              if (needed < 0) {
-                break;
-              }
+        got += memoryManager.acquireExecutionMemory(size - got, taskAttemptId);
+      }
+
+      if (got < size) {
+        long needed = size - got;
+        // call spill() on other consumers to release memory
+        for (MemoryConsumer c: consumers.keySet()) {
+          if (c != null && c != consumer) {
+            needed -= c.spill(size - got);
+            if (needed < 0) {
+              break;
             }
           }
-          got += acquireExecutionMemory(size - got);
         }
+        got += memoryManager.acquireExecutionMemory(size - got, taskAttemptId);
       }
 
       long old = 0L;
@@ -165,30 +160,35 @@ public long acquireExecutionMemory(long size, MemoryConsumer consumer) throws IO
     }
   }
 
-  /**
-   * Release N bytes of execution memory.
-   */
-  public void releaseExecutionMemory(long size) {
-    memoryManager.releaseExecutionMemory(size, taskAttemptId);
-  }
-
   /**
    * Release N bytes of execution memory for a MemoryConsumer.
    */
   public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
     synchronized (this) {
-      if (consumer != null && consumers.containsKey(consumer)) {
+      if (consumers.containsKey(consumer)) {
         long old = consumers.get(consumer);
         if (old > size) {
           consumers.put(consumer, old - size);
         } else {
           if (old < size) {
-            // TODO
+            if (Utils.isTesting()) {
+              Platform.throwException(
+                new SparkException("Release more memory " + size + "than acquired " + old + " for "
+                  + consumer));
+            } else {
+              logger.warn("Release more memory " + size + " than acquired " + old + "for "
+                + consumer);
+            }
           }
           consumers.remove(consumer);
         }
       } else {
-        // TODO
+        if (Utils.isTesting()) {
+          Platform.throwException(
+            new SparkException("Release memory " + size + " for not existed " + consumer));
+        } else {
+          logger.warn("Release memory " + size + " for not existed " + consumer);
+        }
       }
       memoryManager.releaseExecutionMemory(size, taskAttemptId);
     }
@@ -198,16 +198,6 @@ public long pageSizeBytes() {
     return memoryManager.pageSizeBytes();
   }
 
-  /**
-   * Allocate a block of memory that will be tracked in the MemoryManager's page table; this is
-   * intended for allocating large blocks of Tungsten memory that will be shared between operators.
-   *
-   * Returns `null` if there was not enough memory to allocate the page.
-   */
-  public MemoryBlock allocatePage(long size) throws IOException {
-    return allocatePage(size, null);
-  }
-
   /**
    * Allocate a block of memory that will be tracked in the MemoryManager's page table; this is
    * intended for allocating large blocks of Tungsten memory that will be shared between operators.
@@ -244,14 +234,7 @@ public MemoryBlock allocatePage(long size, MemoryConsumer consumer) throws IOExc
   }
 
   /**
-   * Free a block of memory allocated via {@link TaskMemoryManager#allocatePage(long)}.
-   */
-  public void freePage(MemoryBlock page) {
-    freePage(page, null);
-  }
-
-  /**
-   * Free a block of memory allocated via {@link TaskMemoryManager#allocatePage(long)}.
+   * Free a block of memory allocated via {@link TaskMemoryManager#allocatePage}.
    */
   public void freePage(MemoryBlock page, MemoryConsumer consumer) {
     assert (page.pageNumber != -1) :
@@ -273,7 +256,7 @@ public void freePage(MemoryBlock page, MemoryConsumer consumer) {
    * Given a memory page and offset within that page, encode this address into a 64-bit long.
    * This address will remain valid as long as the corresponding page has not been freed.
    *
-   * @param page a data page allocated by {@link TaskMemoryManager#allocatePage(long)}/
+   * @param page a data page allocated by {@link TaskMemoryManager#allocatePage}/
    * @param offsetInPage an offset in this page which incorporates the base offset. In other words,
    *                     this should be the value that you would pass as the base offset into an
    *                     UNSAFE call (e.g. page.baseOffset() + something).
@@ -349,7 +332,7 @@ public long cleanUpAllAllocatedMemory() {
     for (MemoryBlock page : pageTable) {
       if (page != null) {
         freedBytes += page.size();
-        freePage(page);
+        freePage(page, null);
       }
     }
 
diff --git a/core/src/main/scala/org/apache/spark/util/collection/Spillable.scala b/core/src/main/scala/org/apache/spark/util/collection/Spillable.scala
index a76891acf0baf..9e002621a6909 100644
--- a/core/src/main/scala/org/apache/spark/util/collection/Spillable.scala
+++ b/core/src/main/scala/org/apache/spark/util/collection/Spillable.scala
@@ -78,7 +78,7 @@ private[spark] trait Spillable[C] extends Logging {
     if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
       // Claim up to double our current memory from the shuffle memory pool
       val amountToRequest = 2 * currentMemory - myMemoryThreshold
-      val granted = taskMemoryManager.acquireExecutionMemory(amountToRequest)
+      val granted = taskMemoryManager.acquireExecutionMemory(amountToRequest, null)
       myMemoryThreshold += granted
       // If we were granted too little memory to grow further (either tryToAcquire returned 0,
       // or we already had more memory than myMemoryThreshold), spill the current collection
@@ -107,7 +107,7 @@ private[spark] trait Spillable[C] extends Logging {
    */
   def releaseMemory(): Unit = {
     // The amount we requested does not include the initial memory tracking threshold
-    taskMemoryManager.releaseExecutionMemory(myMemoryThreshold - initialMemoryThreshold)
+    taskMemoryManager.releaseExecutionMemory(myMemoryThreshold - initialMemoryThreshold, null)
     myMemoryThreshold = initialMemoryThreshold
   }
 
diff --git a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
index f381db0c62653..ba6039994fa31 100644
--- a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
+++ b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
@@ -17,6 +17,8 @@
 
 package org.apache.spark.memory;
 
+import java.io.IOException;
+
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -26,18 +28,18 @@
 public class TaskMemoryManagerSuite {
 
   @Test
-  public void leakedPageMemoryIsDetected() {
+  public void leakedPageMemoryIsDetected() throws IOException {
     final TaskMemoryManager manager = new TaskMemoryManager(
       new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
-    manager.allocatePage(4096);  // leak memory
+    manager.allocatePage(4096, null);  // leak memory
     Assert.assertEquals(4096, manager.cleanUpAllAllocatedMemory());
   }
 
   @Test
-  public void encodePageNumberAndOffsetOffHeap() {
+  public void encodePageNumberAndOffsetOffHeap() throws IOException {
     final TaskMemoryManager manager = new TaskMemoryManager(
       new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "true")), 0);
-    final MemoryBlock dataPage = manager.allocatePage(256);
+    final MemoryBlock dataPage = manager.allocatePage(256, null);
     // In off-heap mode, an offset is an absolute address that may require more than 51 bits to
     // encode. This test exercises that corner-case:
     final long offset = ((1L << TaskMemoryManager.OFFSET_BITS) + 10);
@@ -47,10 +49,10 @@ public void encodePageNumberAndOffsetOffHeap() {
   }
 
   @Test
-  public void encodePageNumberAndOffsetOnHeap() {
+  public void encodePageNumberAndOffsetOnHeap() throws IOException {
     final TaskMemoryManager manager = new TaskMemoryManager(
       new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
-    final MemoryBlock dataPage = manager.allocatePage(256);
+    final MemoryBlock dataPage = manager.allocatePage(256, null);
     final long encodedAddress = manager.encodePageNumberAndOffset(dataPage, 64);
     Assert.assertEquals(dataPage.getBaseObject(), manager.getPage(encodedAddress));
     Assert.assertEquals(64, manager.getOffsetInPage(encodedAddress));
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java
index 4ba4d7ab9641b..b5c781001687b 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java
@@ -38,8 +38,8 @@ public void heap() throws IOException {
     final SparkConf conf = new SparkConf().set("spark.unsafe.offHeap", "false");
     final TaskMemoryManager memoryManager =
       new TaskMemoryManager(new GrantEverythingMemoryManager(conf), 0);
-    final MemoryBlock page0 = memoryManager.allocatePage(128);
-    final MemoryBlock page1 = memoryManager.allocatePage(128);
+    final MemoryBlock page0 = memoryManager.allocatePage(128, null);
+    final MemoryBlock page1 = memoryManager.allocatePage(128, null);
     final long addressInPage1 = memoryManager.encodePageNumberAndOffset(page1,
       page1.getBaseOffset() + 42);
     PackedRecordPointer packedPointer = new PackedRecordPointer();
@@ -57,8 +57,8 @@ public void offHeap() throws IOException {
     final SparkConf conf = new SparkConf().set("spark.unsafe.offHeap", "true");
     final TaskMemoryManager memoryManager =
       new TaskMemoryManager(new GrantEverythingMemoryManager(conf), 0);
-    final MemoryBlock page0 = memoryManager.allocatePage(128);
-    final MemoryBlock page1 = memoryManager.allocatePage(128);
+    final MemoryBlock page0 = memoryManager.allocatePage(128, null);
+    final MemoryBlock page1 = memoryManager.allocatePage(128, null);
     final long addressInPage1 = memoryManager.encodePageNumberAndOffset(page1,
       page1.getBaseOffset() + 42);
     PackedRecordPointer packedPointer = new PackedRecordPointer();
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java
index 5049a5306ff21..719646ec2fbde 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java
@@ -61,7 +61,7 @@ public void testBasicSorting() throws Exception {
     final SparkConf conf = new SparkConf().set("spark.unsafe.offHeap", "false");
     final TaskMemoryManager memoryManager =
       new TaskMemoryManager(new GrantEverythingMemoryManager(conf), 0);
-    final MemoryBlock dataPage = memoryManager.allocatePage(2048);
+    final MemoryBlock dataPage = memoryManager.allocatePage(2048, null);
     final Object baseObject = dataPage.getBaseObject();
     final ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(4);
     final HashPartitioner hashPartitioner = new HashPartitioner(4);
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
index 00f3c1eea7545..a824e5b900830 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
@@ -402,7 +402,7 @@ public void writeEnoughDataToTriggerSpill() throws Exception {
       .doCallRealMethod() // allocate initial data page
       .doReturn(0L) // deny request to allocate new page
       .doCallRealMethod() // grant new sort buffer and data page
-      .when(taskMemoryManager).acquireExecutionMemory(anyLong());
+      .when(taskMemoryManager).acquireExecutionMemory(anyLong(), null);
     final UnsafeShuffleWriter writer = createWriter(false);
     final ArrayList> dataToWrite = new ArrayList>();
     final byte[] bigByteArray = new byte[PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES / 128];
@@ -430,7 +430,7 @@ public void writeEnoughRecordsToTriggerSortBufferExpansionAndSpill() throws Exce
       .doCallRealMethod() // allocate initial data page
       .doReturn(0L) // deny request to allocate new page
       .doCallRealMethod() // grant new sort buffer and data page
-      .when(taskMemoryManager).acquireExecutionMemory(anyLong());
+      .when(taskMemoryManager).acquireExecutionMemory(anyLong(), null);
     final UnsafeShuffleWriter writer = createWriter(false);
     final ArrayList> dataToWrite = new ArrayList<>();
     for (int i = 0; i < UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE; i++) {
diff --git a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java
index d5de56a0512f9..6a2649e1a19b8 100644
--- a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java
+++ b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java
@@ -20,17 +20,19 @@
 import java.util.Arrays;
 
 import org.junit.Test;
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.*;
-import static org.junit.Assert.*;
-import static org.mockito.Mockito.mock;
 
 import org.apache.spark.HashPartitioner;
 import org.apache.spark.SparkConf;
-import org.apache.spark.unsafe.Platform;
 import org.apache.spark.memory.GrantEverythingMemoryManager;
-import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.memory.TaskMemoryManager;
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.memory.MemoryBlock;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.isIn;
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
 
 public class UnsafeInMemorySorterSuite {
 
@@ -67,7 +69,7 @@ public void testSortingOnlyByIntegerPrefix() throws Exception {
     };
     final TaskMemoryManager memoryManager = new TaskMemoryManager(
       new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
-    final MemoryBlock dataPage = memoryManager.allocatePage(2048);
+    final MemoryBlock dataPage = memoryManager.allocatePage(2048, null);
     final Object baseObject = dataPage.getBaseObject();
     // Write the records into the data page:
     long position = dataPage.getBaseOffset();
diff --git a/core/src/test/scala/org/apache/spark/FailureSuite.scala b/core/src/test/scala/org/apache/spark/FailureSuite.scala
index 0242cbc9244a8..203dab934ca1f 100644
--- a/core/src/test/scala/org/apache/spark/FailureSuite.scala
+++ b/core/src/test/scala/org/apache/spark/FailureSuite.scala
@@ -149,7 +149,7 @@ class FailureSuite extends SparkFunSuite with LocalSparkContext {
     // cause is preserved
     val thrownDueToTaskFailure = intercept[SparkException] {
       sc.parallelize(Seq(0)).mapPartitions { iter =>
-        TaskContext.get().taskMemoryManager().allocatePage(128)
+        TaskContext.get().taskMemoryManager().allocatePage(128, null)
         throw new Exception("intentional task failure")
         iter
       }.count()
@@ -159,7 +159,7 @@ class FailureSuite extends SparkFunSuite with LocalSparkContext {
     // If the task succeeded but memory was leaked, then the task should fail due to that leak
     val thrownDueToMemoryLeak = intercept[SparkException] {
       sc.parallelize(Seq(0)).mapPartitions { iter =>
-        TaskContext.get().taskMemoryManager().allocatePage(128)
+        TaskContext.get().taskMemoryManager().allocatePage(128, null)
         iter
       }.count()
     }
diff --git a/core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala b/core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala
index 1265087743a98..4a9479cf490fb 100644
--- a/core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala
@@ -145,20 +145,20 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
     val manager = createMemoryManager(1000L)
     val taskMemoryManager = new TaskMemoryManager(manager, 0)
 
-    assert(taskMemoryManager.acquireExecutionMemory(100L) === 100L)
-    assert(taskMemoryManager.acquireExecutionMemory(400L) === 400L)
-    assert(taskMemoryManager.acquireExecutionMemory(400L) === 400L)
-    assert(taskMemoryManager.acquireExecutionMemory(200L) === 100L)
-    assert(taskMemoryManager.acquireExecutionMemory(100L) === 0L)
-    assert(taskMemoryManager.acquireExecutionMemory(100L) === 0L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L, null) === 100L)
+    assert(taskMemoryManager.acquireExecutionMemory(400L, null) === 400L)
+    assert(taskMemoryManager.acquireExecutionMemory(400L, null) === 400L)
+    assert(taskMemoryManager.acquireExecutionMemory(200L, null) === 100L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L, null) === 0L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L, null) === 0L)
 
-    taskMemoryManager.releaseExecutionMemory(500L)
-    assert(taskMemoryManager.acquireExecutionMemory(300L) === 300L)
-    assert(taskMemoryManager.acquireExecutionMemory(300L) === 200L)
+    taskMemoryManager.releaseExecutionMemory(500L, null)
+    assert(taskMemoryManager.acquireExecutionMemory(300L, null) === 300L)
+    assert(taskMemoryManager.acquireExecutionMemory(300L, null) === 200L)
 
     taskMemoryManager.cleanUpAllAllocatedMemory()
-    assert(taskMemoryManager.acquireExecutionMemory(1000L) === 1000L)
-    assert(taskMemoryManager.acquireExecutionMemory(100L) === 0L)
+    assert(taskMemoryManager.acquireExecutionMemory(1000L, null) === 1000L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L, null) === 0L)
   }
 
   test("two tasks requesting full execution memory") {
@@ -168,15 +168,15 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
     val futureTimeout: Duration = 20.seconds
 
     // Have both tasks request 500 bytes, then wait until both requests have been granted:
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(500L) }
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(500L, null) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
     assert(Await.result(t1Result1, futureTimeout) === 500L)
     assert(Await.result(t2Result1, futureTimeout) === 500L)
 
     // Have both tasks each request 500 bytes more; both should immediately return 0 as they are
     // both now at 1 / N
-    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L) }
-    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L, null) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
     assert(Await.result(t1Result2, 200.millis) === 0L)
     assert(Await.result(t2Result2, 200.millis) === 0L)
   }
@@ -188,15 +188,15 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
     val futureTimeout: Duration = 20.seconds
 
     // Have both tasks request 250 bytes, then wait until both requests have been granted:
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(250L) }
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(250L, null) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L, null) }
     assert(Await.result(t1Result1, futureTimeout) === 250L)
     assert(Await.result(t2Result1, futureTimeout) === 250L)
 
     // Have both tasks each request 500 bytes more.
     // We should only grant 250 bytes to each of them on this second request
-    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L) }
-    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L, null) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
     assert(Await.result(t1Result2, futureTimeout) === 250L)
     assert(Await.result(t2Result2, futureTimeout) === 250L)
   }
@@ -208,17 +208,17 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
     val futureTimeout: Duration = 20.seconds
 
     // t1 grabs 1000 bytes and then waits until t2 is ready to make a request.
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L, null) }
     assert(Await.result(t1Result1, futureTimeout) === 1000L)
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L, null) }
     // Make sure that t2 didn't grab the memory right away. This is hacky but it would be difficult
     // to make sure the other thread blocks for some time otherwise.
     Thread.sleep(300)
-    t1MemManager.releaseExecutionMemory(250L)
+    t1MemManager.releaseExecutionMemory(250L, null)
     // The memory freed from t1 should now be granted to t2.
     assert(Await.result(t2Result1, futureTimeout) === 250L)
     // Further requests by t2 should be denied immediately because it now has 1 / 2N of the memory.
-    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(100L) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(100L, null) }
     assert(Await.result(t2Result2, 200.millis) === 0L)
   }
 
@@ -229,18 +229,18 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
     val futureTimeout: Duration = 20.seconds
 
     // t1 grabs 1000 bytes and then waits until t2 is ready to make a request.
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L, null) }
     assert(Await.result(t1Result1, futureTimeout) === 1000L)
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
     // Make sure that t2 didn't grab the memory right away. This is hacky but it would be difficult
     // to make sure the other thread blocks for some time otherwise.
     Thread.sleep(300)
     // t1 releases all of its memory, so t2 should be able to grab all of the memory
     t1MemManager.cleanUpAllAllocatedMemory()
     assert(Await.result(t2Result1, futureTimeout) === 500L)
-    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
     assert(Await.result(t2Result2, futureTimeout) === 500L)
-    val t2Result3 = Future { t2MemManager.acquireExecutionMemory(500L) }
+    val t2Result3 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
     assert(Await.result(t2Result3, 200.millis) === 0L)
   }
 
@@ -251,13 +251,13 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
     val t2MemManager = new TaskMemoryManager(memoryManager, 2)
     val futureTimeout: Duration = 20.seconds
 
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(700L) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(700L, null) }
     assert(Await.result(t1Result1, futureTimeout) === 700L)
 
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(300L) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(300L, null) }
     assert(Await.result(t2Result1, futureTimeout) === 300L)
 
-    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(300L) }
+    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(300L, null) }
     assert(Await.result(t1Result2, 200.millis) === 0L)
   }
 }

From 49b8135bf37347fb057f363398c17367f0f0c991 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Tue, 27 Oct 2015 10:48:16 -0700
Subject: [PATCH 03/17] fix Spillable

---
 .../main/java/org/apache/spark/memory/TaskMemoryManager.java   | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 2ec75c9424f0c..ac2e4e887b06c 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -164,6 +164,9 @@ public long acquireExecutionMemory(long size, MemoryConsumer consumer) throws IO
    * Release N bytes of execution memory for a MemoryConsumer.
    */
   public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
+    if (size == 0) {
+      return;
+    }
     synchronized (this) {
       if (consumers.containsKey(consumer)) {
         long old = consumers.get(consumer);

From 7087f2f4ca9841716b9b4740c037b410607cf719 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Tue, 27 Oct 2015 16:19:32 -0700
Subject: [PATCH 04/17] fix tests

---
 .../apache/spark/memory/MemoryConsumer.java   |  12 +-
 .../spark/memory/TaskMemoryManager.java       |   8 +-
 .../shuffle/sort/ShuffleExternalSorter.java   |  41 +++---
 .../shuffle/sort/ShuffleInMemorySorter.java   |  49 ++++---
 .../shuffle/sort/UnsafeShuffleWriter.java     |   8 +-
 .../spark/unsafe/map/BytesToBytesMap.java     |   6 +-
 .../unsafe/sort/UnsafeExternalSorter.java     | 133 ++++++++----------
 .../unsafe/sort/UnsafeInMemorySorter.java     |  42 +++---
 .../apache/spark/memory/MemoryManager.scala   |   2 +-
 .../spark/memory/TaskMemoryManagerSuite.java  |   6 +-
 .../sort/PackedRecordPointerSuite.java        |   6 +-
 .../sort/ShuffleInMemorySorterSuite.java      |   4 +-
 .../sort/UnsafeShuffleWriterSuite.java        |  35 ++---
 .../map/AbstractBytesToBytesMapSuite.java     |  12 +-
 .../sort/UnsafeExternalSorterSuite.java       |  20 +--
 .../sort/UnsafeInMemorySorterSuite.java       |   6 +-
 ...yManager.scala => TestMemoryManager.scala} |  28 ++--
 .../execution/UnsafeExternalRowSorter.java    |   7 +-
 .../sql/execution/UnsafeKVExternalSorter.java |   2 +-
 .../UnsafeFixedWidthAggregationMapSuite.scala |   6 +-
 .../UnsafeKVExternalSorterSuite.scala         |   4 +-
 21 files changed, 216 insertions(+), 221 deletions(-)
 rename core/src/test/scala/org/apache/spark/memory/{GrantEverythingMemoryManager.scala => TestMemoryManager.scala} (78%)

diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
index e308c4b74a579..6d7421cf8b86a 100644
--- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
+++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
@@ -40,15 +40,25 @@ protected MemoryConsumer(TaskMemoryManager memoryManager) {
     this(memoryManager, memoryManager.pageSizeBytes());
   }
 
+  /**
+   * Force spill during during.
+   *
+   * For testing.
+   */
+  public long spill() throws IOException {
+    return spill(Long.MAX_VALUE, this);
+  }
+
   /**
    * Spill some data to disk to release memory, which will be called by TaskMemoryManager
    * when there is not enough memory for the task.
    *
    * @param size the amount of memory should be released
+   * @param trigger the MemoryConsumer that trigger this spilling
    * @return the amount of released memory in bytes
    * @throws IOException
    */
-  public long spill(long size) throws IOException {
+  public long spill(long size, MemoryConsumer trigger) throws IOException {
     return 0L;
   }
 
diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index ac2e4e887b06c..0303e93d77123 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -132,7 +132,7 @@ public long acquireExecutionMemory(long size, MemoryConsumer consumer) throws IO
 
       // call spill() on itself to release some memory
       if (got < size && consumer != null) {
-        consumer.spill(size - got);
+        consumer.spill(size - got, consumer);
         got += memoryManager.acquireExecutionMemory(size - got, taskAttemptId);
       }
 
@@ -141,7 +141,7 @@ public long acquireExecutionMemory(long size, MemoryConsumer consumer) throws IO
         // call spill() on other consumers to release memory
         for (MemoryConsumer c: consumers.keySet()) {
           if (c != null && c != consumer) {
-            needed -= c.spill(size - got);
+            needed -= c.spill(size - got, consumer);
             if (needed < 0) {
               break;
             }
@@ -227,11 +227,11 @@ public MemoryBlock allocatePage(long size, MemoryConsumer consumer) throws IOExc
       }
       allocatedPages.set(pageNumber);
     }
-    final MemoryBlock page = memoryManager.tungstenMemoryAllocator().allocate(size);
+    final MemoryBlock page = memoryManager.tungstenMemoryAllocator().allocate(acquired);
     page.pageNumber = pageNumber;
     pageTable[pageNumber] = page;
     if (logger.isTraceEnabled()) {
-      logger.trace("Allocate page number {} ({} bytes)", pageNumber, size);
+      logger.trace("Allocate page number {} ({} bytes)", pageNumber, acquired);
     }
     return page;
   }
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
index 69ac697a03511..d5696979036d8 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
@@ -66,14 +66,10 @@ final class ShuffleExternalSorter extends MemoryConsumer {
   static final int DISK_WRITE_BUFFER_SIZE = 1024 * 1024;
 
   private final int numPartitions;
-  private final int pageSizeBytes;
-  @VisibleForTesting
-  final int maxRecordSizeBytes;
   private final TaskMemoryManager taskMemoryManager;
   private final BlockManager blockManager;
   private final TaskContext taskContext;
   private final ShuffleWriteMetrics writeMetrics;
-  private long numRecordsInsertedSinceLastSpill = 0;
 
   /** Force this sorter to spill when there are this many elements in memory. For testing only */
   private final long numElementsForSpillThreshold;
@@ -117,13 +113,9 @@ public ShuffleExternalSorter(
     this.fileBufferSizeBytes = (int) conf.getSizeAsKb("spark.shuffle.file.buffer", "32k") * 1024;
     this.numElementsForSpillThreshold =
       conf.getLong("spark.shuffle.spill.numElementsForceSpillThreshold", Long.MAX_VALUE);
-    this.pageSizeBytes = (int) Math.min(
-      PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES, taskMemoryManager.pageSizeBytes());
-    this.maxRecordSizeBytes = pageSizeBytes - 4;
     this.writeMetrics = writeMetrics;
     acquireMemory(initialSize * 8L);
     this.inMemSorter = new ShuffleInMemorySorter(initialSize);
-    numRecordsInsertedSinceLastSpill = 0;
   }
 
   /**
@@ -248,9 +240,12 @@ private void writeSortedFile(boolean isLastFile) throws IOException {
   /**
    * Sort and spill the current records in response to memory pressure.
    */
-  @VisibleForTesting
   @Override
-  public long spill(long size) throws IOException {
+  public long spill(long size, MemoryConsumer trigger) throws IOException {
+    if (trigger != this || inMemSorter == null || inMemSorter.numRecords() == 0) {
+      return 0L;
+    }
+
     logger.info("Thread {} spilling sort data of {} to disk ({} {} so far)",
       Thread.currentThread().getId(),
       Utils.bytesToString(getMemoryUsage()),
@@ -290,8 +285,8 @@ private long freeMemory() {
     updatePeakMemoryUsed();
     long memoryFreed = 0;
     for (MemoryBlock block : allocatedPages) {
-      freePage(block);
       memoryFreed += block.size();
+      freePage(block);
     }
     allocatedPages.clear();
     currentPage = null;
@@ -325,9 +320,15 @@ private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
       logger.debug("Attempting to expand sort pointer array");
-      final long oldPointerArrayMemoryUsage = inMemSorter.getMemoryUsage();
-      acquireMemory(oldPointerArrayMemoryUsage);  //TODO actual memory may be less
-      inMemSorter.expandPointerArray();
+      long used = inMemSorter.getMemoryUsage();
+      long needed = inMemSorter.getMemoryToExpand();
+      acquireMemory(used + needed);  // could trigger spilling
+      if (inMemSorter.hasSpaceForAnotherRecord()) {
+        releaseMemory(used + needed);
+      } else {
+        inMemSorter.expandPointerArray();
+        releaseMemory(used);
+      }
     }
   }
 
@@ -341,11 +342,10 @@ private void growPointerArrayIfNecessary() throws IOException {
    *                      special overflow pages).
    */
   private void acquireNewPageIfNecessary(int required) throws IOException {
-    growPointerArrayIfNecessary();
     if (currentPage == null ||
       pageCursor + required > currentPage.getBaseOffset() + currentPage.size() ) {
       // TODO: try to find space in previous pages
-      currentPage = allocatePage(pageSizeBytes);
+      currentPage = allocatePage(required);
       pageCursor = currentPage.getBaseOffset();
       allocatedPages.add(currentPage);
     }
@@ -358,22 +358,23 @@ public void insertRecord(Object recordBase, long recordOffset, int length, int p
     throws IOException {
 
     // for tests
-    if (numRecordsInsertedSinceLastSpill > numElementsForSpillThreshold) {
-      spill(0);
+    assert(inMemSorter != null);
+    if (inMemSorter.numRecords() > numElementsForSpillThreshold) {
+      spill();
     }
 
+    growPointerArrayIfNecessary();
     final int required = length + 4;
     acquireNewPageIfNecessary(required);
 
+    assert(currentPage != null);
     final Object base = currentPage.getBaseObject();
     final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
     Platform.putInt(base, pageCursor, length);
     pageCursor += 4;
     Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
     pageCursor += length;
-    assert(inMemSorter != null);
     inMemSorter.insertRecord(recordAddress, partitionId);
-    numRecordsInsertedSinceLastSpill += 1;
   }
 
   /**
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
index ca18e510e8c62..2b49737fed413 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
@@ -37,37 +37,52 @@ public int compare(PackedRecordPointer left, PackedRecordPointer right) {
    * {@link PackedRecordPointer}. The sort operates on this array instead of directly manipulating
    * records.
    */
-  private long[] pointerArray;
+  private long[] array;
 
   /**
    * The position in the pointer array where new records can be inserted.
    */
-  private int pointerArrayInsertPosition = 0;
+  private int pos = 0;
 
   public ShuffleInMemorySorter(int initialSize) {
     assert (initialSize > 0);
-    this.pointerArray = new long[initialSize];
-    this.sorter = new Sorter(ShuffleSortDataFormat.INSTANCE);
+    this.array = new long[initialSize];
+    this.sorter = new Sorter<>(ShuffleSortDataFormat.INSTANCE);
+  }
+
+  public int numRecords() {
+    return pos;
   }
 
   public void reset() {
-    pointerArrayInsertPosition = 0;
+    pos = 0;
   }
 
-  public void expandPointerArray() {
-    final long[] oldArray = pointerArray;
+  private int newLength() {
     // Guard against overflow:
-    final int newLength = oldArray.length * 2 > 0 ? (oldArray.length * 2) : Integer.MAX_VALUE;
-    pointerArray = new long[newLength];
-    System.arraycopy(oldArray, 0, pointerArray, 0, oldArray.length);
+    return array.length <= Integer.MAX_VALUE / 2 ?
+      (array.length * 2) : Integer.MAX_VALUE;
+  }
+
+  /**
+   * Returns the memory needed to expand
+   */
+  public long getMemoryToExpand() {
+    return ((long) (newLength() - array.length)) * 8;
+  }
+
+  public void expandPointerArray() {
+    final long[] oldArray = array;
+    array = new long[newLength()];
+    System.arraycopy(oldArray, 0, array, 0, oldArray.length);
   }
 
   public boolean hasSpaceForAnotherRecord() {
-    return pointerArrayInsertPosition < pointerArray.length;
+    return pos < array.length;
   }
 
   public long getMemoryUsage() {
-    return pointerArray.length * 8L;
+    return array.length * 8L;
   }
 
   /**
@@ -82,15 +97,15 @@ public long getMemoryUsage() {
    */
   public void insertRecord(long recordPointer, int partitionId) {
     if (!hasSpaceForAnotherRecord()) {
-      if (pointerArray.length == Integer.MAX_VALUE) {
+      if (array.length == Integer.MAX_VALUE) {
         throw new IllegalStateException("Sort pointer array has reached maximum size");
       } else {
         expandPointerArray();
       }
     }
-    pointerArray[pointerArrayInsertPosition] =
+    array[pos] =
         PackedRecordPointer.packPointer(recordPointer, partitionId);
-    pointerArrayInsertPosition++;
+    pos++;
   }
 
   /**
@@ -122,7 +137,7 @@ public void loadNext() {
    * Return an iterator over record pointers in sorted order.
    */
   public ShuffleSorterIterator getSortedIterator() {
-    sorter.sort(pointerArray, 0, pointerArrayInsertPosition, SORT_COMPARATOR);
-    return new ShuffleSorterIterator(pointerArrayInsertPosition, pointerArray);
+    sorter.sort(array, 0, pos, SORT_COMPARATOR);
+    return new ShuffleSorterIterator(pos, array);
   }
 }
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
index fcb3f3064794c..e19b37864293c 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
@@ -127,12 +127,6 @@ public UnsafeShuffleWriter(
     open();
   }
 
-  @VisibleForTesting
-  public int maxRecordSizeBytes() {
-    assert(sorter != null);
-    return sorter.maxRecordSizeBytes;
-  }
-
   private void updatePeakMemoryUsed() {
     // sorter can be null if this writer is closed
     if (sorter != null) {
@@ -245,7 +239,7 @@ void insertRecordIntoSorter(Product2 record) throws IOException {
   @VisibleForTesting
   void forceSorterToSpill() throws IOException {
     assert (sorter != null);
-    sorter.spill(Long.MAX_VALUE);
+    sorter.spill();
   }
 
   /**
diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 42c6ac34dc22c..04144f599c967 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -166,7 +166,7 @@ public final class BytesToBytesMap extends MemoryConsumer {
   private long peakMemoryUsedBytes = 0L;
 
   private final BlockManager blockManager;
-  private MapIterator destructiveIterator = null;
+  private volatile MapIterator destructiveIterator = null;
 
   public BytesToBytesMap(
       TaskMemoryManager taskMemoryManager,
@@ -695,8 +695,8 @@ private boolean acquireNewPage(long required) {
   }
 
   @Override
-  public long spill(long size) throws IOException {
-    if (destructiveIterator != null) {
+  public long spill(long size, MemoryConsumer trigger) throws IOException {
+    if (trigger != this && destructiveIterator != null) {
       return destructiveIterator.spill(size);
     }
     return 0L;
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index e554695f4e3e5..cabe29fcc1675 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -66,14 +66,14 @@ public final class UnsafeExternalSorter extends MemoryConsumer {
   private final LinkedList spillWriters = new LinkedList<>();
 
   // These variables are reset after spilling:
-  @Nullable private UnsafeInMemorySorter inMemSorter;
+  @Nullable private volatile UnsafeInMemorySorter inMemSorter;
   // The acquired memory for in-memory sorter
   private long acquiredMem = 0L;
 
   private MemoryBlock currentPage = null;
   private long pageCursor = -1;
   private long peakMemoryUsedBytes = 0;
-  private SpillableIterator readingIterator = null;
+  private volatile SpillableIterator readingIterator = null;
 
   public static UnsafeExternalSorter createWithExistingInMemorySorter(
       TaskMemoryManager taskMemoryManager,
@@ -144,8 +144,6 @@ public BoxedUnit apply() {
     });
   }
 
-
-
   /**
    * Marks the current page as no-more-space-available, and as a result, either allocate a
    * new page or spill when we see the next record.
@@ -161,18 +159,24 @@ public void closeCurrentPage() {
    * Sort and spill the current records in response to memory pressure.
    */
   @Override
-  public long spill(long size) throws IOException {
-    assert(inMemSorter != null);
+  public long spill(long size, MemoryConsumer trigger) throws IOException {
+    if (trigger != this) {
+      if (readingIterator != null) {
+        return readingIterator.spill();
+      }
+      return 0L;
+    }
+
+    if (inMemSorter == null || inMemSorter.numRecords() <= 0) {
+      return 0L;
+    }
+
     logger.info("Thread {} spilling sort data of {} to disk ({} {} so far)",
       Thread.currentThread().getId(),
       Utils.bytesToString(getMemoryUsage()),
       spillWriters.size(),
       spillWriters.size() > 1 ? " times" : " time");
 
-    if (readingIterator != null) {
-      return readingIterator.spill();
-    }
-
     // We only write out contents of the inMemSorter if it is not empty.
     if (inMemSorter.numRecords() > 0) {
       final UnsafeSorterSpillWriter spillWriter =
@@ -269,12 +273,14 @@ private void deleteSpillFiles() {
    * Frees this sorter's in-memory data structures and cleans up its spill files.
    */
   public void cleanupResources() {
-    deleteSpillFiles();
-    freeMemory();
-    if (inMemSorter != null) {
-      inMemSorter = null;
-      releaseMemory(acquiredMem);
-      acquiredMem = 0;
+    synchronized (this) {
+      deleteSpillFiles();
+      freeMemory();
+      if (inMemSorter != null) {
+        inMemSorter = null;
+        releaseMemory(acquiredMem);
+        acquiredMem = 0;
+      }
     }
   }
 
@@ -286,14 +292,15 @@ public void cleanupResources() {
   private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
-      // assume that the memory of array will be doubled
-      long needed = inMemSorter.getMemoryUsage();
-      acquireMemory(needed);  // could trigger spilling
+      long used = inMemSorter.getMemoryUsage();
+      long needed = inMemSorter.getMemoryToExpand();
+      acquireMemory(used + needed);  // could trigger spilling
       if (inMemSorter.hasSpaceForAnotherRecord()) {
-        releaseMemory(needed);
+        releaseMemory(used + needed);
       } else {
-        acquiredMem += needed;
         inMemSorter.expandPointerArray();
+        releaseMemory(used);
+        acquiredMem += needed;
       }
     }
   }
@@ -391,16 +398,14 @@ public UnsafeSorterIterator getSortedIterator() throws IOException {
     }
   }
 
+  /**
+   * A
+   */
   class SpillableIterator extends UnsafeSorterIterator {
-    private Object baseObject = null;
-    private long baseOffset = 0;
-    private long keyPrefix = 0;
-    private int recordLength = 0;
-
-    private boolean cached = false;
+    private volatile UnsafeSorterIterator upstream;
+    private UnsafeSorterIterator nextUpstream = null;
     private MemoryBlock lastPage = null;
-
-    private UnsafeSorterIterator upstream;
+    private boolean loaded = false;
 
     public SpillableIterator(UnsafeInMemorySorter.SortedIterator inMemIterator) {
       this.upstream = inMemIterator;
@@ -415,12 +420,6 @@ public long spill() throws IOException {
         UnsafeInMemorySorter.SortedIterator inMemIterator =
           (UnsafeInMemorySorter.SortedIterator) upstream;
 
-        // backup the current record
-        baseObject = inMemIterator.getBaseObject();
-        baseOffset = inMemIterator.getBaseOffset();
-        keyPrefix = inMemIterator.getKeyPrefix();
-        recordLength = inMemIterator.getRecordLength();
-
         final UnsafeSorterSpillWriter spillWriter =
           new UnsafeSorterSpillWriter(blockManager, fileBufferSizeBytes, writeMetrics,
             inMemIterator.numRecordsLeft());
@@ -432,41 +431,47 @@ public long spill() throws IOException {
           spillWriter.write(baseObject, baseOffset, recordLength, inMemIterator.getKeyPrefix());
         }
         spillWriter.close();
-        upstream = spillWriter.getReader(blockManager);
+        nextUpstream = spillWriter.getReader(blockManager);
 
-        cached = true;
-        // release the pages except the one that is used
         long released = 0L;
-        for (MemoryBlock block : allocatedPages) {
-          if (block.getBaseObject() != baseObject) {
-            freePage(block);
-            released += block.size();
-          } else {
-            lastPage = block;
+        synchronized (UnsafeExternalSorter.this) {
+          // release the pages except the one that is used
+          for (MemoryBlock page : allocatedPages) {
+            if (!loaded || page.getBaseObject() != inMemIterator.getBaseObject()) {
+              released += page.size();
+              freePage(page);
+            } else {
+              lastPage = page;
+            }
           }
+          allocatedPages.clear();
+          assert(inMemSorter != null);
+          long used = inMemSorter.getMemoryUsage();
+          inMemSorter = null;
+          releaseMemory(used);
         }
-        allocatedPages.clear();
+
         return released;
       }
     }
 
     @Override
     public boolean hasNext() {
-      synchronized (this) {
-        return upstream.hasNext();
-      }
+      return upstream.hasNext();
     }
 
     @Override
     public void loadNext() throws IOException {
       synchronized (this) {
-        if (cached) {
+        loaded = true;
+        if (nextUpstream != null) {
           // Just consumed the last record from in memory iterator
           if (lastPage != null) {
             freePage(lastPage);
             lastPage = null;
           }
-          cached = false;
+          upstream = nextUpstream;
+          nextUpstream = null;
         }
         upstream.loadNext();
       }
@@ -474,42 +479,22 @@ public void loadNext() throws IOException {
 
     @Override
     public Object getBaseObject() {
-      synchronized (this) {
-        if (cached) {
-          return baseObject;
-        }
-        return upstream.getBaseObject();
-      }
+      return upstream.getBaseObject();
     }
 
     @Override
     public long getBaseOffset() {
-      synchronized (this) {
-        if (cached) {
-          return baseOffset;
-        }
-        return upstream.getBaseOffset();
-      }
+      return upstream.getBaseOffset();
     }
 
     @Override
     public int getRecordLength() {
-      synchronized (this) {
-        if (cached) {
-          return recordLength;
-        }
-        return upstream.getRecordLength();
-      }
+      return upstream.getRecordLength();
     }
 
     @Override
     public long getKeyPrefix() {
-      synchronized (this) {
-        if (cached) {
-          return keyPrefix;
-        }
-        return upstream.getKeyPrefix();
-      }
+      return upstream.getKeyPrefix();
     }
   }
 }
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
index 45596fba8b2fe..a6fc7cb13a6bb 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
@@ -70,12 +70,12 @@ public int compare(RecordPointerAndKeyPrefix r1, RecordPointerAndKeyPrefix r2) {
    * Within this buffer, position {@code 2 * i} holds a pointer pointer to the record at
    * index {@code i}, while position {@code 2 * i + 1} in the array holds an 8-byte key prefix.
    */
-  private long[] pointerArray;
+  private long[] array;
 
   /**
    * The position in the sort buffer where new records can be inserted.
    */
-  private int pointerArrayInsertPosition = 0;
+  private int pos = 0;
 
   public UnsafeInMemorySorter(
       final TaskMemoryManager memoryManager,
@@ -83,37 +83,43 @@ public UnsafeInMemorySorter(
       final PrefixComparator prefixComparator,
       int initialSize) {
     assert (initialSize > 0);
-    this.pointerArray = new long[initialSize * 2];
+    this.array = new long[initialSize * 2];
     this.memoryManager = memoryManager;
     this.sorter = new Sorter<>(UnsafeSortDataFormat.INSTANCE);
     this.sortComparator = new SortComparator(recordComparator, prefixComparator, memoryManager);
   }
 
   public void reset() {
-    pointerArrayInsertPosition = 0;
+    pos = 0;
   }
 
   /**
    * @return the number of records that have been inserted into this sorter.
    */
   public int numRecords() {
-    return pointerArrayInsertPosition / 2;
+    return pos / 2;
+  }
+
+  private int newLength() {
+    return array.length < Integer.MAX_VALUE / 2 ? (array.length * 2) : Integer.MAX_VALUE;
+  }
+
+  public long getMemoryToExpand() {
+    return (long) (newLength() - array.length) * 8L;
   }
 
   public long getMemoryUsage() {
-    return pointerArray.length * 8L;
+    return array.length * 8L;
   }
 
   public boolean hasSpaceForAnotherRecord() {
-    return pointerArrayInsertPosition + 2 <= pointerArray.length;
+    return pos + 2 <= array.length;
   }
 
   public void expandPointerArray() {
-    final long[] oldArray = pointerArray;
-    // Guard against overflow:
-    final int newLength = oldArray.length * 2 > 0 ? (oldArray.length * 2) : Integer.MAX_VALUE;
-    pointerArray = new long[newLength];
-    System.arraycopy(oldArray, 0, pointerArray, 0, oldArray.length);
+    final long[] oldArray = array;
+    array = new long[newLength()];
+    System.arraycopy(oldArray, 0, array, 0, oldArray.length);
   }
 
   /**
@@ -127,10 +133,10 @@ public void insertRecord(long recordPointer, long keyPrefix) {
     if (!hasSpaceForAnotherRecord()) {
       expandPointerArray();
     }
-    pointerArray[pointerArrayInsertPosition] = recordPointer;
-    pointerArrayInsertPosition++;
-    pointerArray[pointerArrayInsertPosition] = keyPrefix;
-    pointerArrayInsertPosition++;
+    array[pos] = recordPointer;
+    pos++;
+    array[pos] = keyPrefix;
+    pos++;
   }
 
   public static final class SortedIterator extends UnsafeSorterIterator {
@@ -191,7 +197,7 @@ public void loadNext() {
    * {@code next()} will return the same mutable object.
    */
   public SortedIterator getSortedIterator() {
-    sorter.sort(pointerArray, 0, pointerArrayInsertPosition / 2, sortComparator);
-    return new SortedIterator(memoryManager, pointerArrayInsertPosition, pointerArray);
+    sorter.sort(array, 0, pos / 2, sortComparator);
+    return new SortedIterator(memoryManager, pos, array);
   }
 }
diff --git a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala
index 6c9a71c3855b0..e37797b55654a 100644
--- a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala
+++ b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala
@@ -135,7 +135,7 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) exte
    * that control global sharing of memory between execution and storage.
    */
   private[memory]
-  final def acquireExecutionMemory(numBytes: Long, taskAttemptId: Long): Long = synchronized {
+  def acquireExecutionMemory(numBytes: Long, taskAttemptId: Long): Long = synchronized {
     assert(numBytes > 0, "invalid number of bytes requested: " + numBytes)
 
     // Add this task to the taskMemory map just so we can keep an accurate count of the number
diff --git a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
index ba6039994fa31..0fc04ace215cb 100644
--- a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
+++ b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
@@ -30,7 +30,7 @@ public class TaskMemoryManagerSuite {
   @Test
   public void leakedPageMemoryIsDetected() throws IOException {
     final TaskMemoryManager manager = new TaskMemoryManager(
-      new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
+      new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
     manager.allocatePage(4096, null);  // leak memory
     Assert.assertEquals(4096, manager.cleanUpAllAllocatedMemory());
   }
@@ -38,7 +38,7 @@ public void leakedPageMemoryIsDetected() throws IOException {
   @Test
   public void encodePageNumberAndOffsetOffHeap() throws IOException {
     final TaskMemoryManager manager = new TaskMemoryManager(
-      new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "true")), 0);
+      new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "true")), 0);
     final MemoryBlock dataPage = manager.allocatePage(256, null);
     // In off-heap mode, an offset is an absolute address that may require more than 51 bits to
     // encode. This test exercises that corner-case:
@@ -51,7 +51,7 @@ public void encodePageNumberAndOffsetOffHeap() throws IOException {
   @Test
   public void encodePageNumberAndOffsetOnHeap() throws IOException {
     final TaskMemoryManager manager = new TaskMemoryManager(
-      new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
+      new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
     final MemoryBlock dataPage = manager.allocatePage(256, null);
     final long encodedAddress = manager.encodePageNumberAndOffset(dataPage, 64);
     Assert.assertEquals(dataPage.getBaseObject(), manager.getPage(encodedAddress));
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java
index b5c781001687b..9a43f1f3a9235 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/PackedRecordPointerSuite.java
@@ -22,7 +22,7 @@
 import org.junit.Test;
 
 import org.apache.spark.SparkConf;
-import org.apache.spark.memory.GrantEverythingMemoryManager;
+import org.apache.spark.memory.TestMemoryManager;
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 
@@ -37,7 +37,7 @@ public class PackedRecordPointerSuite {
   public void heap() throws IOException {
     final SparkConf conf = new SparkConf().set("spark.unsafe.offHeap", "false");
     final TaskMemoryManager memoryManager =
-      new TaskMemoryManager(new GrantEverythingMemoryManager(conf), 0);
+      new TaskMemoryManager(new TestMemoryManager(conf), 0);
     final MemoryBlock page0 = memoryManager.allocatePage(128, null);
     final MemoryBlock page1 = memoryManager.allocatePage(128, null);
     final long addressInPage1 = memoryManager.encodePageNumberAndOffset(page1,
@@ -56,7 +56,7 @@ public void heap() throws IOException {
   public void offHeap() throws IOException {
     final SparkConf conf = new SparkConf().set("spark.unsafe.offHeap", "true");
     final TaskMemoryManager memoryManager =
-      new TaskMemoryManager(new GrantEverythingMemoryManager(conf), 0);
+      new TaskMemoryManager(new TestMemoryManager(conf), 0);
     final MemoryBlock page0 = memoryManager.allocatePage(128, null);
     final MemoryBlock page1 = memoryManager.allocatePage(128, null);
     final long addressInPage1 = memoryManager.encodePageNumberAndOffset(page1,
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java
index 719646ec2fbde..2293b1bbc113e 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java
@@ -26,7 +26,7 @@
 import org.apache.spark.HashPartitioner;
 import org.apache.spark.SparkConf;
 import org.apache.spark.unsafe.Platform;
-import org.apache.spark.memory.GrantEverythingMemoryManager;
+import org.apache.spark.memory.TestMemoryManager;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.memory.TaskMemoryManager;
 
@@ -60,7 +60,7 @@ public void testBasicSorting() throws Exception {
     };
     final SparkConf conf = new SparkConf().set("spark.unsafe.offHeap", "false");
     final TaskMemoryManager memoryManager =
-      new TaskMemoryManager(new GrantEverythingMemoryManager(conf), 0);
+      new TaskMemoryManager(new TestMemoryManager(conf), 0);
     final MemoryBlock dataPage = memoryManager.allocatePage(2048, null);
     final Object baseObject = dataPage.getBaseObject();
     final ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(4);
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
index a824e5b900830..a19e345c417b1 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
@@ -49,19 +49,19 @@
 import org.apache.spark.io.SnappyCompressionCodec;
 import org.apache.spark.executor.ShuffleWriteMetrics;
 import org.apache.spark.executor.TaskMetrics;
-import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.network.util.LimitedInputStream;
 import org.apache.spark.serializer.*;
 import org.apache.spark.scheduler.MapStatus;
 import org.apache.spark.shuffle.IndexShuffleBlockResolver;
 import org.apache.spark.storage.*;
-import org.apache.spark.memory.GrantEverythingMemoryManager;
+import org.apache.spark.memory.TestMemoryManager;
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.util.Utils;
 
 public class UnsafeShuffleWriterSuite {
 
   static final int NUM_PARTITITONS = 4;
+  TestMemoryManager memoryManager;
   TaskMemoryManager taskMemoryManager;
   final HashPartitioner hashPartitioner = new HashPartitioner(NUM_PARTITITONS);
   File mergedOutputFile;
@@ -107,10 +107,11 @@ public void setUp() throws IOException {
     partitionSizesInMergedFile = null;
     spillFilesCreated.clear();
     conf = new SparkConf()
-      .set("spark.buffer.pageSize", "128m")
+      .set("spark.buffer.pageSize", "1m")
       .set("spark.unsafe.offHeap", "false");
     taskMetrics = new TaskMetrics();
-    taskMemoryManager =  new TaskMemoryManager(new GrantEverythingMemoryManager(conf), 0);
+    memoryManager = new TestMemoryManager(conf);
+    taskMemoryManager =  new TaskMemoryManager(memoryManager, 0);
 
     when(blockManager.diskBlockManager()).thenReturn(diskBlockManager);
     when(blockManager.getDiskWriter(
@@ -397,20 +398,14 @@ public void mergeSpillsWithFileStreamAndNoCompression() throws Exception {
 
   @Test
   public void writeEnoughDataToTriggerSpill() throws Exception {
-    taskMemoryManager = spy(taskMemoryManager);
-    doCallRealMethod() // initialize sort buffer
-      .doCallRealMethod() // allocate initial data page
-      .doReturn(0L) // deny request to allocate new page
-      .doCallRealMethod() // grant new sort buffer and data page
-      .when(taskMemoryManager).acquireExecutionMemory(anyLong(), null);
+    memoryManager.limit(PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES);
     final UnsafeShuffleWriter writer = createWriter(false);
     final ArrayList> dataToWrite = new ArrayList>();
-    final byte[] bigByteArray = new byte[PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES / 128];
-    for (int i = 0; i < 128 + 1; i++) {
+    final byte[] bigByteArray = new byte[PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES / 10];
+    for (int i = 0; i < 10 + 1; i++) {
       dataToWrite.add(new Tuple2(i, bigByteArray));
     }
     writer.write(dataToWrite.iterator());
-    verify(taskMemoryManager, times(3)).acquireExecutionMemory(anyLong(), any(MemoryConsumer.class));
     assertEquals(2, spillFilesCreated.size());
     writer.stop(true);
     readRecordsFromFile();
@@ -425,19 +420,13 @@ public void writeEnoughDataToTriggerSpill() throws Exception {
 
   @Test
   public void writeEnoughRecordsToTriggerSortBufferExpansionAndSpill() throws Exception {
-    taskMemoryManager = spy(taskMemoryManager);
-    doCallRealMethod() // initialize sort buffer
-      .doCallRealMethod() // allocate initial data page
-      .doReturn(0L) // deny request to allocate new page
-      .doCallRealMethod() // grant new sort buffer and data page
-      .when(taskMemoryManager).acquireExecutionMemory(anyLong(), null);
+    memoryManager.limit(UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE * 16);
     final UnsafeShuffleWriter writer = createWriter(false);
     final ArrayList> dataToWrite = new ArrayList<>();
-    for (int i = 0; i < UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE; i++) {
+    for (int i = 0; i < UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE * 2; i++) {
       dataToWrite.add(new Tuple2(i, i));
     }
     writer.write(dataToWrite.iterator());
-    verify(taskMemoryManager, times(2)).acquireExecutionMemory(anyLong(), any(MemoryConsumer.class));
     assertEquals(2, spillFilesCreated.size());
     writer.stop(true);
     readRecordsFromFile();
@@ -472,11 +461,11 @@ public void writeRecordsThatAreBiggerThanMaxRecordSize() throws Exception {
     final ArrayList> dataToWrite = new ArrayList>();
     dataToWrite.add(new Tuple2(1, ByteBuffer.wrap(new byte[1])));
     // We should be able to write a record that's right _at_ the max record size
-    final byte[] atMaxRecordSize = new byte[writer.maxRecordSizeBytes()];
+    final byte[] atMaxRecordSize = new byte[(int) taskMemoryManager.pageSizeBytes() - 4];
     new Random(42).nextBytes(atMaxRecordSize);
     dataToWrite.add(new Tuple2(2, ByteBuffer.wrap(atMaxRecordSize)));
     // Inserting a record that's larger than the max record size
-    final byte[] exceedsMaxRecordSize = new byte[writer.maxRecordSizeBytes() + 1];
+    final byte[] exceedsMaxRecordSize = new byte[(int) taskMemoryManager.pageSizeBytes()];
     new Random(42).nextBytes(exceedsMaxRecordSize);
     dataToWrite.add(new Tuple2(3, ByteBuffer.wrap(exceedsMaxRecordSize)));
     writer.write(dataToWrite.iterator());
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index 3c70bdce666f2..454db433d0a4c 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -39,7 +39,7 @@
 
 import org.apache.spark.SparkConf;
 import org.apache.spark.executor.ShuffleWriteMetrics;
-import org.apache.spark.memory.GrantEverythingMemoryManager;
+import org.apache.spark.memory.TestMemoryManager;
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.serializer.SerializerInstance;
 import org.apache.spark.storage.*;
@@ -62,7 +62,7 @@ public abstract class AbstractBytesToBytesMapSuite {
 
   private final Random rand = new Random(42);
 
-  private GrantEverythingMemoryManager memoryManager;
+  private TestMemoryManager memoryManager;
   private TaskMemoryManager taskMemoryManager;
   private final long PAGE_SIZE_BYTES = 1L << 26; // 64 megabytes
 
@@ -82,7 +82,7 @@ public OutputStream apply(OutputStream stream) {
   @Before
   public void setup() {
     memoryManager =
-      new GrantEverythingMemoryManager(
+      new TestMemoryManager(
         new SparkConf().set("spark.unsafe.offHeap", "" + useOffHeapMemoryAllocator()));
     taskMemoryManager = new TaskMemoryManager(memoryManager, 0);
 
@@ -492,7 +492,7 @@ public void randomizedTestWithRecordsLargerThanPageSize() {
 
   @Test
   public void failureToAllocateFirstPage() {
-    memoryManager.markExecutionAsOutOfMemoryAlways();
+    memoryManager.limit(0);
     BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, 1, PAGE_SIZE_BYTES);
     try {
       final long[] emptyArray = new long[0];
@@ -503,7 +503,6 @@ public void failureToAllocateFirstPage() {
         emptyArray, Platform.LONG_ARRAY_OFFSET, 0, emptyArray, Platform.LONG_ARRAY_OFFSET, 0));
     } finally {
       map.free();
-      memoryManager.resetOutOfMemory();
     }
   }
 
@@ -516,7 +515,7 @@ public void failureToGrow() {
       int i;
       for (i = 0; i < 127; i++) {
         if (i > 0) {
-          memoryManager.markExecutionAsOutOfMemoryAlways();
+          memoryManager.limit(0);
         }
         final long[] arr = new long[]{i};
         final BytesToBytesMap.Location loc = map.lookup(arr, Platform.LONG_ARRAY_OFFSET, 8);
@@ -530,7 +529,6 @@ public void failureToGrow() {
       Assert.assertFalse(success);
     } finally {
       map.free();
-      memoryManager.resetOutOfMemory();
     }
   }
 
diff --git a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
index 3c7d0791ac09a..58d7744f14c48 100644
--- a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
+++ b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
@@ -41,7 +41,7 @@
 import org.apache.spark.TaskContext;
 import org.apache.spark.executor.ShuffleWriteMetrics;
 import org.apache.spark.executor.TaskMetrics;
-import org.apache.spark.memory.GrantEverythingMemoryManager;
+import org.apache.spark.memory.TestMemoryManager;
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.serializer.SerializerInstance;
 import org.apache.spark.storage.*;
@@ -57,8 +57,8 @@
 public class UnsafeExternalSorterSuite {
 
   final LinkedList spillFilesCreated = new LinkedList();
-  final GrantEverythingMemoryManager memoryManager =
-    new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false"));
+  final TestMemoryManager memoryManager =
+    new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false"));
   final TaskMemoryManager taskMemoryManager = new TaskMemoryManager(memoryManager, 0);
   // Use integer comparison for comparing prefixes (which are partition ids, in this case)
   final PrefixComparator prefixComparator = new PrefixComparator() {
@@ -184,9 +184,9 @@ public void testSortingOnlyByPrefix() throws Exception {
     insertNumber(sorter, 5);
     insertNumber(sorter, 1);
     insertNumber(sorter, 3);
-    sorter.spill(0);
+    sorter.spill();
     insertNumber(sorter, 4);
-    sorter.spill(0);
+    sorter.spill();
     insertNumber(sorter, 2);
 
     UnsafeSorterIterator iter = sorter.getSortedIterator();
@@ -207,9 +207,9 @@ public void testSortingEmptyArrays() throws Exception {
     final UnsafeExternalSorter sorter = newSorter();
     sorter.insertRecord(null, 0, 0, 0);
     sorter.insertRecord(null, 0, 0, 0);
-    sorter.spill(0L);
+    sorter.spill();
     sorter.insertRecord(null, 0, 0, 0);
-    sorter.spill(0L);
+    sorter.spill();
     sorter.insertRecord(null, 0, 0, 0);
     sorter.insertRecord(null, 0, 0, 0);
 
@@ -275,9 +275,9 @@ public void sortingRecordsThatExceedPageSize() throws Exception {
     Arrays.fill(smallRecord, 123);
 
     insertRecord(sorter, largeRecord, 456);
-    sorter.spill(0L);
+    sorter.spill();
     insertRecord(sorter, smallRecord, 123);
-    sorter.spill(0L);
+    sorter.spill();
     insertRecord(sorter, smallRecord, 123);
     insertRecord(sorter, largeRecord, 456);
 
@@ -345,7 +345,7 @@ public void testPeakMemoryUsed() throws Exception {
       }
 
       // Spilling should not change peak memory
-      sorter.spill(0L);
+      sorter.spill();
       newPeakMemory = sorter.getPeakMemoryUsedBytes();
       assertEquals(previousPeakMemory, newPeakMemory);
       for (int i = 0; i < numRecordsPerPage; i++) {
diff --git a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java
index 6a2649e1a19b8..642f6585f8a15 100644
--- a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java
+++ b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java
@@ -23,7 +23,7 @@
 
 import org.apache.spark.HashPartitioner;
 import org.apache.spark.SparkConf;
-import org.apache.spark.memory.GrantEverythingMemoryManager;
+import org.apache.spark.memory.TestMemoryManager;
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.memory.MemoryBlock;
@@ -46,7 +46,7 @@ private static String getStringFromDataPage(Object baseObject, long baseOffset,
   public void testSortingEmptyInput() {
     final UnsafeInMemorySorter sorter = new UnsafeInMemorySorter(
       new TaskMemoryManager(
-        new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0),
+        new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0),
       mock(RecordComparator.class),
       mock(PrefixComparator.class),
       100);
@@ -68,7 +68,7 @@ public void testSortingOnlyByIntegerPrefix() throws Exception {
       "Mango"
     };
     final TaskMemoryManager memoryManager = new TaskMemoryManager(
-      new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
+      new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
     final MemoryBlock dataPage = memoryManager.allocatePage(2048, null);
     final Object baseObject = dataPage.getBaseObject();
     // Write the records into the data page:
diff --git a/core/src/test/scala/org/apache/spark/memory/GrantEverythingMemoryManager.scala b/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala
similarity index 78%
rename from core/src/test/scala/org/apache/spark/memory/GrantEverythingMemoryManager.scala
rename to core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala
index cde7b48b461dc..2bb4452dde74c 100644
--- a/core/src/test/scala/org/apache/spark/memory/GrantEverythingMemoryManager.scala
+++ b/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala
@@ -22,18 +22,20 @@ import scala.collection.mutable
 import org.apache.spark.SparkConf
 import org.apache.spark.storage.{BlockStatus, BlockId}
 
-class GrantEverythingMemoryManager(conf: SparkConf) extends MemoryManager(conf, numCores = 1) {
+class TestMemoryManager(conf: SparkConf) extends MemoryManager(conf, numCores = 1) {
   private[memory] override def doAcquireExecutionMemory(
       numBytes: Long,
       evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Long = synchronized {
-    if (oomAlways) {
-      0
-    } else if (oomOnce) {
+    if (oomOnce) {
       oomOnce = false
       0
-    } else {
+    } else if (available >= numBytes) {
       _executionMemoryUsed += numBytes // To suppress warnings when freeing unallocated memory
+      available -= numBytes
       numBytes
+    } else {
+      _executionMemoryUsed += available
+      available
     }
   }
   override def acquireStorageMemory(
@@ -44,22 +46,22 @@ class GrantEverythingMemoryManager(conf: SparkConf) extends MemoryManager(conf,
       blockId: BlockId,
       numBytes: Long,
       evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = true
-  override def releaseStorageMemory(numBytes: Long): Unit = { }
+  override def releaseExecutionMemory(numBytes: Long): Unit = {
+    available += numBytes
+    _executionMemoryUsed -= numBytes
+  }
+  override def releaseStorageMemory(numBytes: Long): Unit = {}
   override def maxExecutionMemory: Long = Long.MaxValue
   override def maxStorageMemory: Long = Long.MaxValue
 
   private var oomOnce = false
-  private var oomAlways = false
+  private var available = Long.MaxValue
 
   def markExecutionAsOutOfMemoryOnce(): Unit = {
     oomOnce = true
   }
 
-  def markExecutionAsOutOfMemoryAlways(): Unit = {
-    oomAlways = true
-  }
-
-  def resetOutOfMemory(): Unit = {
-    oomAlways = false
+  def limit(avail: Long): Unit = {
+    available = avail
   }
 }
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java b/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
index a56d31c2ef369..f7063d1e5c829 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
@@ -96,15 +96,10 @@ void insertRow(UnsafeRow row) throws IOException {
     );
     numRowsInserted++;
     if (testSpillFrequency > 0 && (numRowsInserted % testSpillFrequency) == 0) {
-      spill();
+      sorter.spill();
     }
   }
 
-  @VisibleForTesting
-  void spill() throws IOException {
-    sorter.spill(Long.MAX_VALUE);
-  }
-
   /**
    * Return the peak memory used so far, in bytes.
    */
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
index bd5f91f36726b..4bfaf60257488 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
@@ -125,7 +125,7 @@ public UnsafeKVExternalSorter(
         pageSizeBytes,
         inMemSorter);
 
-      sorter.spill(Long.MAX_VALUE);
+      sorter.spill();
       map.free();
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
index 23e66adcaffc5..4cd55cf96d848 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
@@ -24,7 +24,7 @@ import scala.util.{Try, Random}
 import org.scalatest.Matchers
 
 import org.apache.spark.{SparkConf, TaskContextImpl, TaskContext, SparkFunSuite}
-import org.apache.spark.memory.{TaskMemoryManager, GrantEverythingMemoryManager}
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{UnsafeRow, UnsafeProjection}
 import org.apache.spark.sql.test.SharedSQLContext
@@ -48,7 +48,7 @@ class UnsafeFixedWidthAggregationMapSuite
   private def emptyAggregationBuffer: InternalRow = InternalRow(0)
   private val PAGE_SIZE_BYTES: Long = 1L << 26; // 64 megabytes
 
-  private var memoryManager: GrantEverythingMemoryManager = null
+  private var memoryManager: TestMemoryManager = null
   private var taskMemoryManager: TaskMemoryManager = null
 
   def testWithMemoryLeakDetection(name: String)(f: => Unit) {
@@ -62,7 +62,7 @@ class UnsafeFixedWidthAggregationMapSuite
 
     test(name) {
       val conf = new SparkConf().set("spark.unsafe.offHeap", "false")
-      memoryManager = new GrantEverythingMemoryManager(conf)
+      memoryManager = new TestMemoryManager(conf)
       taskMemoryManager = new TaskMemoryManager(memoryManager, 0)
 
       TaskContext.setTaskContext(new TaskContextImpl(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeKVExternalSorterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeKVExternalSorterSuite.scala
index 261861fbf81a5..7b80963ec8708 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeKVExternalSorterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeKVExternalSorterSuite.scala
@@ -20,7 +20,7 @@ package org.apache.spark.sql.execution
 import scala.util.Random
 
 import org.apache.spark._
-import org.apache.spark.memory.{TaskMemoryManager, GrantEverythingMemoryManager}
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
 import org.apache.spark.sql.{RandomDataGenerator, Row}
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
 import org.apache.spark.sql.catalyst.expressions.{InterpretedOrdering, UnsafeRow, UnsafeProjection}
@@ -109,7 +109,7 @@ class UnsafeKVExternalSorterSuite extends SparkFunSuite with SharedSQLContext {
       pageSize: Long,
       spill: Boolean): Unit = {
     val memoryManager =
-      new GrantEverythingMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false"))
+      new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false"))
     val taskMemMgr = new TaskMemoryManager(memoryManager, 0)
     TaskContext.setTaskContext(new TaskContextImpl(
       stageId = 0,

From ce24f039f07bde78500191a9afec7a6298a6b559 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Tue, 27 Oct 2015 18:03:36 -0700
Subject: [PATCH 05/17] add more tests

---
 .../spark/unsafe/map/BytesToBytesMap.java     | 45 ++++++++++------
 .../unsafe/sort/UnsafeExternalSorter.java     |  4 +-
 .../unsafe/sort/UnsafeInMemorySorter.java     | 10 ++++
 .../map/AbstractBytesToBytesMapSuite.java     |  1 +
 .../sort/UnsafeExternalSorterSuite.java       | 54 +++++++++++++++++++
 .../UnsafeFixedWidthAggregationMapSuite.scala |  2 +-
 6 files changed, 98 insertions(+), 18 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 04144f599c967..5cbf234f102e5 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -18,15 +18,20 @@
 package org.apache.spark.unsafe.map;
 
 import javax.annotation.Nullable;
+import java.io.File;
 import java.io.IOException;
 import java.util.Iterator;
 import java.util.LinkedList;
 
+import scala.runtime.AbstractFunction0;
+import scala.runtime.BoxedUnit;
+
 import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import org.apache.spark.SparkEnv;
+import org.apache.spark.TaskContext;
 import org.apache.spark.executor.ShuffleWriteMetrics;
 import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.memory.TaskMemoryManager;
@@ -235,8 +240,7 @@ public final class MapIterator implements Iterator {
     // If this iterator destructive or not. When it is true, it frees each page as it moves onto
     // next one.
     private boolean destructive = false;
-    private LinkedList spillWriters =
-      new LinkedList();
+    private LinkedList spillWriters = new LinkedList<>();
     private UnsafeSorterSpillReader reader = null;
 
     private MapIterator(int numRecords, Location loc, boolean destructive) {
@@ -264,8 +268,17 @@ private void advanceToNextPage() {
           offsetInPage += 4;
         } else {
           currentPage = null;
+          if (reader != null) {
+            // remove the spill file from disk
+            File file = spillWriters.removeFirst().getFile();
+            if (file != null && file.exists()) {
+              if (!file.delete()) {
+                logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
+              }
+            }
+          }
           try {
-            reader = spillWriters.removeFirst().getReader(blockManager);
+            reader = spillWriters.getFirst().getReader(blockManager);
             recordsInPage = -1;
           } catch (IOException e) {
             // Scala iterator does not handle exception
@@ -338,18 +351,20 @@ public long spill(long numBytes) throws IOException {
           }
           writer.close();
           spillWriters.add(writer);
-//          taskContext.addOnCompleteCallback(new AbstractFunction0() {
-//            @Override
-//            public BoxedUnit apply() {
-//              File file = writer.getFile();
-//              if (file != null && file.exists()) {
-//                if (!file.delete()) {
-//                  logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
-//                }
-//              }
-//              return null;
-//            }
-//          });
+          if (TaskContext.get() != null) {
+            TaskContext.get().addOnCompleteCallback(new AbstractFunction0() {
+              @Override
+              public BoxedUnit apply() {
+                File file = writer.getFile();
+                if (file != null && file.exists()) {
+                  if (!file.delete()) {
+                    logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
+                  }
+                }
+                return null;
+              }
+            });
+          }
 
           dataPages.removeLast();
           freePage(block);
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index cabe29fcc1675..af4af97c87a00 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -418,7 +418,7 @@ public long spill() throws IOException {
         }
 
         UnsafeInMemorySorter.SortedIterator inMemIterator =
-          (UnsafeInMemorySorter.SortedIterator) upstream;
+          ((UnsafeInMemorySorter.SortedIterator) upstream).clone();
 
         final UnsafeSorterSpillWriter spillWriter =
           new UnsafeSorterSpillWriter(blockManager, fileBufferSizeBytes, writeMetrics,
@@ -431,6 +431,7 @@ public long spill() throws IOException {
           spillWriter.write(baseObject, baseOffset, recordLength, inMemIterator.getKeyPrefix());
         }
         spillWriter.close();
+        spillWriters.add(spillWriter);
         nextUpstream = spillWriter.getReader(blockManager);
 
         long released = 0L;
@@ -450,7 +451,6 @@ public long spill() throws IOException {
           inMemSorter = null;
           releaseMemory(used);
         }
-
         return released;
       }
     }
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
index a6fc7cb13a6bb..f9a6d49398744 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
@@ -159,6 +159,16 @@ private SortedIterator(
       this.sortBuffer = sortBuffer;
     }
 
+    public SortedIterator clone () {
+      SortedIterator iter = new SortedIterator(memoryManager, sortBufferInsertPosition, sortBuffer);
+      iter.position = position;
+      iter.baseObject = baseObject;
+      iter.baseOffset = baseOffset;
+      iter.keyPrefix = keyPrefix;
+      iter.recordLength = recordLength;
+      return iter;
+    }
+
     @Override
     public boolean hasNext() {
       return position < sortBufferInsertPosition;
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index 454db433d0a4c..6cdedcfcb9a8d 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -551,6 +551,7 @@ public void spillInIterator() throws IOException {
       for (i = 100; i < 1024; i++) {
         iter.next();
       }
+
       BytesToBytesMap.MapIterator iter2 = map.destructiveIterator();
       for (i = 0; i < 100; i++) {
         iter2.next();
diff --git a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
index 58d7744f14c48..dc654998e2950 100644
--- a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
+++ b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
@@ -312,6 +312,60 @@ public void sortingRecordsThatExceedPageSize() throws Exception {
     assertSpillFilesWereCleanedUp();
   }
 
+  @Test
+  public void forcedSpillingWithReadIterator() throws Exception {
+    final UnsafeExternalSorter sorter = newSorter();
+    long[] record = new long[100];
+    int recordSize = record.length * 8;
+    int n = (int) pageSizeBytes / recordSize * 3;
+    for (int i = 0; i < n; i++) {
+      record[0] = (long) i;
+      sorter.insertRecord(record, Platform.LONG_ARRAY_OFFSET, recordSize, 0);
+    }
+    assert(sorter.getNumberOfAllocatedPages() >= 2);
+    UnsafeExternalSorter.SpillableIterator iter =
+      (UnsafeExternalSorter.SpillableIterator) sorter.getSortedIterator();
+    int lastv = 0;
+    for (int i = 0; i < n / 3; i++) {
+      iter.hasNext();
+      iter.loadNext();
+      assert(Platform.getLong(iter.getBaseObject(), iter.getBaseOffset()) == i);
+      lastv = i;
+    }
+    iter.spill();
+    assert(Platform.getLong(iter.getBaseObject(), iter.getBaseOffset()) == lastv);
+    for (int i = n / 3; i < n; i++) {
+      iter.hasNext();
+      iter.loadNext();
+      assert(Platform.getLong(iter.getBaseObject(), iter.getBaseOffset()) == i);
+    }
+    sorter.cleanupResources();
+    assertSpillFilesWereCleanedUp();
+  }
+
+  @Test
+  public void forcedSpillingWithNotReadIterator() throws Exception {
+    final UnsafeExternalSorter sorter = newSorter();
+    long[] record = new long[100];
+    int recordSize = record.length * 8;
+    int n = (int) pageSizeBytes / recordSize * 3;
+    for (int i = 0; i < n; i++) {
+      record[0] = (long) i;
+      sorter.insertRecord(record, Platform.LONG_ARRAY_OFFSET, recordSize, 0);
+    }
+    assert(sorter.getNumberOfAllocatedPages() >= 2);
+    UnsafeExternalSorter.SpillableIterator iter =
+      (UnsafeExternalSorter.SpillableIterator) sorter.getSortedIterator();
+    iter.spill();
+    for (int i = 0; i < n; i++) {
+      iter.hasNext();
+      iter.loadNext();
+      assert(Platform.getLong(iter.getBaseObject(), iter.getBaseOffset()) == i);
+    }
+    sorter.cleanupResources();
+    assertSpillFilesWereCleanedUp();
+  }
+
   @Test
   public void testPeakMemoryUsed() throws Exception {
     final long recordLengthBytes = 8;
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
index 4cd55cf96d848..e36d9599d2fe4 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
@@ -342,7 +342,7 @@ class UnsafeFixedWidthAggregationMapSuite
       buf.setInt(0, str.length)
     }
     // Simulate running out of space
-    memoryManager.markExecutionAsOutOfMemoryAlways()
+    memoryManager.limit(0)
     val str = rand.nextString(4096)
     val buf = map.getAggregationBuffer(InternalRow(UTF8String.fromString(str)))
     assert(buf == null)

From 8470fc9ddd37c525cf648c00e019d968410fa66f Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Tue, 27 Oct 2015 23:07:59 -0700
Subject: [PATCH 06/17] address coments

---
 .../main/java/org/apache/spark/memory/MemoryConsumer.java   | 6 ++----
 .../java/org/apache/spark/memory/TaskMemoryManager.java     | 6 ++++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
index 6d7421cf8b86a..7e375b200e901 100644
--- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
+++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
@@ -26,7 +26,7 @@
 /**
  * An memory consumer of TaskMemoryManager, which support spilling.
  */
-public class MemoryConsumer {
+public abstract class MemoryConsumer {
 
   private TaskMemoryManager memoryManager;
   private long pageSize;
@@ -58,9 +58,7 @@ public long spill() throws IOException {
    * @return the amount of released memory in bytes
    * @throws IOException
    */
-  public long spill(long size, MemoryConsumer trigger) throws IOException {
-    return 0L;
-  }
+  public abstract long spill(long size, MemoryConsumer trigger) throws IOException;
 
   /**
    * Acquire `size` bytes memory.
diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 0303e93d77123..4f70e3909bba0 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -108,7 +108,7 @@ public class TaskMemoryManager {
   /**
    * The size of memory granted to each consumer.
    */
-  private HashMap consumers;
+  private final HashMap consumers;
 
   /**
    * Construct a new TaskMemoryManager.
@@ -127,6 +127,7 @@ public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) {
    * @return number of bytes successfully granted (<= N).
    */
   public long acquireExecutionMemory(long size, MemoryConsumer consumer) throws IOException {
+    assert(size >= 0);
     synchronized (this) {
       long got = memoryManager.acquireExecutionMemory(size, taskAttemptId);
 
@@ -164,6 +165,7 @@ public long acquireExecutionMemory(long size, MemoryConsumer consumer) throws IO
    * Release N bytes of execution memory for a MemoryConsumer.
    */
   public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
+    assert(size >= 0);
     if (size == 0) {
       return;
     }
@@ -179,7 +181,7 @@ public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
                 new SparkException("Release more memory " + size + "than acquired " + old + " for "
                   + consumer));
             } else {
-              logger.warn("Release more memory " + size + " than acquired " + old + "for "
+              logger.warn("Release more memory " + size + " than acquired " + old + " for "
                 + consumer);
             }
           }

From 827d4f0a6d42e4f2af37eedacf3098fceb896c32 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Wed, 28 Oct 2015 14:14:48 -0700
Subject: [PATCH 07/17] add tests for TaskMemoryManager

---
 .../apache/spark/memory/MemoryConsumer.java   | 20 ++--
 .../spark/memory/TaskMemoryManager.java       | 91 ++++++++++++-------
 .../spark/memory/TaskMemoryManagerSuite.java  | 68 ++++++++++++++
 .../spark/memory/TestMemoryManager.scala      |  5 +-
 4 files changed, 143 insertions(+), 41 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
index 7e375b200e901..7f8cf90883572 100644
--- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
+++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
@@ -41,18 +41,22 @@ protected MemoryConsumer(TaskMemoryManager memoryManager) {
   }
 
   /**
-   * Force spill during during.
+   * Force spill during building.
    *
    * For testing.
    */
-  public long spill() throws IOException {
-    return spill(Long.MAX_VALUE, this);
+  public void spill() throws IOException {
+    spill(Long.MAX_VALUE, this);
   }
 
   /**
    * Spill some data to disk to release memory, which will be called by TaskMemoryManager
    * when there is not enough memory for the task.
    *
+   * This should be implemented by subclass.
+   *
+   * Note: In order to avoid possible deadlock, should not call acquireMemory() from spill().
+   *
    * @param size the amount of memory should be released
    * @param trigger the MemoryConsumer that trigger this spilling
    * @return the amount of released memory in bytes
@@ -70,12 +74,13 @@ public long spill() throws IOException {
   protected void acquireMemory(long size) throws IOException {
     long got = memoryManager.acquireExecutionMemory(size, this);
     if (got < size) {
-      throw new IOException("Could not acquire " + size + " bytes of memory " + got);
+      memoryManager.showMemoryUsage();
+      throw new IOException("Could not acquire " + size + " bytes of memory, got " + got);
     }
   }
 
   /**
-   * Release amount of memory.
+   * Release `size` bytes memory.
    */
   protected void releaseMemory(long size) {
     memoryManager.releaseExecutionMemory(size, this);
@@ -91,10 +96,13 @@ protected void releaseMemory(long size) {
   protected MemoryBlock allocatePage(long required) throws IOException {
     MemoryBlock page = memoryManager.allocatePage(Math.max(pageSize, required), this);
     if (page == null || page.size() < required) {
+      long got = 0;
       if (page != null) {
+        got = page.size();
         freePage(page);
       }
-      throw new IOException("Unable to acquire " + required + " bytes of memory");
+      memoryManager.showMemoryUsage();
+      throw new IOException("Unable to acquire " + required + " bytes of memory, got " + got);
     }
     return page;
   }
diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 4f70e3909bba0..14553f4c35272 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -126,37 +126,42 @@ public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) {
    *
    * @return number of bytes successfully granted (<= N).
    */
-  public long acquireExecutionMemory(long size, MemoryConsumer consumer) throws IOException {
-    assert(size >= 0);
+  public long acquireExecutionMemory(long required, MemoryConsumer consumer) throws IOException {
+    assert(required >= 0);
     synchronized (this) {
-      long got = memoryManager.acquireExecutionMemory(size, taskAttemptId);
-
-      // call spill() on itself to release some memory
-      if (got < size && consumer != null) {
-        consumer.spill(size - got, consumer);
-        got += memoryManager.acquireExecutionMemory(size - got, taskAttemptId);
-      }
-
-      if (got < size) {
-        long needed = size - got;
-        // call spill() on other consumers to release memory
-        for (MemoryConsumer c: consumers.keySet()) {
-          if (c != null && c != consumer) {
-            needed -= c.spill(size - got, consumer);
-            if (needed < 0) {
-              break;
+      long got = memoryManager.acquireExecutionMemory(required, taskAttemptId);
+
+      if (got < required) {
+        // consumers could be modified by spill(), so we should have a copy here.
+        MemoryConsumer[] cs = new MemoryConsumer[consumers.size()];
+        consumers.keySet().toArray(cs);
+        // Call spill() on other consumers to release memory
+        for (MemoryConsumer c: cs) {
+          if (c != null) {
+            long released = c.spill(required - got, consumer);
+            logger.info("released " + Utils.bytesToString(released) + " from " + consumer);
+            if (released >= 0) {
+              got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId);
+              if (got >= required) {
+                break;
+              }
             }
           }
         }
-        got += memoryManager.acquireExecutionMemory(size - got, taskAttemptId);
       }
 
-      long old = 0L;
-      if (consumers.containsKey(consumer)) {
-        old = consumers.get(consumer);
+      // Update the accounting, even consumer is null
+      if (got > 0) {
+        long old = 0L;
+        if (consumers.containsKey(consumer)) {
+          old = consumers.get(consumer);
+        }
+        consumers.put(consumer, got + old);
       }
-      consumers.put(consumer, got + old);
 
+      if (logger.isTraceEnabled()) {
+        logger.trace("Acquire {} for {}", Utils.bytesToString(got), consumer);
+      }
       return got;
     }
   }
@@ -176,29 +181,45 @@ public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
           consumers.put(consumer, old - size);
         } else {
           if (old < size) {
+            String msg = "Release " + size + " bytes memory (more than acquired " + old + ") for "
+              + consumer;
+            logger.warn(msg);
             if (Utils.isTesting()) {
-              Platform.throwException(
-                new SparkException("Release more memory " + size + "than acquired " + old + " for "
-                  + consumer));
-            } else {
-              logger.warn("Release more memory " + size + " than acquired " + old + " for "
-                + consumer);
+              Platform.throwException(new SparkException(msg));
             }
           }
           consumers.remove(consumer);
         }
       } else {
+        String msg = "Release " + size + " bytes memory for non-existent " + consumer;
+        logger.warn(msg);
         if (Utils.isTesting()) {
-          Platform.throwException(
-            new SparkException("Release memory " + size + " for not existed " + consumer));
-        } else {
-          logger.warn("Release memory " + size + " for not existed " + consumer);
+          Platform.throwException(new SparkException(msg));
         }
       }
-      memoryManager.releaseExecutionMemory(size, taskAttemptId);
+    }
+
+    if (logger.isTraceEnabled()) {
+      logger.trace("Release {} from {}", Utils.bytesToString(size), consumer);
+    }
+    memoryManager.releaseExecutionMemory(size, taskAttemptId);
+  }
+
+  /**
+   * Dump the memory usage of all consumers.
+   */
+  public void showMemoryUsage() {
+    logger.info("Memory used in task " + taskAttemptId);
+    synchronized (this) {
+      for (MemoryConsumer c: consumers.keySet()) {
+        logger.info("acquired by " + c + ": " + Utils.bytesToString(consumers.get(c)));
+      }
     }
   }
 
+  /**
+   * Return the page size in bytes.
+   */
   public long pageSizeBytes() {
     return memoryManager.pageSizeBytes();
   }
@@ -224,6 +245,7 @@ public MemoryBlock allocatePage(long size, MemoryConsumer consumer) throws IOExc
     synchronized (this) {
       pageNumber = allocatedPages.nextClearBit(0);
       if (pageNumber >= PAGE_TABLE_SIZE) {
+        releaseExecutionMemory(acquired, consumer);
         throw new IllegalStateException(
           "Have already allocated a maximum of " + PAGE_TABLE_SIZE + " pages");
       }
@@ -340,6 +362,7 @@ public long cleanUpAllAllocatedMemory() {
         freePage(page, null);
       }
     }
+    consumers.clear();
 
     freedBytes += memoryManager.releaseAllExecutionMemoryForTask(taskAttemptId);
 
diff --git a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
index 0fc04ace215cb..3e7351b37a8c5 100644
--- a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
+++ b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
@@ -27,6 +27,32 @@
 
 public class TaskMemoryManagerSuite {
 
+  class TestMemoryConsumer extends MemoryConsumer {
+    volatile long used = 0L;
+
+    TestMemoryConsumer(TaskMemoryManager memoryManager) {
+      super(memoryManager);
+    }
+
+    @Override
+    public long spill(long size, MemoryConsumer trigger) throws IOException {
+      releaseMemory(used);
+      long released = used;
+      used = 0;
+      return released;
+    }
+
+    void use(long size) throws IOException {
+      acquireMemory(size);
+      used += size;
+    }
+
+    void free(long size) throws IOException {
+      releaseMemory(size);
+      used -= size;
+    }
+  }
+
   @Test
   public void leakedPageMemoryIsDetected() throws IOException {
     final TaskMemoryManager manager = new TaskMemoryManager(
@@ -58,4 +84,46 @@ public void encodePageNumberAndOffsetOnHeap() throws IOException {
     Assert.assertEquals(64, manager.getOffsetInPage(encodedAddress));
   }
 
+  @Test
+  public void cooperativeSpilling() throws IOException {
+    final TestMemoryManager memoryManager = new TestMemoryManager(new SparkConf());
+    memoryManager.limit(100);
+    final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0);
+
+    TestMemoryConsumer c1 = new TestMemoryConsumer(manager);
+    TestMemoryConsumer c2 = new TestMemoryConsumer(manager);
+    c1.use(100);
+    assert(c1.used == 100);
+    c2.use(100);
+    assert(c2.used == 100);
+    assert(c1.used == 0);  // spilled
+    c1.use(100);
+    assert(c1.used == 100);
+    assert(c2.used == 0);  // spilled
+
+    c1.use(50);
+    assert(c1.used == 50);  // spilled
+    assert(c2.used == 0);
+    c2.use(50);
+    assert(c1.used == 50);
+    assert(c2.used == 50);
+
+    c1.use(100);
+    assert(c1.used == 100);
+    assert(c2.used == 0);  // spilled
+
+    c1.free(20);
+    assert(c1.used == 80);
+    c2.use(10);
+    assert(c1.used == 80);
+    assert(c2.used == 10);
+    c2.use(100);
+    assert(c2.used == 100);
+    assert(c1.used == 0);  // spilled
+
+    c1.free(0);
+    c2.free(100);
+    assert(manager.cleanUpAllAllocatedMemory() == 0);
+  }
+
 }
diff --git a/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala b/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala
index 2bb4452dde74c..77e43554ee27c 100644
--- a/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala
+++ b/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala
@@ -35,7 +35,9 @@ class TestMemoryManager(conf: SparkConf) extends MemoryManager(conf, numCores =
       numBytes
     } else {
       _executionMemoryUsed += available
-      available
+      val grant = available
+      available = 0
+      grant
     }
   }
   override def acquireStorageMemory(
@@ -64,4 +66,5 @@ class TestMemoryManager(conf: SparkConf) extends MemoryManager(conf, numCores =
   def limit(avail: Long): Unit = {
     available = avail
   }
+
 }

From a3e01d0c9222bee4bc8eb48e2542a4344e72aa33 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Wed, 28 Oct 2015 22:54:41 -0700
Subject: [PATCH 08/17] counting the longArray in BytesToBytesMap

---
 .../apache/spark/memory/MemoryConsumer.java   | 14 ++---
 .../spark/memory/TaskMemoryManager.java       | 59 +++++++++++++------
 .../shuffle/sort/ShuffleExternalSorter.java   |  6 +-
 .../spark/unsafe/map/BytesToBytesMap.java     | 23 +++++++-
 .../unsafe/sort/UnsafeExternalSorter.java     | 41 ++++++-------
 .../unsafe/sort/UnsafeInMemorySorter.java     |  2 +-
 .../apache/spark/memory/MemoryManager.scala   | 11 +++-
 .../spark/memory/TaskMemoryManagerSuite.java  | 12 ++--
 .../sort/UnsafeShuffleWriterSuite.java        |  2 +-
 .../map/AbstractBytesToBytesMapSuite.java     |  2 +-
 .../sort/UnsafeExternalSorterSuite.java       |  6 +-
 .../UnsafeFixedWidthAggregationMapSuite.scala |  2 +-
 12 files changed, 111 insertions(+), 69 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
index 7f8cf90883572..b051e79ec4133 100644
--- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
+++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
@@ -67,15 +67,13 @@ public void spill() throws IOException {
   /**
    * Acquire `size` bytes memory.
    *
-   * If there is not enough memory, throws IOException.
-   *
-   * @throws IOException
+   * If there is not enough memory, throws OutOfMemoryError.
    */
-  protected void acquireMemory(long size) throws IOException {
+  protected void acquireMemory(long size) {
     long got = memoryManager.acquireExecutionMemory(size, this);
     if (got < size) {
       memoryManager.showMemoryUsage();
-      throw new IOException("Could not acquire " + size + " bytes of memory, got " + got);
+      throw new OutOfMemoryError("Could not acquire " + size + " bytes of memory, got " + got);
     }
   }
 
@@ -91,9 +89,9 @@ protected void releaseMemory(long size) {
    *
    * Throws IOException if there is not enough memory.
    *
-   * @throws IOException
+   * @throws OutOfMemoryError
    */
-  protected MemoryBlock allocatePage(long required) throws IOException {
+  protected MemoryBlock allocatePage(long required) {
     MemoryBlock page = memoryManager.allocatePage(Math.max(pageSize, required), this);
     if (page == null || page.size() < required) {
       long got = 0;
@@ -102,7 +100,7 @@ protected MemoryBlock allocatePage(long required) throws IOException {
         freePage(page);
       }
       memoryManager.showMemoryUsage();
-      throw new IOException("Unable to acquire " + required + " bytes of memory, got " + got);
+      throw new OutOfMemoryError("Unable to acquire " + required + " bytes of memory, got " + got);
     }
     return page;
   }
diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 14553f4c35272..8c1ac7ec2f3f2 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -18,6 +18,7 @@
 package org.apache.spark.memory;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.BitSet;
 import java.util.HashMap;
 
@@ -126,7 +127,7 @@ public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) {
    *
    * @return number of bytes successfully granted (<= N).
    */
-  public long acquireExecutionMemory(long required, MemoryConsumer consumer) throws IOException {
+  public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
     assert(required >= 0);
     synchronized (this) {
       long got = memoryManager.acquireExecutionMemory(required, taskAttemptId);
@@ -138,13 +139,19 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) throw
         // Call spill() on other consumers to release memory
         for (MemoryConsumer c: cs) {
           if (c != null) {
-            long released = c.spill(required - got, consumer);
-            logger.info("released " + Utils.bytesToString(released) + " from " + consumer);
-            if (released >= 0) {
-              got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId);
-              if (got >= required) {
-                break;
+            try {
+              long released = c.spill(required - got, consumer);
+              if (released > 0) {
+                logger.info("released {} from {} for {}", Utils.bytesToString(released), c, consumer);
+                got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId);
+                if (got >= required) {
+                  break;
+                }
               }
+            } catch (IOException e) {
+              logger.error("error while calling spill() on " + c, e);
+              throw new OutOfMemoryError("error while calling spill() on " + c + " : "
+                + e.getMessage());
             }
           }
         }
@@ -205,6 +212,26 @@ public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
     memoryManager.releaseExecutionMemory(size, taskAttemptId);
   }
 
+  public void transferOwnership(long size, MemoryConsumer from, MemoryConsumer to) {
+    assert(size >= 0);
+    synchronized (this) {
+      if (consumers.containsKey(from)) {
+        long old = consumers.get(from);
+        if (old > size) {
+          consumers.put(from, old - size);
+        } else {
+          consumers.remove(from);
+        }
+        if (consumers.containsKey(to)) {
+          old = consumers.get(to);
+        } else {
+          old = 0L;
+        }
+        consumers.put(to, old + size);
+      }
+    }
+  }
+
   /**
    * Dump the memory usage of all consumers.
    */
@@ -230,7 +257,7 @@ public long pageSizeBytes() {
    *
    * Returns `null` if there was not enough memory to allocate the page.
    */
-  public MemoryBlock allocatePage(long size, MemoryConsumer consumer) throws IOException {
+  public MemoryBlock allocatePage(long size, MemoryConsumer consumer) {
     if (size > MAXIMUM_PAGE_SIZE_BYTES) {
       throw new IllegalArgumentException(
         "Cannot allocate a page with more than " + MAXIMUM_PAGE_SIZE_BYTES + " bytes");
@@ -355,18 +382,14 @@ public long getOffsetInPage(long pagePlusOffsetAddress) {
    * value can be used to detect memory leaks.
    */
   public long cleanUpAllAllocatedMemory() {
-    long freedBytes = 0;
-    for (MemoryBlock page : pageTable) {
-      if (page != null) {
-        freedBytes += page.size();
-        freePage(page, null);
+    synchronized (this) {
+      Arrays.fill(pageTable, null);
+      for (MemoryConsumer c: consumers.keySet()) {
+        logger.warn("leak " + Utils.bytesToString(consumers.get(c)) + " memory from " + c);
       }
+      consumers.clear();
     }
-    consumers.clear();
-
-    freedBytes += memoryManager.releaseAllExecutionMemoryForTask(taskAttemptId);
-
-    return freedBytes;
+    return memoryManager.releaseAllExecutionMemoryForTask(taskAttemptId);
   }
 
   /**
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
index d5696979036d8..5631686ada060 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
@@ -102,7 +102,7 @@ public ShuffleExternalSorter(
       int initialSize,
       int numPartitions,
       SparkConf conf,
-      ShuffleWriteMetrics writeMetrics) throws IOException {
+      ShuffleWriteMetrics writeMetrics) {
     super(memoryManager);
     this.taskMemoryManager = memoryManager;
     this.blockManager = blockManager;
@@ -316,7 +316,7 @@ public void cleanupResources() {
    * array and grows the array if additional space is required. If the required space cannot be
    * obtained, then the in-memory data will be spilled to disk.
    */
-  private void growPointerArrayIfNecessary() throws IOException {
+  private void growPointerArrayIfNecessary() {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
       logger.debug("Attempting to expand sort pointer array");
@@ -341,7 +341,7 @@ private void growPointerArrayIfNecessary() throws IOException {
    *                      that exceed the page size are handled via a different code path which uses
    *                      special overflow pages).
    */
-  private void acquireNewPageIfNecessary(int required) throws IOException {
+  private void acquireNewPageIfNecessary(int required) {
     if (currentPage == null ||
       pageCursor + required > currentPage.getBaseOffset() + currentPage.size() ) {
       // TODO: try to find space in previous pages
diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 5cbf234f102e5..27fdc22e03496 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -199,6 +199,7 @@ public BytesToBytesMap(
         TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES);
     }
     allocate(initialCapacity);
+    acquireMemory(longArray.memoryBlock().size());
   }
 
   public BytesToBytesMap(
@@ -662,6 +663,16 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
           return false;
         }
       }
+      boolean toGrow = numElements > growthThreshold && longArray.size() < MAX_CAPACITY;
+      if (toGrow) {
+        long capacity = Math.min(MAX_CAPACITY,
+          ByteArrayMethods.nextPowerOf2(growthStrategy.nextCapacity((int) bitset.capacity())));
+        try {
+          acquireMemory(capacity * 16L);
+        } catch (OutOfMemoryError e) {
+          return false;
+        }
+      }
 
       // --- Append the key and value data to the current data page --------------------------------
       final Object base = currentPage.getBaseObject();
@@ -686,8 +697,11 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
       longArray.set(pos * 2 + 1, keyHashcode);
       updateAddressesAndSizes(storedKeyAddress);
       isDefined = true;
-      if (numElements > growthThreshold && longArray.size() < MAX_CAPACITY) {
+
+      if (toGrow) {
+        long usedMemory = longArray.memoryBlock().size();
         growAndRehash();
+        releaseMemory(usedMemory);
       }
       return true;
     }
@@ -704,7 +718,7 @@ private boolean acquireNewPage(long required) {
       Platform.putInt(currentPage.getBaseObject(), currentPage.getBaseOffset(), 0);
       pageCursor = 4;
       return true;
-    } catch (IOException e) {
+    } catch (OutOfMemoryError e) {
       return false;
     }
   }
@@ -743,7 +757,10 @@ private void allocate(int capacity) {
    */
   public void free() {
     updatePeakMemoryUsed();
-    longArray = null;
+    if (longArray != null) {
+      releaseMemory(longArray.memoryBlock().size());
+      longArray = null;
+    }
     bitset = null;
     Iterator dataPagesIterator = dataPages.iterator();
     while (dataPagesIterator.hasNext()) {
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index af4af97c87a00..6b9ae2eda061c 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -67,8 +67,6 @@ public final class UnsafeExternalSorter extends MemoryConsumer {
 
   // These variables are reset after spilling:
   @Nullable private volatile UnsafeInMemorySorter inMemSorter;
-  // The acquired memory for in-memory sorter
-  private long acquiredMem = 0L;
 
   private MemoryBlock currentPage = null;
   private long pageCursor = -1;
@@ -83,7 +81,7 @@ public static UnsafeExternalSorter createWithExistingInMemorySorter(
       PrefixComparator prefixComparator,
       int initialSize,
       long pageSizeBytes,
-      UnsafeInMemorySorter inMemorySorter) throws IOException {
+      UnsafeInMemorySorter inMemorySorter) {
     return new UnsafeExternalSorter(taskMemoryManager, blockManager,
       taskContext, recordComparator, prefixComparator, initialSize, pageSizeBytes, inMemorySorter);
   }
@@ -95,7 +93,7 @@ public static UnsafeExternalSorter create(
       RecordComparator recordComparator,
       PrefixComparator prefixComparator,
       int initialSize,
-      long pageSizeBytes) throws IOException {
+      long pageSizeBytes) {
     return new UnsafeExternalSorter(taskMemoryManager, blockManager,
       taskContext, recordComparator, prefixComparator, initialSize, pageSizeBytes, null);
   }
@@ -108,7 +106,7 @@ private UnsafeExternalSorter(
       PrefixComparator prefixComparator,
       int initialSize,
       long pageSizeBytes,
-      @Nullable UnsafeInMemorySorter existingInMemorySorter) throws IOException {
+      @Nullable UnsafeInMemorySorter existingInMemorySorter) {
     super(taskMemoryManager, pageSizeBytes);
     this.taskMemoryManager = taskMemoryManager;
     this.blockManager = blockManager;
@@ -125,12 +123,10 @@ private UnsafeExternalSorter(
     if (existingInMemorySorter == null) {
       this.inMemSorter =
         new UnsafeInMemorySorter(taskMemoryManager, recordComparator, prefixComparator, initialSize);
-      acquireMemory(inMemSorter.getMemoryUsage());
-      acquiredMem = inMemSorter.getMemoryUsage();
     } else {
-      acquiredMem = 0;
       this.inMemSorter = existingInMemorySorter;
     }
+    acquireMemory(inMemSorter.getMemoryUsage());
 
     // Register a cleanup task with TaskContext to ensure that memory is guaranteed to be freed at
     // the end of the task. This is necessary to avoid memory leaks in when the downstream operator
@@ -277,9 +273,9 @@ public void cleanupResources() {
       deleteSpillFiles();
       freeMemory();
       if (inMemSorter != null) {
+        long used = inMemSorter.getMemoryUsage();
         inMemSorter = null;
-        releaseMemory(acquiredMem);
-        acquiredMem = 0;
+        releaseMemory(used);
       }
     }
   }
@@ -289,7 +285,7 @@ public void cleanupResources() {
    * array and grows the array if additional space is required. If the required space cannot be
    * obtained, then the in-memory data will be spilled to disk.
    */
-  private void growPointerArrayIfNecessary() throws IOException {
+  private void growPointerArrayIfNecessary() {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
       long used = inMemSorter.getMemoryUsage();
@@ -300,7 +296,6 @@ private void growPointerArrayIfNecessary() throws IOException {
       } else {
         inMemSorter.expandPointerArray();
         releaseMemory(used);
-        acquiredMem += needed;
       }
     }
   }
@@ -314,7 +309,7 @@ private void growPointerArrayIfNecessary() throws IOException {
    *                      that exceed the page size are handled via a different code path which uses
    *                      special overflow pages).
    */
-  private void acquireNewPageIfNecessary(int required) throws IOException {
+  private void acquireNewPageIfNecessary(int required) {
     if (currentPage == null ||
       pageCursor + required > currentPage.getBaseOffset() + currentPage.size()) {
       // TODO: try to find space on previous pages
@@ -327,8 +322,7 @@ private void acquireNewPageIfNecessary(int required) throws IOException {
   /**
    * Write a record to the sorter.
    */
-  public void insertRecord(Object recordBase, long recordOffset, int length, long prefix)
-    throws IOException {
+  public void insertRecord(Object recordBase, long recordOffset, int length, long prefix) {
 
     growPointerArrayIfNecessary();
     // Need 4 bytes to store the record length.
@@ -354,7 +348,7 @@ public void insertRecord(Object recordBase, long recordOffset, int length, long
    * record length = key length + value length + 4
    */
   public void insertKVRecord(Object keyBase, long keyOffset, int keyLen,
-      Object valueBase, long valueOffset, int valueLen, long prefix) throws IOException {
+      Object valueBase, long valueOffset, int valueLen, long prefix) {
 
     growPointerArrayIfNecessary();
     final int required = keyLen + valueLen + 4 + 4;
@@ -399,21 +393,24 @@ public UnsafeSorterIterator getSortedIterator() throws IOException {
   }
 
   /**
-   * A
+   * An UnsafeSorterIterator that support spilling.
    */
   class SpillableIterator extends UnsafeSorterIterator {
-    private volatile UnsafeSorterIterator upstream;
+    private UnsafeSorterIterator upstream;
     private UnsafeSorterIterator nextUpstream = null;
     private MemoryBlock lastPage = null;
     private boolean loaded = false;
+    private int numRecords = 0;
 
     public SpillableIterator(UnsafeInMemorySorter.SortedIterator inMemIterator) {
       this.upstream = inMemIterator;
+      this.numRecords = inMemIterator.numRecordsLeft();
     }
 
     public long spill() throws IOException {
       synchronized (this) {
-        if (!(upstream instanceof UnsafeInMemorySorter.SortedIterator && upstream.hasNext())) {
+        if (!(upstream instanceof UnsafeInMemorySorter.SortedIterator && nextUpstream == null
+          && numRecords > 0)) {
           return 0L;
         }
 
@@ -421,8 +418,7 @@ public long spill() throws IOException {
           ((UnsafeInMemorySorter.SortedIterator) upstream).clone();
 
         final UnsafeSorterSpillWriter spillWriter =
-          new UnsafeSorterSpillWriter(blockManager, fileBufferSizeBytes, writeMetrics,
-            inMemIterator.numRecordsLeft());
+          new UnsafeSorterSpillWriter(blockManager, fileBufferSizeBytes, writeMetrics, numRecords);
         while (inMemIterator.hasNext()) {
           inMemIterator.loadNext();
           final Object baseObject = inMemIterator.getBaseObject();
@@ -457,7 +453,7 @@ public long spill() throws IOException {
 
     @Override
     public boolean hasNext() {
-      return upstream.hasNext();
+      return numRecords > 0;
     }
 
     @Override
@@ -473,6 +469,7 @@ public void loadNext() throws IOException {
           upstream = nextUpstream;
           nextUpstream = null;
         }
+        numRecords--;
         upstream.loadNext();
       }
     }
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
index f9a6d49398744..1480f0681ed9c 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
@@ -175,7 +175,7 @@ public boolean hasNext() {
     }
 
     public int numRecordsLeft() {
-      return sortBufferInsertPosition - position;
+      return (sortBufferInsertPosition - position) / 2;
     }
 
     @Override
diff --git a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala
index e37797b55654a..b0cf2696a397f 100644
--- a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala
+++ b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala
@@ -24,6 +24,7 @@ import scala.collection.mutable.ArrayBuffer
 
 import com.google.common.annotations.VisibleForTesting
 
+import org.apache.spark.util.Utils
 import org.apache.spark.{SparkException, TaskContext, SparkConf, Logging}
 import org.apache.spark.storage.{BlockId, BlockStatus, MemoryStore}
 import org.apache.spark.unsafe.array.ByteArrayMethods
@@ -135,7 +136,7 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) exte
    * that control global sharing of memory between execution and storage.
    */
   private[memory]
-  def acquireExecutionMemory(numBytes: Long, taskAttemptId: Long): Long = synchronized {
+  final def acquireExecutionMemory(numBytes: Long, taskAttemptId: Long): Long = synchronized {
     assert(numBytes > 0, "invalid number of bytes requested: " + numBytes)
 
     // Add this task to the taskMemory map just so we can keep an accurate count of the number
@@ -215,8 +216,12 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) exte
   final def releaseExecutionMemory(numBytes: Long, taskAttemptId: Long): Unit = synchronized {
     val curMem = executionMemoryForTask.getOrElse(taskAttemptId, 0L)
     if (curMem < numBytes) {
-      throw new SparkException(
-        s"Internal error: release called on $numBytes bytes but task only has $curMem")
+      if (Utils.isTesting) {
+        throw new SparkException(
+          s"Internal error: release called on $numBytes bytes but task only has $curMem")
+      } else {
+        logWarning(s"Internal error: release called on $numBytes bytes but task only has $curMem")
+      }
     }
     if (executionMemoryForTask.contains(taskAttemptId)) {
       executionMemoryForTask(taskAttemptId) -= numBytes
diff --git a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
index 3e7351b37a8c5..3c360ba14a31f 100644
--- a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
+++ b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
@@ -42,19 +42,19 @@ public long spill(long size, MemoryConsumer trigger) throws IOException {
       return released;
     }
 
-    void use(long size) throws IOException {
+    void use(long size) {
       acquireMemory(size);
       used += size;
     }
 
-    void free(long size) throws IOException {
+    void free(long size) {
       releaseMemory(size);
       used -= size;
     }
   }
 
   @Test
-  public void leakedPageMemoryIsDetected() throws IOException {
+  public void leakedPageMemoryIsDetected() {
     final TaskMemoryManager manager = new TaskMemoryManager(
       new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
     manager.allocatePage(4096, null);  // leak memory
@@ -62,7 +62,7 @@ public void leakedPageMemoryIsDetected() throws IOException {
   }
 
   @Test
-  public void encodePageNumberAndOffsetOffHeap() throws IOException {
+  public void encodePageNumberAndOffsetOffHeap() {
     final TaskMemoryManager manager = new TaskMemoryManager(
       new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "true")), 0);
     final MemoryBlock dataPage = manager.allocatePage(256, null);
@@ -75,7 +75,7 @@ public void encodePageNumberAndOffsetOffHeap() throws IOException {
   }
 
   @Test
-  public void encodePageNumberAndOffsetOnHeap() throws IOException {
+  public void encodePageNumberAndOffsetOnHeap() {
     final TaskMemoryManager manager = new TaskMemoryManager(
       new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
     final MemoryBlock dataPage = manager.allocatePage(256, null);
@@ -85,7 +85,7 @@ public void encodePageNumberAndOffsetOnHeap() throws IOException {
   }
 
   @Test
-  public void cooperativeSpilling() throws IOException {
+  public void cooperativeSpilling() {
     final TestMemoryManager memoryManager = new TestMemoryManager(new SparkConf());
     memoryManager.limit(100);
     final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0);
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
index a19e345c417b1..4763395d7d401 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
@@ -423,7 +423,7 @@ public void writeEnoughRecordsToTriggerSortBufferExpansionAndSpill() throws Exce
     memoryManager.limit(UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE * 16);
     final UnsafeShuffleWriter writer = createWriter(false);
     final ArrayList> dataToWrite = new ArrayList<>();
-    for (int i = 0; i < UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE * 2; i++) {
+    for (int i = 0; i < UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE; i++) {
       dataToWrite.add(new Tuple2(i, i));
     }
     writer.write(dataToWrite.iterator());
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index 6cdedcfcb9a8d..c44c34d924dc9 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -492,7 +492,7 @@ public void randomizedTestWithRecordsLargerThanPageSize() {
 
   @Test
   public void failureToAllocateFirstPage() {
-    memoryManager.limit(0);
+    memoryManager.limit(1024);  // longArray
     BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, 1, PAGE_SIZE_BYTES);
     try {
       final long[] emptyArray = new long[0];
diff --git a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
index dc654998e2950..cfead0e5924b8 100644
--- a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
+++ b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
@@ -332,7 +332,8 @@ public void forcedSpillingWithReadIterator() throws Exception {
       assert(Platform.getLong(iter.getBaseObject(), iter.getBaseOffset()) == i);
       lastv = i;
     }
-    iter.spill();
+    assert(iter.spill() > 0);
+    assert(iter.spill() == 0);
     assert(Platform.getLong(iter.getBaseObject(), iter.getBaseOffset()) == lastv);
     for (int i = n / 3; i < n; i++) {
       iter.hasNext();
@@ -356,7 +357,8 @@ public void forcedSpillingWithNotReadIterator() throws Exception {
     assert(sorter.getNumberOfAllocatedPages() >= 2);
     UnsafeExternalSorter.SpillableIterator iter =
       (UnsafeExternalSorter.SpillableIterator) sorter.getSortedIterator();
-    iter.spill();
+    assert(iter.spill() > 0);
+    assert(iter.spill() == 0);
     for (int i = 0; i < n; i++) {
       iter.hasNext();
       iter.loadNext();
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
index e36d9599d2fe4..ae4e5b314ec26 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
@@ -343,7 +343,7 @@ class UnsafeFixedWidthAggregationMapSuite
     }
     // Simulate running out of space
     memoryManager.limit(0)
-    val str = rand.nextString(4096)
+    val str = rand.nextString(1024)
     val buf = map.getAggregationBuffer(InternalRow(UTF8String.fromString(str)))
     assert(buf == null)
 

From 7bf76e5a29b49c97fd2d2e22ba244e2dba74df51 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Wed, 28 Oct 2015 23:52:35 -0700
Subject: [PATCH 09/17] preserve the memory in BytsToBytesMap

---
 .../org/apache/spark/unsafe/map/BytesToBytesMap.java   | 10 ++++++----
 .../collection/unsafe/sort/UnsafeExternalSorter.java   |  3 ++-
 .../spark/unsafe/map/AbstractBytesToBytesMapSuite.java |  2 +-
 .../spark/sql/execution/UnsafeKVExternalSorter.java    |  7 ++-----
 4 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 27fdc22e03496..db69e52352274 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -199,7 +199,8 @@ public BytesToBytesMap(
         TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES);
     }
     allocate(initialCapacity);
-    acquireMemory(longArray.memoryBlock().size());
+    // preserve the same amount of memory for UnsafeInMemorySorter
+    acquireMemory(longArray.memoryBlock().size() * 2);
   }
 
   public BytesToBytesMap(
@@ -668,7 +669,8 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
         long capacity = Math.min(MAX_CAPACITY,
           ByteArrayMethods.nextPowerOf2(growthStrategy.nextCapacity((int) bitset.capacity())));
         try {
-          acquireMemory(capacity * 16L);
+          // preserve the same amount of memory for UnsafeInMemorySorter
+          acquireMemory(capacity * 16L * 2);
         } catch (OutOfMemoryError e) {
           return false;
         }
@@ -701,7 +703,7 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
       if (toGrow) {
         long usedMemory = longArray.memoryBlock().size();
         growAndRehash();
-        releaseMemory(usedMemory);
+        releaseMemory(usedMemory * 2);
       }
       return true;
     }
@@ -758,7 +760,7 @@ private void allocate(int capacity) {
   public void free() {
     updatePeakMemoryUsed();
     if (longArray != null) {
-      releaseMemory(longArray.memoryBlock().size());
+      releaseMemory(longArray.memoryBlock().size() * 2);
       longArray = null;
     }
     bitset = null;
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index 6b9ae2eda061c..1c03f5a33f3b7 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -123,10 +123,11 @@ private UnsafeExternalSorter(
     if (existingInMemorySorter == null) {
       this.inMemSorter =
         new UnsafeInMemorySorter(taskMemoryManager, recordComparator, prefixComparator, initialSize);
+      acquireMemory(inMemSorter.getMemoryUsage());
     } else {
       this.inMemSorter = existingInMemorySorter;
+      // will acquire after free the map
     }
-    acquireMemory(inMemSorter.getMemoryUsage());
 
     // Register a cleanup task with TaskContext to ensure that memory is guaranteed to be freed at
     // the end of the task. This is necessary to avoid memory leaks in when the downstream operator
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index c44c34d924dc9..f46a69b640ce9 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -492,7 +492,7 @@ public void randomizedTestWithRecordsLargerThanPageSize() {
 
   @Test
   public void failureToAllocateFirstPage() {
-    memoryManager.limit(1024);  // longArray
+    memoryManager.limit(2048);  // longArray
     BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, 1, PAGE_SIZE_BYTES);
     try {
       final long[] emptyArray = new long[0];
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
index 4bfaf60257488..8e277472e9100 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
@@ -83,11 +83,7 @@ public UnsafeKVExternalSorter(
         /* initialSize */ 4096,
         pageSizeBytes);
     } else {
-      // Insert the records into the in-memory sorter.
-      // We will use the number of elements in the map as the initialSize of the
-      // UnsafeInMemorySorter. Because UnsafeInMemorySorter does not accept 0 as the initialSize,
-      // we will use 1 as its initial size if the map is empty.
-      // TODO: track pointer array memory used by this in-memory sorter! (SPARK-10474)
+      // The memory have been preserved in BytesToBytesMap
       final UnsafeInMemorySorter inMemSorter = new UnsafeInMemorySorter(
         taskMemoryManager, recordComparator, prefixComparator, Math.max(1, map.numElements()));
 
@@ -127,6 +123,7 @@ public UnsafeKVExternalSorter(
 
       sorter.spill();
       map.free();
+      taskContext.taskMemoryManager().acquireExecutionMemory(inMemSorter.getMemoryUsage(), sorter);
     }
   }
 

From 51278f89d73d46f8cf22377c3bd0b7aa3ad5c02e Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Thu, 29 Oct 2015 10:15:50 -0700
Subject: [PATCH 10/17] use the memory of longArray for in-memory sorter

---
 .../spark/unsafe/map/BytesToBytesMap.java     | 28 +++++++++++--------
 .../map/AbstractBytesToBytesMapSuite.java     |  2 +-
 .../sql/execution/UnsafeKVExternalSorter.java |  9 ++++--
 .../UnsafeFixedWidthAggregationMapSuite.scala |  8 ------
 4 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index db69e52352274..5fc925885eec4 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -199,8 +199,7 @@ public BytesToBytesMap(
         TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES);
     }
     allocate(initialCapacity);
-    // preserve the same amount of memory for UnsafeInMemorySorter
-    acquireMemory(longArray.memoryBlock().size() * 2);
+    acquireMemory(longArray.memoryBlock().size());
   }
 
   public BytesToBytesMap(
@@ -669,8 +668,7 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
         long capacity = Math.min(MAX_CAPACITY,
           ByteArrayMethods.nextPowerOf2(growthStrategy.nextCapacity((int) bitset.capacity())));
         try {
-          // preserve the same amount of memory for UnsafeInMemorySorter
-          acquireMemory(capacity * 16L * 2);
+          acquireMemory(capacity * 16L);
         } catch (OutOfMemoryError e) {
           return false;
         }
@@ -703,7 +701,7 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
       if (toGrow) {
         long usedMemory = longArray.memoryBlock().size();
         growAndRehash();
-        releaseMemory(usedMemory * 2);
+        releaseMemory(usedMemory);
       }
       return true;
     }
@@ -752,18 +750,26 @@ private void allocate(int capacity) {
   }
 
   /**
-   * Free all allocated memory associated with this map, including the storage for keys and values
-   * as well as the hash map array itself.
-   *
-   * This method is idempotent and can be called multiple times.
+   * Free the memory used by longArray.
    */
-  public void free() {
+  public void freeArray() {
     updatePeakMemoryUsed();
     if (longArray != null) {
-      releaseMemory(longArray.memoryBlock().size() * 2);
+      long used = longArray.memoryBlock().size();
       longArray = null;
+      releaseMemory(used);
     }
     bitset = null;
+  }
+
+  /**
+   * Free all allocated memory associated with this map, including the storage for keys and values
+   * as well as the hash map array itself.
+   *
+   * This method is idempotent and can be called multiple times.
+   */
+  public void free() {
+    freeArray();
     Iterator dataPagesIterator = dataPages.iterator();
     while (dataPagesIterator.hasNext()) {
       MemoryBlock dataPage = dataPagesIterator.next();
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index f46a69b640ce9..c44c34d924dc9 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -492,7 +492,7 @@ public void randomizedTestWithRecordsLargerThanPageSize() {
 
   @Test
   public void failureToAllocateFirstPage() {
-    memoryManager.limit(2048);  // longArray
+    memoryManager.limit(1024);  // longArray
     BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, 1, PAGE_SIZE_BYTES);
     try {
       final long[] emptyArray = new long[0];
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
index 8e277472e9100..845f2ae6859b7 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
@@ -83,7 +83,9 @@ public UnsafeKVExternalSorter(
         /* initialSize */ 4096,
         pageSizeBytes);
     } else {
-      // The memory have been preserved in BytesToBytesMap
+      // The memory needed for UnsafeInMemorySorter should be less than longArray in map.
+      map.freeArray();
+      // The memory used by UnsafeInMemorySorter will be counted later (end of this block)
       final UnsafeInMemorySorter inMemSorter = new UnsafeInMemorySorter(
         taskMemoryManager, recordComparator, prefixComparator, Math.max(1, map.numElements()));
 
@@ -112,7 +114,7 @@ public UnsafeKVExternalSorter(
       }
 
       sorter = UnsafeExternalSorter.createWithExistingInMemorySorter(
-        taskContext.taskMemoryManager(),
+        taskMemoryManager,
         blockManager,
         taskContext,
         new KVComparator(ordering, keySchema.length()),
@@ -123,7 +125,8 @@ public UnsafeKVExternalSorter(
 
       sorter.spill();
       map.free();
-      taskContext.taskMemoryManager().acquireExecutionMemory(inMemSorter.getMemoryUsage(), sorter);
+      // counting the memory used UnsafeInMemorySorter
+      taskMemoryManager.acquireExecutionMemory(inMemSorter.getMemoryUsage(), sorter);
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
index ae4e5b314ec26..a38623623a441 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
@@ -193,10 +193,6 @@ class UnsafeFixedWidthAggregationMapSuite
     // Convert the map into a sorter
     val sorter = map.destructAndCreateExternalSorter()
 
-    withClue(s"destructAndCreateExternalSorter should release memory used by the map") {
-      assert(taskMemoryManager.getMemoryConsumptionForThisTask() === initialMemoryConsumption)
-    }
-
     // Add more keys to the sorter and make sure the results come out sorted.
     val additionalKeys = randomStrings(1024)
     val keyConverter = UnsafeProjection.create(groupKeySchema)
@@ -294,10 +290,6 @@ class UnsafeFixedWidthAggregationMapSuite
     // Convert the map into a sorter. Right now, it contains one record.
     val sorter = map.destructAndCreateExternalSorter()
 
-    withClue(s"destructAndCreateExternalSorter should release memory used by the map") {
-      assert(taskMemoryManager.getMemoryConsumptionForThisTask() === initialMemoryConsumption)
-    }
-
     // Add more keys to the sorter and make sure the results come out sorted.
     (1 to 4096).foreach { i =>
       sorter.insertKV(UnsafeRow.createFromByteArray(0, 0), UnsafeRow.createFromByteArray(0, 0))

From 27ff4fc859b40afdad379068a6ac0d06f9c09eef Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Thu, 29 Oct 2015 11:54:04 -0700
Subject: [PATCH 11/17] improve spilling

---
 .../apache/spark/memory/MemoryConsumer.java   |  1 +
 .../shuffle/sort/ShuffleExternalSorter.java   | 22 ++++++++----
 .../spark/unsafe/map/BytesToBytesMap.java     | 34 +++++++++---------
 .../unsafe/sort/UnsafeExternalSorter.java     | 36 ++++++++++++-------
 4 files changed, 55 insertions(+), 38 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
index b051e79ec4133..17a469cc15b37 100644
--- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
+++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
@@ -72,6 +72,7 @@ public void spill() throws IOException {
   protected void acquireMemory(long size) {
     long got = memoryManager.acquireExecutionMemory(size, this);
     if (got < size) {
+      memoryManager.releaseExecutionMemory(got, this);
       memoryManager.showMemoryUsage();
       throw new OutOfMemoryError("Could not acquire " + size + " bytes of memory, got " + got);
     }
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
index 5631686ada060..90972d99f435a 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
@@ -316,18 +316,26 @@ public void cleanupResources() {
    * array and grows the array if additional space is required. If the required space cannot be
    * obtained, then the in-memory data will be spilled to disk.
    */
-  private void growPointerArrayIfNecessary() {
+  private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
       logger.debug("Attempting to expand sort pointer array");
       long used = inMemSorter.getMemoryUsage();
       long needed = inMemSorter.getMemoryToExpand();
-      acquireMemory(used + needed);  // could trigger spilling
-      if (inMemSorter.hasSpaceForAnotherRecord()) {
-        releaseMemory(used + needed);
-      } else {
-        inMemSorter.expandPointerArray();
-        releaseMemory(used);
+      try {
+        acquireMemory(used + needed);  // could trigger spilling
+        if (inMemSorter.hasSpaceForAnotherRecord()) {
+          releaseMemory(used + needed);
+        } else {
+          logger.debug("Expand sort pointer array");
+          inMemSorter.expandPointerArray();
+          releaseMemory(used);
+        }
+      } catch (OutOfMemoryError oom) {
+        // spilling should be triggered
+        if (!inMemSorter.hasSpaceForAnotherRecord()) {
+          spill();  // Just in case that JVM had run out of memory
+        }
       }
     }
   }
diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 5fc925885eec4..9d102338ddc19 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -122,6 +122,11 @@ public final class BytesToBytesMap extends MemoryConsumer {
   // full base addresses in the page table for off-heap mode so that we can reconstruct the full
   // absolute memory addresses.
 
+  /**
+   * Whether or not the longArray can grow. We will not insert more elements if it's false.
+   */
+  private boolean canGrowArray = true;
+
   /**
    * A {@link BitSet} used to track location of the map where the key is set.
    * Size of the bitset should be half of the size of the long array.
@@ -199,7 +204,6 @@ public BytesToBytesMap(
         TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES);
     }
     allocate(initialCapacity);
-    acquireMemory(longArray.memoryBlock().size());
   }
 
   public BytesToBytesMap(
@@ -649,8 +653,8 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
       assert(bitset != null);
       assert(longArray != null);
 
-      if (numElements == MAX_CAPACITY) {
-        throw new IllegalStateException("BytesToBytesMap has reached maximum capacity");
+      if (numElements == MAX_CAPACITY || !canGrowArray) {
+        return false;
       }
 
       // Here, we'll copy the data into our data pages. Because we only store a relative offset from
@@ -663,16 +667,6 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
           return false;
         }
       }
-      boolean toGrow = numElements > growthThreshold && longArray.size() < MAX_CAPACITY;
-      if (toGrow) {
-        long capacity = Math.min(MAX_CAPACITY,
-          ByteArrayMethods.nextPowerOf2(growthStrategy.nextCapacity((int) bitset.capacity())));
-        try {
-          acquireMemory(capacity * 16L);
-        } catch (OutOfMemoryError e) {
-          return false;
-        }
-      }
 
       // --- Append the key and value data to the current data page --------------------------------
       final Object base = currentPage.getBaseObject();
@@ -698,10 +692,12 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
       updateAddressesAndSizes(storedKeyAddress);
       isDefined = true;
 
-      if (toGrow) {
-        long usedMemory = longArray.memoryBlock().size();
-        growAndRehash();
-        releaseMemory(usedMemory);
+      if (numElements > growthThreshold && longArray.size() < MAX_CAPACITY) {
+        try {
+          growAndRehash();
+        } catch (OutOfMemoryError oom) {
+          canGrowArray = false;
+        }
       }
       return true;
     }
@@ -742,6 +738,7 @@ private void allocate(int capacity) {
     // The capacity needs to be divisible by 64 so that our bit set can be sized properly
     capacity = Math.max((int) Math.min(MAX_CAPACITY, ByteArrayMethods.nextPowerOf2(capacity)), 64);
     assert (capacity <= MAX_CAPACITY);
+    acquireMemory(capacity * 16);
     longArray = new LongArray(MemoryBlock.fromLongArray(new long[capacity * 2]));
     bitset = new BitSet(MemoryBlock.fromLongArray(new long[capacity / 64]));
 
@@ -758,8 +755,8 @@ public void freeArray() {
       long used = longArray.memoryBlock().size();
       longArray = null;
       releaseMemory(used);
+      bitset = null;
     }
-    bitset = null;
   }
 
   /**
@@ -889,6 +886,7 @@ void growAndRehash() {
         }
       }
     }
+    releaseMemory(oldLongArray.memoryBlock().size());
 
     if (enablePerfMetrics) {
       timeSpentResizingNs += System.nanoTime() - resizeStartTime;
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index 1c03f5a33f3b7..50c1c990eb0c2 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -286,17 +286,24 @@ public void cleanupResources() {
    * array and grows the array if additional space is required. If the required space cannot be
    * obtained, then the in-memory data will be spilled to disk.
    */
-  private void growPointerArrayIfNecessary() {
+  private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
       long used = inMemSorter.getMemoryUsage();
       long needed = inMemSorter.getMemoryToExpand();
-      acquireMemory(used + needed);  // could trigger spilling
-      if (inMemSorter.hasSpaceForAnotherRecord()) {
-        releaseMemory(used + needed);
-      } else {
-        inMemSorter.expandPointerArray();
-        releaseMemory(used);
+      try {
+        acquireMemory(used + needed);  // could trigger spilling
+        if (inMemSorter.hasSpaceForAnotherRecord()) {
+          releaseMemory(used + needed);
+        } else {
+          inMemSorter.expandPointerArray();
+          releaseMemory(used);
+        }
+      } catch (OutOfMemoryError oom) {
+        // spilling should be triggered
+        if (!inMemSorter.hasSpaceForAnotherRecord()) {
+          spill();  // Just in case that JVM had run out of memory
+        }
       }
     }
   }
@@ -323,7 +330,8 @@ private void acquireNewPageIfNecessary(int required) {
   /**
    * Write a record to the sorter.
    */
-  public void insertRecord(Object recordBase, long recordOffset, int length, long prefix) {
+  public void insertRecord(Object recordBase, long recordOffset, int length, long prefix)
+    throws IOException {
 
     growPointerArrayIfNecessary();
     // Need 4 bytes to store the record length.
@@ -349,7 +357,8 @@ public void insertRecord(Object recordBase, long recordOffset, int length, long
    * record length = key length + value length + 4
    */
   public void insertKVRecord(Object keyBase, long keyOffset, int keyLen,
-      Object valueBase, long valueOffset, int valueLen, long prefix) {
+      Object valueBase, long valueOffset, int valueLen, long prefix)
+    throws IOException {
 
     growPointerArrayIfNecessary();
     final int required = keyLen + valueLen + 4 + 4;
@@ -443,10 +452,6 @@ public long spill() throws IOException {
             }
           }
           allocatedPages.clear();
-          assert(inMemSorter != null);
-          long used = inMemSorter.getMemoryUsage();
-          inMemSorter = null;
-          releaseMemory(used);
         }
         return released;
       }
@@ -469,6 +474,11 @@ public void loadNext() throws IOException {
           }
           upstream = nextUpstream;
           nextUpstream = null;
+
+          assert(inMemSorter != null);
+          long used = inMemSorter.getMemoryUsage();
+          inMemSorter = null;
+          releaseMemory(used);
         }
         numRecords--;
         upstream.loadNext();

From 4491013c9b669d7b03fbe16c7d68f62f28175f7e Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Thu, 29 Oct 2015 12:10:44 -0700
Subject: [PATCH 12/17] release from other consumers first

---
 .../spark/memory/TaskMemoryManager.java       | 33 ++++++++++++++-----
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 8c1ac7ec2f3f2..f1291333e06bb 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -132,17 +132,20 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
     synchronized (this) {
       long got = memoryManager.acquireExecutionMemory(required, taskAttemptId);
 
+      // try to release memory from other consumers first, then we can reduce the frequency of
+      // spilling, avoid to have too many spilled files.
       if (got < required) {
         // consumers could be modified by spill(), so we should have a copy here.
         MemoryConsumer[] cs = new MemoryConsumer[consumers.size()];
         consumers.keySet().toArray(cs);
         // Call spill() on other consumers to release memory
         for (MemoryConsumer c: cs) {
-          if (c != null) {
+          if (c != null && c != consumer) {
             try {
               long released = c.spill(required - got, consumer);
               if (released > 0) {
-                logger.info("released {} from {} for {}", Utils.bytesToString(released), c, consumer);
+                logger.info("Task {} released {} from {} for {}", taskAttemptId,
+                  Utils.bytesToString(released), c, consumer);
                 got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId);
                 if (got >= required) {
                   break;
@@ -157,6 +160,22 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
         }
       }
 
+      // call spill() on itself
+      if (got < required) {
+        try {
+          long released = consumer.spill(required - got, consumer);
+          if (released > 0) {
+            logger.info("Task {} released {} from itself ({})", taskAttemptId,
+              Utils.bytesToString(released), consumer);
+            got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId);
+          }
+        } catch (IOException e) {
+          logger.error("error while calling spill() on " + consumer, e);
+          throw new OutOfMemoryError("error while calling spill() on " + consumer + " : "
+            + e.getMessage());
+        }
+      }
+
       // Update the accounting, even consumer is null
       if (got > 0) {
         long old = 0L;
@@ -166,9 +185,7 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
         consumers.put(consumer, got + old);
       }
 
-      if (logger.isTraceEnabled()) {
-        logger.trace("Acquire {} for {}", Utils.bytesToString(got), consumer);
-      }
+      logger.debug("Task {} acquire {} for {}", taskAttemptId, Utils.bytesToString(got), consumer);
       return got;
     }
   }
@@ -206,9 +223,7 @@ public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
       }
     }
 
-    if (logger.isTraceEnabled()) {
-      logger.trace("Release {} from {}", Utils.bytesToString(size), consumer);
-    }
+    logger.debug("Task {} release {} from {}", taskAttemptId, Utils.bytesToString(size), consumer);
     memoryManager.releaseExecutionMemory(size, taskAttemptId);
   }
 
@@ -239,7 +254,7 @@ public void showMemoryUsage() {
     logger.info("Memory used in task " + taskAttemptId);
     synchronized (this) {
       for (MemoryConsumer c: consumers.keySet()) {
-        logger.info("acquired by " + c + ": " + Utils.bytesToString(consumers.get(c)));
+        logger.info("Acquired by " + c + ": " + Utils.bytesToString(consumers.get(c)));
       }
     }
   }

From c044afe62a36d98934d4e9f36665e2a4837783c7 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Thu, 29 Oct 2015 13:14:43 -0700
Subject: [PATCH 13/17] fix npe

---
 .../main/java/org/apache/spark/memory/TaskMemoryManager.java    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index f1291333e06bb..73ec258b1e7df 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -161,7 +161,7 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
       }
 
       // call spill() on itself
-      if (got < required) {
+      if (got < required && consumer != null) {
         try {
           long released = consumer.spill(required - got, consumer);
           if (released > 0) {

From afc8c7c9b0e92f9db6a5f72c14d8484a87311a51 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Thu, 29 Oct 2015 15:27:44 -0700
Subject: [PATCH 14/17] address comments

---
 .../apache/spark/memory/MemoryConsumer.java   | 42 ++++++---
 .../spark/memory/TaskMemoryManager.java       | 87 ++++---------------
 .../shuffle/sort/ShuffleExternalSorter.java   |  6 +-
 .../unsafe/sort/UnsafeExternalSorter.java     |  1 +
 .../spark/memory/TaskMemoryManagerSuite.java  | 41 ++++-----
 5 files changed, 67 insertions(+), 110 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
index 17a469cc15b37..5992ad178318b 100644
--- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
+++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
@@ -28,16 +28,28 @@
  */
 public abstract class MemoryConsumer {
 
-  private TaskMemoryManager memoryManager;
-  private long pageSize;
+  private final TaskMemoryManager taskMemoryManager;
+  private final long pageSize;
+  private long used;
 
-  protected MemoryConsumer(TaskMemoryManager memoryManager, long pageSize) {
-    this.memoryManager = memoryManager;
+  protected MemoryConsumer(TaskMemoryManager taskMemoryManager, long pageSize) {
+    this.taskMemoryManager = taskMemoryManager;
+    if (pageSize == 0) {
+      pageSize = taskMemoryManager.pageSizeBytes();
+    }
     this.pageSize = pageSize;
+    this.used = 0;
+  }
+
+  protected MemoryConsumer(TaskMemoryManager taskMemoryManager) {
+    this(taskMemoryManager, taskMemoryManager.pageSizeBytes());
   }
 
-  protected MemoryConsumer(TaskMemoryManager memoryManager) {
-    this(memoryManager, memoryManager.pageSizeBytes());
+  /**
+   * Returns the size of used memory in bytes.
+   */
+  long getUsed() {
+    return used;
   }
 
   /**
@@ -70,19 +82,21 @@ public void spill() throws IOException {
    * If there is not enough memory, throws OutOfMemoryError.
    */
   protected void acquireMemory(long size) {
-    long got = memoryManager.acquireExecutionMemory(size, this);
+    long got = taskMemoryManager.acquireExecutionMemory(size, this);
     if (got < size) {
-      memoryManager.releaseExecutionMemory(got, this);
-      memoryManager.showMemoryUsage();
+      taskMemoryManager.releaseExecutionMemory(got, this);
+      taskMemoryManager.showMemoryUsage();
       throw new OutOfMemoryError("Could not acquire " + size + " bytes of memory, got " + got);
     }
+    used += got;
   }
 
   /**
    * Release `size` bytes memory.
    */
   protected void releaseMemory(long size) {
-    memoryManager.releaseExecutionMemory(size, this);
+    taskMemoryManager.releaseExecutionMemory(size, this);
+    used -= size;
   }
 
   /**
@@ -93,16 +107,17 @@ protected void releaseMemory(long size) {
    * @throws OutOfMemoryError
    */
   protected MemoryBlock allocatePage(long required) {
-    MemoryBlock page = memoryManager.allocatePage(Math.max(pageSize, required), this);
+    MemoryBlock page = taskMemoryManager.allocatePage(Math.max(pageSize, required), this);
     if (page == null || page.size() < required) {
       long got = 0;
       if (page != null) {
         got = page.size();
         freePage(page);
       }
-      memoryManager.showMemoryUsage();
+      taskMemoryManager.showMemoryUsage();
       throw new OutOfMemoryError("Unable to acquire " + required + " bytes of memory, got " + got);
     }
+    used += page.size();
     return page;
   }
 
@@ -110,6 +125,7 @@ protected MemoryBlock allocatePage(long required) {
    * Free a memory block.
    */
   protected void freePage(MemoryBlock page) {
-    memoryManager.freePage(page, this);
+    taskMemoryManager.freePage(page, this);
+    used -= page.size();
   }
 }
diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 73ec258b1e7df..45aecacaa23db 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -20,14 +20,12 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.BitSet;
-import java.util.HashMap;
+import java.util.HashSet;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import org.apache.spark.SparkException;
-import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.util.Utils;
 
@@ -109,7 +107,7 @@ public class TaskMemoryManager {
   /**
    * The size of memory granted to each consumer.
    */
-  private final HashMap consumers;
+  private final HashSet consumers;
 
   /**
    * Construct a new TaskMemoryManager.
@@ -118,7 +116,7 @@ public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) {
     this.inHeap = memoryManager.tungstenMemoryIsAllocatedInHeap();
     this.memoryManager = memoryManager;
     this.taskAttemptId = taskAttemptId;
-    this.consumers = new HashMap<>();
+    this.consumers = new HashSet<>();
   }
 
   /**
@@ -135,12 +133,9 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
       // try to release memory from other consumers first, then we can reduce the frequency of
       // spilling, avoid to have too many spilled files.
       if (got < required) {
-        // consumers could be modified by spill(), so we should have a copy here.
-        MemoryConsumer[] cs = new MemoryConsumer[consumers.size()];
-        consumers.keySet().toArray(cs);
         // Call spill() on other consumers to release memory
-        for (MemoryConsumer c: cs) {
-          if (c != null && c != consumer) {
+        for (MemoryConsumer c: consumers) {
+          if (c != null && c != consumer && c.getUsed() > 0) {
             try {
               long released = c.spill(required - got, consumer);
               if (released > 0) {
@@ -176,15 +171,7 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
         }
       }
 
-      // Update the accounting, even consumer is null
-      if (got > 0) {
-        long old = 0L;
-        if (consumers.containsKey(consumer)) {
-          old = consumers.get(consumer);
-        }
-        consumers.put(consumer, got + old);
-      }
-
+      consumers.add(consumer);
       logger.debug("Task {} acquire {} for {}", taskAttemptId, Utils.bytesToString(got), consumer);
       return got;
     }
@@ -194,67 +181,20 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
    * Release N bytes of execution memory for a MemoryConsumer.
    */
   public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
-    assert(size >= 0);
-    if (size == 0) {
-      return;
-    }
-    synchronized (this) {
-      if (consumers.containsKey(consumer)) {
-        long old = consumers.get(consumer);
-        if (old > size) {
-          consumers.put(consumer, old - size);
-        } else {
-          if (old < size) {
-            String msg = "Release " + size + " bytes memory (more than acquired " + old + ") for "
-              + consumer;
-            logger.warn(msg);
-            if (Utils.isTesting()) {
-              Platform.throwException(new SparkException(msg));
-            }
-          }
-          consumers.remove(consumer);
-        }
-      } else {
-        String msg = "Release " + size + " bytes memory for non-existent " + consumer;
-        logger.warn(msg);
-        if (Utils.isTesting()) {
-          Platform.throwException(new SparkException(msg));
-        }
-      }
-    }
-
     logger.debug("Task {} release {} from {}", taskAttemptId, Utils.bytesToString(size), consumer);
     memoryManager.releaseExecutionMemory(size, taskAttemptId);
   }
 
-  public void transferOwnership(long size, MemoryConsumer from, MemoryConsumer to) {
-    assert(size >= 0);
-    synchronized (this) {
-      if (consumers.containsKey(from)) {
-        long old = consumers.get(from);
-        if (old > size) {
-          consumers.put(from, old - size);
-        } else {
-          consumers.remove(from);
-        }
-        if (consumers.containsKey(to)) {
-          old = consumers.get(to);
-        } else {
-          old = 0L;
-        }
-        consumers.put(to, old + size);
-      }
-    }
-  }
-
   /**
    * Dump the memory usage of all consumers.
    */
   public void showMemoryUsage() {
     logger.info("Memory used in task " + taskAttemptId);
     synchronized (this) {
-      for (MemoryConsumer c: consumers.keySet()) {
-        logger.info("Acquired by " + c + ": " + Utils.bytesToString(consumers.get(c)));
+      for (MemoryConsumer c: consumers) {
+        if (c.getUsed() > 0) {
+          logger.info("Acquired by " + c + ": " + Utils.bytesToString(c.getUsed()));
+        }
       }
     }
   }
@@ -399,8 +339,11 @@ public long getOffsetInPage(long pagePlusOffsetAddress) {
   public long cleanUpAllAllocatedMemory() {
     synchronized (this) {
       Arrays.fill(pageTable, null);
-      for (MemoryConsumer c: consumers.keySet()) {
-        logger.warn("leak " + Utils.bytesToString(consumers.get(c)) + " memory from " + c);
+      for (MemoryConsumer c: consumers) {
+        if (c != null && c.getUsed() > 0) {
+          // In case of failed task, it's normal to see leaked memory
+          logger.warn("leak " + Utils.bytesToString(c.getUsed()) + " memory from " + c);
+        }
       }
       consumers.clear();
     }
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
index 90972d99f435a..4e433d97d7ed8 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
@@ -103,11 +103,11 @@ public ShuffleExternalSorter(
       int numPartitions,
       SparkConf conf,
       ShuffleWriteMetrics writeMetrics) {
-    super(memoryManager);
+    super(memoryManager, (int) Math.min(PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES,
+      memoryManager.pageSizeBytes()));
     this.taskMemoryManager = memoryManager;
     this.blockManager = blockManager;
     this.taskContext = taskContext;
-    this.peakMemoryUsedBytes = initialSize;
     this.numPartitions = numPartitions;
     // Use getSizeAsKb (not bytes) to maintain backwards compatibility if no units are provided
     this.fileBufferSizeBytes = (int) conf.getSizeAsKb("spark.shuffle.file.buffer", "32k") * 1024;
@@ -116,6 +116,7 @@ public ShuffleExternalSorter(
     this.writeMetrics = writeMetrics;
     acquireMemory(initialSize * 8L);
     this.inMemSorter = new ShuffleInMemorySorter(initialSize);
+    this.peakMemoryUsedBytes = getMemoryUsage();
   }
 
   /**
@@ -372,6 +373,7 @@ public void insertRecord(Object recordBase, long recordOffset, int length, int p
     }
 
     growPointerArrayIfNecessary();
+    // Need 4 bytes to store the record length.
     final int required = length + 4;
     acquireNewPageIfNecessary(required);
 
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index 50c1c990eb0c2..6c18039b19168 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -128,6 +128,7 @@ private UnsafeExternalSorter(
       this.inMemSorter = existingInMemorySorter;
       // will acquire after free the map
     }
+    this.peakMemoryUsedBytes = getMemoryUsage();
 
     // Register a cleanup task with TaskContext to ensure that memory is guaranteed to be freed at
     // the end of the task. This is necessary to avoid memory leaks in when the downstream operator
diff --git a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
index 3c360ba14a31f..dab7b0592cb4e 100644
--- a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
+++ b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
@@ -28,28 +28,23 @@
 public class TaskMemoryManagerSuite {
 
   class TestMemoryConsumer extends MemoryConsumer {
-    volatile long used = 0L;
-
     TestMemoryConsumer(TaskMemoryManager memoryManager) {
       super(memoryManager);
     }
 
     @Override
     public long spill(long size, MemoryConsumer trigger) throws IOException {
+      long used = getUsed();
       releaseMemory(used);
-      long released = used;
-      used = 0;
-      return released;
+      return used;
     }
 
     void use(long size) {
       acquireMemory(size);
-      used += size;
     }
 
     void free(long size) {
       releaseMemory(size);
-      used -= size;
     }
   }
 
@@ -93,33 +88,33 @@ public void cooperativeSpilling() {
     TestMemoryConsumer c1 = new TestMemoryConsumer(manager);
     TestMemoryConsumer c2 = new TestMemoryConsumer(manager);
     c1.use(100);
-    assert(c1.used == 100);
+    assert(c1.getUsed() == 100);
     c2.use(100);
-    assert(c2.used == 100);
-    assert(c1.used == 0);  // spilled
+    assert(c2.getUsed() == 100);
+    assert(c1.getUsed() == 0);  // spilled
     c1.use(100);
-    assert(c1.used == 100);
-    assert(c2.used == 0);  // spilled
+    assert(c1.getUsed() == 100);
+    assert(c2.getUsed() == 0);  // spilled
 
     c1.use(50);
-    assert(c1.used == 50);  // spilled
-    assert(c2.used == 0);
+    assert(c1.getUsed() == 50);  // spilled
+    assert(c2.getUsed() == 0);
     c2.use(50);
-    assert(c1.used == 50);
-    assert(c2.used == 50);
+    assert(c1.getUsed() == 50);
+    assert(c2.getUsed() == 50);
 
     c1.use(100);
-    assert(c1.used == 100);
-    assert(c2.used == 0);  // spilled
+    assert(c1.getUsed() == 100);
+    assert(c2.getUsed() == 0);  // spilled
 
     c1.free(20);
-    assert(c1.used == 80);
+    assert(c1.getUsed() == 80);
     c2.use(10);
-    assert(c1.used == 80);
-    assert(c2.used == 10);
+    assert(c1.getUsed() == 80);
+    assert(c2.getUsed() == 10);
     c2.use(100);
-    assert(c2.used == 100);
-    assert(c1.used == 0);  // spilled
+    assert(c2.getUsed() == 100);
+    assert(c1.getUsed() == 0);  // spilled
 
     c1.free(0);
     c2.free(100);

From cda4b2afa45acef529eb24afa7ec0c5370de237f Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Thu, 29 Oct 2015 16:52:12 -0700
Subject: [PATCH 15/17] address more comments

---
 .../apache/spark/memory/MemoryConsumer.java   |  7 ++--
 .../spark/memory/TaskMemoryManager.java       |  2 ++
 .../spark/unsafe/map/BytesToBytesMap.java     | 34 +++++++++----------
 .../unsafe/sort/UnsafeExternalSorter.java     | 19 +++++------
 4 files changed, 30 insertions(+), 32 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
index 5992ad178318b..008799cc77395 100644
--- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
+++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
@@ -34,9 +34,6 @@ public abstract class MemoryConsumer {
 
   protected MemoryConsumer(TaskMemoryManager taskMemoryManager, long pageSize) {
     this.taskMemoryManager = taskMemoryManager;
-    if (pageSize == 0) {
-      pageSize = taskMemoryManager.pageSizeBytes();
-    }
     this.pageSize = pageSize;
     this.used = 0;
   }
@@ -95,8 +92,8 @@ protected void acquireMemory(long size) {
    * Release `size` bytes memory.
    */
   protected void releaseMemory(long size) {
-    taskMemoryManager.releaseExecutionMemory(size, this);
     used -= size;
+    taskMemoryManager.releaseExecutionMemory(size, this);
   }
 
   /**
@@ -125,7 +122,7 @@ protected MemoryBlock allocatePage(long required) {
    * Free a memory block.
    */
   protected void freePage(MemoryBlock page) {
-    taskMemoryManager.freePage(page, this);
     used -= page.size();
+    taskMemoryManager.freePage(page, this);
   }
 }
diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 45aecacaa23db..4230575446d31 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -17,6 +17,7 @@
 
 package org.apache.spark.memory;
 
+import javax.annotation.concurrent.GuardedBy;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.BitSet;
@@ -107,6 +108,7 @@ public class TaskMemoryManager {
   /**
    * The size of memory granted to each consumer.
    */
+  @GuardedBy("this")
   private final HashSet consumers;
 
   /**
diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 9d102338ddc19..d38010417c11b 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -23,9 +23,6 @@
 import java.util.Iterator;
 import java.util.LinkedList;
 
-import scala.runtime.AbstractFunction0;
-import scala.runtime.BoxedUnit;
-
 import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -43,6 +40,7 @@
 import org.apache.spark.unsafe.hash.Murmur3_x86_32;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.unsafe.memory.MemoryLocation;
+import org.apache.spark.util.TaskCompletionListener;
 import org.apache.spark.util.collection.unsafe.sort.UnsafeSorterSpillReader;
 import org.apache.spark.util.collection.unsafe.sort.UnsafeSorterSpillWriter;
 
@@ -349,33 +347,35 @@ public long spill(long numBytes) throws IOException {
           offset += 4;
           final UnsafeSorterSpillWriter writer =
             new UnsafeSorterSpillWriter(blockManager, 32 * 1024, writeMetrics, numRecords);
-          while (numRecords-- > 0) {
+          while (numRecords > 0) {
             int length = Platform.getInt(base, offset);
             writer.write(base, offset + 4, length, 0);
             offset += 4 + length;
+            numRecords--;
           }
           writer.close();
           spillWriters.add(writer);
           if (TaskContext.get() != null) {
-            TaskContext.get().addOnCompleteCallback(new AbstractFunction0() {
-              @Override
-              public BoxedUnit apply() {
-                File file = writer.getFile();
-                if (file != null && file.exists()) {
-                  if (!file.delete()) {
-                    logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
+            TaskContext.get().addTaskCompletionListener(
+              new TaskCompletionListener() {
+                @Override
+                public void onTaskCompletion(TaskContext context) {
+                  File file = writer.getFile();
+                  if (file != null && file.exists()) {
+                    if (!file.delete()) {
+                      logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
+                    }
                   }
                 }
-                return null;
               }
-            });
+            );
           }
 
           dataPages.removeLast();
-          freePage(block);
           released += block.size();
+          freePage(block);
 
-          if (released > numBytes) {
+          if (released >= numBytes) {
             break;
           }
         }
@@ -545,8 +545,8 @@ private Location with(MemoryBlock page, long offsetInPage) {
     }
 
     /**
-     +     * This is only used for spilling
-     +     */
+     * This is only used for spilling
+     */
     private Location with(Object base, long offset, int length) {
       this.isDefined = true;
       this.memoryPage = null;
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index 6c18039b19168..183da6d1a4bd5 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -22,9 +22,6 @@
 import java.io.IOException;
 import java.util.LinkedList;
 
-import scala.runtime.AbstractFunction0;
-import scala.runtime.BoxedUnit;
-
 import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -36,6 +33,7 @@
 import org.apache.spark.storage.BlockManager;
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.memory.MemoryBlock;
+import org.apache.spark.util.TaskCompletionListener;
 import org.apache.spark.util.Utils;
 
 /**
@@ -133,13 +131,14 @@ private UnsafeExternalSorter(
     // Register a cleanup task with TaskContext to ensure that memory is guaranteed to be freed at
     // the end of the task. This is necessary to avoid memory leaks in when the downstream operator
     // does not fully consume the sorter's output (e.g. sort followed by limit).
-    taskContext.addOnCompleteCallback(new AbstractFunction0() {
-      @Override
-      public BoxedUnit apply() {
-        cleanupResources();
-        return null;
+    TaskContext.get().addTaskCompletionListener(
+      new TaskCompletionListener() {
+        @Override
+        public void onTaskCompletion(TaskContext context) {
+          cleanupResources();
+        }
       }
-    });
+    );
   }
 
   /**
@@ -244,8 +243,8 @@ private long freeMemory() {
     updatePeakMemoryUsed();
     long memoryFreed = 0;
     for (MemoryBlock block : allocatedPages) {
-      freePage(block);
       memoryFreed += block.size();
+      freePage(block);
     }
     allocatedPages.clear();
     currentPage = null;

From 4ee1f42b08b46469d8631a56b01898d251ebf3e7 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Thu, 29 Oct 2015 17:57:57 -0700
Subject: [PATCH 16/17] improve OOM handling

---
 .../shuffle/sort/ShuffleExternalSorter.java   | 28 +++++++++++--------
 .../shuffle/sort/ShuffleInMemorySorter.java   |  3 +-
 .../spark/unsafe/map/BytesToBytesMap.java     | 16 +++++++----
 .../unsafe/sort/UnsafeExternalSorter.java     | 26 ++++++++++-------
 4 files changed, 44 insertions(+), 29 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
index 4e433d97d7ed8..400d8520019b9 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
@@ -320,22 +320,26 @@ public void cleanupResources() {
   private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
-      logger.debug("Attempting to expand sort pointer array");
       long used = inMemSorter.getMemoryUsage();
-      long needed = inMemSorter.getMemoryToExpand();
+      long needed = used + inMemSorter.getMemoryToExpand();
       try {
-        acquireMemory(used + needed);  // could trigger spilling
-        if (inMemSorter.hasSpaceForAnotherRecord()) {
-          releaseMemory(used + needed);
-        } else {
-          logger.debug("Expand sort pointer array");
+        acquireMemory(needed);  // could trigger spilling
+      } catch (OutOfMemoryError e) {
+        // should have trigger spilling
+        assert(inMemSorter.hasSpaceForAnotherRecord());
+        return;
+      }
+      // check if spilling is triggered or not
+      if (inMemSorter.hasSpaceForAnotherRecord()) {
+        releaseMemory(needed);
+      } else {
+        try {
           inMemSorter.expandPointerArray();
           releaseMemory(used);
-        }
-      } catch (OutOfMemoryError oom) {
-        // spilling should be triggered
-        if (!inMemSorter.hasSpaceForAnotherRecord()) {
-          spill();  // Just in case that JVM had run out of memory
+        } catch (OutOfMemoryError oom) {
+          // Just in case that JVM had run out of memory
+          releaseMemory(needed);
+          spill();
         }
       }
     }
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
index 2b49737fed413..e630575d1ae19 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
@@ -60,8 +60,7 @@ public void reset() {
 
   private int newLength() {
     // Guard against overflow:
-    return array.length <= Integer.MAX_VALUE / 2 ?
-      (array.length * 2) : Integer.MAX_VALUE;
+    return array.length <= Integer.MAX_VALUE / 2 ? (array.length * 2) : Integer.MAX_VALUE;
   }
 
   /**
diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index d38010417c11b..1f7107e5cb26a 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -710,13 +710,13 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
   private boolean acquireNewPage(long required) {
     try {
       currentPage = allocatePage(required);
-      dataPages.add(currentPage);
-      Platform.putInt(currentPage.getBaseObject(), currentPage.getBaseOffset(), 0);
-      pageCursor = 4;
-      return true;
     } catch (OutOfMemoryError e) {
       return false;
     }
+    dataPages.add(currentPage);
+    Platform.putInt(currentPage.getBaseObject(), currentPage.getBaseOffset(), 0);
+    pageCursor = 4;
+    return true;
   }
 
   @Override
@@ -862,7 +862,13 @@ void growAndRehash() {
     final int oldCapacity = (int) oldBitSet.capacity();
 
     // Allocate the new data structures
-    allocate(Math.min(growthStrategy.nextCapacity(oldCapacity), MAX_CAPACITY));
+    try {
+      allocate(Math.min(growthStrategy.nextCapacity(oldCapacity), MAX_CAPACITY));
+    } catch (OutOfMemoryError oom) {
+      longArray = oldLongArray;
+      bitset = oldBitSet;
+      throw oom;
+    }
 
     // Re-mask (we don't recompute the hashcode because we stored all 32 bits of it)
     for (int pos = oldBitSet.nextSetBit(0); pos >= 0; pos = oldBitSet.nextSetBit(pos + 1)) {
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index 183da6d1a4bd5..fdc2182db0a8e 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -290,19 +290,25 @@ private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
       long used = inMemSorter.getMemoryUsage();
-      long needed = inMemSorter.getMemoryToExpand();
+      long needed = used + inMemSorter.getMemoryToExpand();
       try {
-        acquireMemory(used + needed);  // could trigger spilling
-        if (inMemSorter.hasSpaceForAnotherRecord()) {
-          releaseMemory(used + needed);
-        } else {
+        acquireMemory(needed);  // could trigger spilling
+      } catch (OutOfMemoryError e) {
+        // should have trigger spilling
+        assert(inMemSorter.hasSpaceForAnotherRecord());
+        return;
+      }
+      // check if spilling is triggered or not
+      if (inMemSorter.hasSpaceForAnotherRecord()) {
+        releaseMemory(needed);
+      } else {
+        try {
           inMemSorter.expandPointerArray();
           releaseMemory(used);
-        }
-      } catch (OutOfMemoryError oom) {
-        // spilling should be triggered
-        if (!inMemSorter.hasSpaceForAnotherRecord()) {
-          spill();  // Just in case that JVM had run out of memory
+        } catch (OutOfMemoryError oom) {
+          // Just in case that JVM had run out of memory
+          releaseMemory(needed);
+          spill();
         }
       }
     }

From e943e745317bd4d13934a14026be5a7915706269 Mon Sep 17 00:00:00 2001
From: Davies Liu 
Date: Thu, 29 Oct 2015 20:13:40 -0700
Subject: [PATCH 17/17] fix tests

---
 .../spark/unsafe/map/BytesToBytesMap.java     | 39 ++++++++++---------
 .../unsafe/sort/UnsafeExternalSorter.java     |  2 +-
 .../map/AbstractBytesToBytesMapSuite.java     |  2 +-
 3 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 1f7107e5cb26a..e36709c6fc849 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -28,7 +28,6 @@
 import org.slf4j.LoggerFactory;
 
 import org.apache.spark.SparkEnv;
-import org.apache.spark.TaskContext;
 import org.apache.spark.executor.ShuffleWriteMetrics;
 import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.memory.TaskMemoryManager;
@@ -40,7 +39,6 @@
 import org.apache.spark.unsafe.hash.Murmur3_x86_32;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.unsafe.memory.MemoryLocation;
-import org.apache.spark.util.TaskCompletionListener;
 import org.apache.spark.util.collection.unsafe.sort.UnsafeSorterSpillReader;
 import org.apache.spark.util.collection.unsafe.sort.UnsafeSorterSpillWriter;
 
@@ -175,6 +173,7 @@ public final class BytesToBytesMap extends MemoryConsumer {
 
   private final BlockManager blockManager;
   private volatile MapIterator destructiveIterator = null;
+  private LinkedList spillWriters = new LinkedList<>();
 
   public BytesToBytesMap(
       TaskMemoryManager taskMemoryManager,
@@ -243,7 +242,6 @@ public final class MapIterator implements Iterator {
     // If this iterator destructive or not. When it is true, it frees each page as it moves onto
     // next one.
     private boolean destructive = false;
-    private LinkedList spillWriters = new LinkedList<>();
     private UnsafeSorterSpillReader reader = null;
 
     private MapIterator(int numRecords, Location loc, boolean destructive) {
@@ -293,6 +291,17 @@ private void advanceToNextPage() {
 
     @Override
     public boolean hasNext() {
+      if (numRecords == 0) {
+        if (reader != null) {
+          // remove the spill file from disk
+          File file = spillWriters.removeFirst().getFile();
+          if (file != null && file.exists()) {
+            if (!file.delete()) {
+              logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
+            }
+          }
+        }
+      }
       return numRecords > 0;
     }
 
@@ -355,21 +364,6 @@ public long spill(long numBytes) throws IOException {
           }
           writer.close();
           spillWriters.add(writer);
-          if (TaskContext.get() != null) {
-            TaskContext.get().addTaskCompletionListener(
-              new TaskCompletionListener() {
-                @Override
-                public void onTaskCompletion(TaskContext context) {
-                  File file = writer.getFile();
-                  if (file != null && file.exists()) {
-                    if (!file.delete()) {
-                      logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
-                    }
-                  }
-                }
-              }
-            );
-          }
 
           dataPages.removeLast();
           released += block.size();
@@ -774,6 +768,15 @@ public void free() {
       freePage(dataPage);
     }
     assert(dataPages.isEmpty());
+
+    while (!spillWriters.isEmpty()) {
+      File file = spillWriters.removeFirst().getFile();
+      if (file != null && file.exists()) {
+        if (!file.delete()) {
+          logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
+        }
+      }
+    }
   }
 
   public TaskMemoryManager getTaskMemoryManager() {
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index fdc2182db0a8e..49a5a4b13b70d 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -131,7 +131,7 @@ private UnsafeExternalSorter(
     // Register a cleanup task with TaskContext to ensure that memory is guaranteed to be freed at
     // the end of the task. This is necessary to avoid memory leaks in when the downstream operator
     // does not fully consume the sorter's output (e.g. sort followed by limit).
-    TaskContext.get().addTaskCompletionListener(
+    taskContext.addTaskCompletionListener(
       new TaskCompletionListener() {
         @Override
         public void onTaskCompletion(TaskContext context) {
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index c44c34d924dc9..92bd45e5fa241 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -556,7 +556,7 @@ public void spillInIterator() throws IOException {
       for (i = 0; i < 100; i++) {
         iter2.next();
       }
-      Assert.assertTrue(iter2.spill(1024) > 1024);
+      Assert.assertTrue(iter2.spill(1024) >= 1024);
       for (i = 100; i < 1024; i++) {
         iter2.next();
       }