fix build

Author: Davies Liu

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);
       }
     }
 

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
   }
 

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));

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();

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);

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<Object, Object> writer = createWriter(false);
     final ArrayList<Product2<Object, Object>> dataToWrite = new ArrayList<Product2<Object, Object>>();
     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<Object, Object> writer = createWriter(false);
     final ArrayList<Product2<Object, Object>> dataToWrite = new ArrayList<>();
     for (int i = 0; i < UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE; i++) {

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();

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()
     }