Copy the memory from off-heap to on-heap and some code style modifica…

…tion

core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleReader.scala

@@ -18,15 +18,15 @@
 package org.apache.spark.shuffle.sort
 
 import java.io.FileOutputStream
+import java.nio.ByteBuffer
 import java.util.Comparator
 
-import org.apache.spark.executor.ShuffleWriteMetrics
-
 import scala.collection.mutable.{ArrayBuffer, HashMap, Queue}
 import scala.util.{Failure, Success, Try}
 
 import org.apache.spark._
-import org.apache.spark.network.buffer.ManagedBuffer
+import org.apache.spark.executor.ShuffleWriteMetrics
+import org.apache.spark.network.buffer.{ManagedBuffer, NioManagedBuffer}
 import org.apache.spark.serializer.Serializer
 import org.apache.spark.shuffle.{BaseShuffleHandle, FetchFailedException, ShuffleReader}
 import org.apache.spark.storage._
@@ -129,15 +129,24 @@ private[spark] class SortShuffleReader[K, C](
       // Try to fit block in memory. If this fails, merge in-memory blocks to disk.
       val blockSize = blockData.size
       val granted = shuffleMemoryManager.tryToAcquire(blockSize)
-      val block = MemoryShuffleBlock(blockId, blockData)
       if (granted >= blockSize) {
-        inMemoryBlocks += MemoryShuffleBlock(blockId, blockData)
+        if (blockData.isDirect) {
+          // If the memory shuffle block is allocated on direct buffer, copy it on heap,
+          // otherwise off heap memory will be increased out of control.
+          val onHeapBuffer = ByteBuffer.allocate(blockSize.toInt)
+          onHeapBuffer.put(blockData.nioByteBuffer)
+
+          inMemoryBlocks += MemoryShuffleBlock(blockId, new NioManagedBuffer(onHeapBuffer))
+          blockData.release()
+        } else {
+          inMemoryBlocks += MemoryShuffleBlock(blockId, blockData)
+        }
       } else {
-        logInfo(s"Granted $granted memory is not enough to store shuffle block id $blockId, " +
+        logDebug(s"Granted $granted memory is not enough to store shuffle block id $blockId, " +
           s"block size $blockSize, spilling in-memory blocks to release the memory")
 
         shuffleMemoryManager.release(granted)
-        spillInMemoryBlocks(block)
+        spillInMemoryBlocks(MemoryShuffleBlock(blockId, blockData))
       }
 
       unfetchedBytes -= shuffleBlockMap(blockId.asInstanceOf[ShuffleBlockId])._2
@@ -164,7 +173,6 @@ private[spark] class SortShuffleReader[K, C](
       }
       inMemoryBlocks.clear()
     }
-
     context.addTaskCompletionListener(_ => releaseFinalShuffleMemory())
 
     // Release the in-memory block when iteration is completed.
@@ -214,7 +222,7 @@ private[spark] class SortShuffleReader[K, C](
       var success = false
 
       try {
-        partialMergedItr.foreach(p => writer.write(p))
+        partialMergedItr.foreach(writer.write)
         success = true
       } finally {
         if (!success) {

core/src/main/scala/org/apache/spark/util/collection/TieredDiskMerger.scala

@@ -61,9 +61,6 @@ private[spark] class TieredDiskMerger[K, C](
   private val blockManager = SparkEnv.get.blockManager
   private val ser = Serializer.getSerializer(dep.serializer)
 
-  /** Number of bytes spilled on disk */
-  private var _diskBytesSpilled: Long = 0L
-
   /** PriorityQueue to store the on-disk merging blocks, blocks are merged by size ordering */
   private val onDiskBlocks = new PriorityBlockingQueue[DiskShuffleBlock]()
 
@@ -77,6 +74,11 @@ private[spark] class TieredDiskMerger[K, C](
 
   @volatile private var doneRegistering = false
 
+ /** Number of bytes spilled on disk */
+  private var _diskBytesSpilled: Long = 0L
+
+  def diskBytesSpilled: Long = _diskBytesSpilled
+
   def registerOnDiskBlock(blockId: BlockId, file: File): Unit = {
     assert(!doneRegistering)
     onDiskBlocks.put(new DiskShuffleBlock(blockId, file, file.length()))
@@ -88,8 +90,6 @@ private[spark] class TieredDiskMerger[K, C](
     }
   }
 
-  def diskBytesSpilled: Long = _diskBytesSpilled
-
   /**
    * Notify the merger that no more on disk blocks will be registered.
    */
@@ -184,9 +184,10 @@ private[spark] class TieredDiskMerger[K, C](
           val blocksToMerge = new ArrayBuffer[DiskShuffleBlock]()
           // Try to pick the smallest merge width that will result in the next merge being the final
           // merge.
-          val mergeFactor = math.min(onDiskBlocks.size - maxMergeFactor + 1, maxMergeFactor)
-          (0 until mergeFactor).foreach {
+          var mergeFactor = math.min(onDiskBlocks.size - maxMergeFactor + 1, maxMergeFactor)
+          while (mergeFactor > 0) {
             blocksToMerge += onDiskBlocks.take()
+            mergeFactor -= 1
           }
 
           // Merge the blocks
@@ -201,7 +202,7 @@ private[spark] class TieredDiskMerger[K, C](
           var success = false
 
           try {
-            partialMergedItr.foreach(p => writer.write(p))
+            partialMergedItr.foreach(writer.write)
             success = true
           } finally {
             if (!success) {

network/common/src/main/java/org/apache/spark/network/buffer/ManagedBuffer.java

@@ -69,9 +69,9 @@ public abstract class ManagedBuffer {
    */
   public abstract Object convertToNetty() throws IOException;
 
-    /**
-     * Tell whether to not this byte buffer is direct
-     * @return
-     */
+  /**
+   * Tell whether to not this byte buffer is direct
+   * @return
+   */
   public abstract boolean isDirect();
 }

network/common/src/test/java/org/apache/spark/network/TestManagedBuffer.java

@@ -101,4 +101,9 @@ public boolean equals(Object other) {
     }
     return false;
   }
+
+  @Override
+  public boolean isDirect() {
+    return underlying.isDirect();
+  }
 }