Another fix to freeing pages

Author: JoshRosen

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

@@ -134,6 +134,6 @@ protected void freePage(MemoryBlock page) {
     } else {
       offHeapMemoryUsed -= page.size();
     }
-    taskMemoryManager.freePage(page);
+    taskMemoryManager.freePage(page, this);
   }
 }

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

@@ -263,7 +263,7 @@ public MemoryBlock allocatePage(long size, MemoryConsumer consumer) {
   /**
    * Free a block of memory allocated via {@link TaskMemoryManager#allocatePage}.
    */
-  public void freePage(MemoryBlock page) {
+  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));
@@ -274,6 +274,9 @@ public void freePage(MemoryBlock page) {
     if (logger.isTraceEnabled()) {
       logger.trace("Freed page number {} ({} bytes)", page.pageNumber, page.size());
     }
+    long pageSize = page.size();
+    memoryManager.tungstenMemoryAllocator().free(page);
+    releaseExecutionMemory(pageSize, tungstenMemoryMode, consumer);
   }
 
   /**