Attempt at fixing merge conflicts; refactor to use MemoryMode.

Author: JoshRosen

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

@@ -17,36 +17,38 @@
 
 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 abstract class MemoryConsumer {
 
-  private final TaskMemoryManager taskMemoryManager;
+  protected final TaskMemoryManager taskMemoryManager;
   private final long pageSize;
-  private long used;
+  private long onHeapMemoryUsed = 0L;
+  private long offHeapMemoryUsed = 0L;
 
   protected MemoryConsumer(TaskMemoryManager taskMemoryManager, long pageSize) {
     this.taskMemoryManager = taskMemoryManager;
     this.pageSize = pageSize;
-    this.used = 0;
   }
 
   protected MemoryConsumer(TaskMemoryManager taskMemoryManager) {
     this(taskMemoryManager, taskMemoryManager.pageSizeBytes());
   }
 
   /**
-   * Returns the size of used memory in bytes.
+   * Returns the size of used on-heap memory in bytes.
    */
-  long getUsed() {
-    return used;
+  long getMemoryUsed(MemoryMode mode) {
+    if (mode == MemoryMode.ON_HEAP) {
+      return onHeapMemoryUsed;
+    } else {
+      return offHeapMemoryUsed;
+    }
   }
 
   /**
@@ -71,29 +73,30 @@ public void spill() throws IOException {
    * @return the amount of released memory in bytes
    * @throws IOException
    */
+  // TODO(josh): clarify assumption that this only frees Tungsten-managed pages (for now).
   public abstract long spill(long size, MemoryConsumer trigger) throws IOException;
 
   /**
-   * Acquire `size` bytes memory.
+   * Acquire `size` bytes of on-heap execution memory.
    *
    * If there is not enough memory, throws OutOfMemoryError.
    */
   protected void acquireMemory(long size) {
-    long got = taskMemoryManager.acquireExecutionMemory(size, this);
+    long got = taskMemoryManager.acquireExecutionMemory(size, MemoryMode.ON_HEAP, this);
     if (got < size) {
-      taskMemoryManager.releaseExecutionMemory(got, this);
+      taskMemoryManager.releaseExecutionMemory(got, MemoryMode.ON_HEAP, this);
       taskMemoryManager.showMemoryUsage();
       throw new OutOfMemoryError("Could not acquire " + size + " bytes of memory, got " + got);
     }
-    used += got;
+    onHeapMemoryUsed += got;
   }
 
   /**
    * Release `size` bytes memory.
    */
   protected void releaseMemory(long size) {
-    used -= size;
-    taskMemoryManager.releaseExecutionMemory(size, this);
+    onHeapMemoryUsed -= size;
+    taskMemoryManager.releaseExecutionMemory(size, MemoryMode.ON_HEAP, this);
   }
 
   /**
@@ -114,15 +117,23 @@ protected MemoryBlock allocatePage(long required) {
       taskMemoryManager.showMemoryUsage();
       throw new OutOfMemoryError("Unable to acquire " + required + " bytes of memory, got " + got);
     }
-    used += page.size();
+    if (taskMemoryManager.tungstenMemoryMode == MemoryMode.ON_HEAP) {
+      onHeapMemoryUsed += page.size();
+    } else {
+      offHeapMemoryUsed += page.size();
+    }
     return page;
   }
 
   /**
    * Free a memory block.
    */
   protected void freePage(MemoryBlock page) {
-    used -= page.size();
+    if (taskMemoryManager.tungstenMemoryMode == MemoryMode.ON_HEAP) {
+      onHeapMemoryUsed -= page.size();
+    } else {
+      offHeapMemoryUsed -= page.size();
+    }
     taskMemoryManager.freePage(page, this);
   }
 }

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

@@ -0,0 +1,26 @@
+/*
+ * 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 org.apache.spark.annotation.Private;
+
+@Private
+public enum MemoryMode {
+  ON_HEAP,
+  OFF_HEAP
+}

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

@@ -103,7 +103,7 @@ public class TaskMemoryManager {
    * without doing any masking or lookups. Since this branching should be well-predicted by the JIT,
    * this extra layer of indirection / abstraction hopefully shouldn't be too expensive.
    */
-  private final boolean inHeap;
+  final MemoryMode tungstenMemoryMode;
 
   /**
    * The size of memory granted to each consumer.
@@ -115,7 +115,7 @@ public class TaskMemoryManager {
    * Construct a new TaskMemoryManager.
    */
   public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) {
-    this.inHeap = memoryManager.tungstenMemoryIsAllocatedInHeap();
+    this.tungstenMemoryMode = memoryManager.tungstenMemoryMode();
     this.memoryManager = memoryManager;
     this.taskAttemptId = taskAttemptId;
     this.consumers = new HashSet<>();
@@ -127,23 +127,33 @@ public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) {
    *
    * @return number of bytes successfully granted (<= N).
    */
-  public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
+  public long acquireExecutionMemory(
+      long required,
+      MemoryMode mode,
+      MemoryConsumer consumer) {
     assert(required >= 0);
+    // TODO(josh): handle spill differently based on type of request (on-heap vs off-heap).
+    // If we are allocating tungsten pages off-heap and receive a request to allocate on-heap
+    // memory here, then it may not make sense to spill since that would only end up freeing
+    // off-heap memory. This is subject to change, though, so it may be risky to make this
+    // optimization now in case we forget to undo it late when making changes.
     synchronized (this) {
-      long got = memoryManager.acquireExecutionMemory(required, taskAttemptId);
+      long got = memoryManager.acquireExecutionMemory(required, taskAttemptId, mode);
 
       // 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) {
         // Call spill() on other consumers to release memory
         for (MemoryConsumer c: consumers) {
-          if (c != null && c != consumer && c.getUsed() > 0) {
+          if (c != null && c != consumer && c.getMemoryUsed(mode) > 0) {
             try {
+              // TODO(josh): subtlety / implementation detail: today, spill() happens to only
+              // release Tungsten pages.
               long released = c.spill(required - got, consumer);
-              if (released > 0) {
+              if (released > 0 && mode == tungstenMemoryMode) {
                 logger.info("Task {} released {} from {} for {}", taskAttemptId,
                   Utils.bytesToString(released), c, consumer);
-                got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId);
+                got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId, mode);
                 if (got >= required) {
                   break;
                 }
@@ -161,10 +171,10 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
       if (got < required && consumer != null) {
         try {
           long released = consumer.spill(required - got, consumer);
-          if (released > 0) {
+          if (released > 0 && mode == tungstenMemoryMode) {
             logger.info("Task {} released {} from itself ({})", taskAttemptId,
               Utils.bytesToString(released), consumer);
-            got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId);
+            got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId, mode);
           }
         } catch (IOException e) {
           logger.error("error while calling spill() on " + consumer, e);
@@ -182,9 +192,9 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
   /**
    * Release N bytes of execution memory for a MemoryConsumer.
    */
-  public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
+  public void releaseExecutionMemory(long size, MemoryMode mode, MemoryConsumer consumer) {
     logger.debug("Task {} release {} from {}", taskAttemptId, Utils.bytesToString(size), consumer);
-    memoryManager.releaseExecutionMemory(size, taskAttemptId);
+    memoryManager.releaseExecutionMemory(size, taskAttemptId, mode);
   }
 
   /**
@@ -194,8 +204,10 @@ public void showMemoryUsage() {
     logger.info("Memory used in task " + taskAttemptId);
     synchronized (this) {
       for (MemoryConsumer c: consumers) {
-        if (c.getUsed() > 0) {
-          logger.info("Acquired by " + c + ": " + Utils.bytesToString(c.getUsed()));
+        long totalMemUsage =
+          c.getMemoryUsed(MemoryMode.OFF_HEAP) + c.getMemoryUsed(MemoryMode.ON_HEAP);
+        if (totalMemUsage > 0) {
+          logger.info("Acquired by " + c + ": " + Utils.bytesToString(totalMemUsage));
         }
       }
     }
@@ -212,15 +224,16 @@ public long 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.
+   * Returns `null` if there was not enough memory to allocate the page. May return a page that
+   * contains fewer bytes than requested, so callers should verify the size of returned pages.
    */
   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");
     }
 
-    long acquired = acquireExecutionMemory(size, consumer);
+    long acquired = acquireExecutionMemory(size, tungstenMemoryMode, consumer);
     if (acquired <= 0) {
       return null;
     }
@@ -229,29 +242,12 @@ public MemoryBlock allocatePage(long size, MemoryConsumer consumer) {
     synchronized (this) {
       pageNumber = allocatedPages.nextClearBit(0);
       if (pageNumber >= PAGE_TABLE_SIZE) {
-        releaseExecutionMemory(acquired, consumer);
+        releaseExecutionMemory(acquired, tungstenMemoryMode, consumer);
         throw new IllegalStateException(
           "Have already allocated a maximum of " + PAGE_TABLE_SIZE + " pages");
       }
       allocatedPages.set(pageNumber);
     }
-    final long acquiredExecutionMemory;
-    if (memoryManager.tungstenMemoryIsAllocatedInHeap()) {
-      acquiredExecutionMemory = acquireOnHeapExecutionMemory(size);
-    } else {
-      acquiredExecutionMemory = acquireOffHeapExecutionMemory(size);
-    }
-    if (acquiredExecutionMemory != size) {
-      if (memoryManager.tungstenMemoryIsAllocatedInHeap()) {
-        releaseOnHeapExecutionMemory(acquiredExecutionMemory);
-      } else {
-        releaseOffHeapExecutionMemory(acquiredExecutionMemory);
-      }
-      synchronized (this) {
-        allocatedPages.clear(pageNumber);
-      }
-      return null;
-    }
     final MemoryBlock page = memoryManager.tungstenMemoryAllocator().allocate(size);
     page.pageNumber = pageNumber;
     pageTable[pageNumber] = page;
@@ -275,13 +271,7 @@ public void freePage(MemoryBlock page, MemoryConsumer consumer) {
     if (logger.isTraceEnabled()) {
       logger.trace("Freed page number {} ({} bytes)", page.pageNumber, page.size());
     }
-    long pageSize = page.size();
-    memoryManager.tungstenMemoryAllocator().free(page);
-    if (memoryManager.tungstenMemoryIsAllocatedInHeap()) {
-      releaseOnHeapExecutionMemory(pageSize);
-    } else {
-      releaseOffHeapExecutionMemory(pageSize);
-    }
+    consumer.freePage(page);
   }
 
   /**
@@ -295,7 +285,7 @@ public void freePage(MemoryBlock page, MemoryConsumer consumer) {
    * @return an encoded page address.
    */
   public long encodePageNumberAndOffset(MemoryBlock page, long offsetInPage) {
-    if (!inHeap) {
+    if (tungstenMemoryMode == MemoryMode.OFF_HEAP) {
       // In off-heap mode, an offset is an absolute address that may require a full 64 bits to
       // encode. Due to our page size limitation, though, we can convert this into an offset that's
       // relative to the page's base offset; this relative offset will fit in 51 bits.
@@ -324,7 +314,7 @@ private static long decodeOffset(long pagePlusOffsetAddress) {
    * {@link TaskMemoryManager#encodePageNumberAndOffset(MemoryBlock, long)}
    */
   public Object getPage(long pagePlusOffsetAddress) {
-    if (inHeap) {
+    if (tungstenMemoryMode == MemoryMode.ON_HEAP) {
       final int pageNumber = decodePageNumber(pagePlusOffsetAddress);
       assert (pageNumber >= 0 && pageNumber < PAGE_TABLE_SIZE);
       final MemoryBlock page = pageTable[pageNumber];
@@ -342,7 +332,7 @@ public Object getPage(long pagePlusOffsetAddress) {
    */
   public long getOffsetInPage(long pagePlusOffsetAddress) {
     final long offsetInPage = decodeOffset(pagePlusOffsetAddress);
-    if (inHeap) {
+    if (tungstenMemoryMode == MemoryMode.ON_HEAP) {
       return offsetInPage;
     } else {
       // In off-heap mode, an offset is an absolute address. In encodePageNumberAndOffset, we
@@ -363,9 +353,15 @@ public long cleanUpAllAllocatedMemory() {
     synchronized (this) {
       Arrays.fill(pageTable, null);
       for (MemoryConsumer c: consumers) {
-        if (c != null && c.getUsed() > 0) {
+        if (c != null && c.getMemoryUsed(MemoryMode.ON_HEAP) > 0) {
+          // In case of failed task, it's normal to see leaked memory
+          logger.warn("leak " + Utils.bytesToString(c.getMemoryUsed(MemoryMode.ON_HEAP)) +
+            " of on-heap memory from " + c);
+        }
+        if (c != null && c.getMemoryUsed(MemoryMode.OFF_HEAP) > 0) {
           // In case of failed task, it's normal to see leaked memory
-          logger.warn("leak " + Utils.bytesToString(c.getUsed()) + " memory from " + c);
+          logger.warn("leak " + Utils.bytesToString(c.getMemoryUsed(MemoryMode.OFF_HEAP)) +
+            " of off-heap memory from " + c);
         }
       }
       consumers.clear();