address more comments

Author: Davies Liu

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

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<MemoryConsumer> consumers;
 
   /**

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<BoxedUnit>() {
-              @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;

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<BoxedUnit>() {
-      @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;