Revert "In memory shuffle (cherry-picked from amplab/graphx#135)"

This reverts commit 5ec645d.

Conflicts:
	graphx/src/main/scala/org/apache/spark/graphx/Pregel.scala

Author: ankurdave

core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala

@@ -29,7 +29,6 @@ import org.apache.spark.rdd.RDDCheckpointData
 import org.apache.spark.serializer.Serializer
 import org.apache.spark.storage._
 import org.apache.spark.util.{MetadataCleaner, MetadataCleanerType, TimeStampedHashMap}
-import java.nio.ByteBuffer
 
 private[spark] object ShuffleMapTask {
 
@@ -169,11 +168,7 @@ private[spark] class ShuffleMapTask(
       var totalBytes = 0L
       var totalTime = 0L
       val compressedSizes: Array[Byte] = shuffle.writers.map { writer: BlockObjectWriter =>
-        // writer.commit()
-        val bytes = writer.commit()
-        if (bytes != null) {
-          blockManager.putBytes(writer.blockId, ByteBuffer.wrap(bytes), StorageLevel.MEMORY_ONLY_SER, tellMaster = false)
-        }
+        writer.commit()
         writer.close()
         val size = writer.fileSegment().length
         totalBytes += size

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -57,7 +57,7 @@ private[spark] class BlockManager(
 
   private val blockInfo = new TimeStampedHashMap[BlockId, BlockInfo]
 
-  private[storage] val memoryStore = new MemoryStore(this, maxMemory)
+  private[storage] val memoryStore: BlockStore = new MemoryStore(this, maxMemory)
   private[storage] val diskStore = new DiskStore(this, diskBlockManager)
 
   // If we use Netty for shuffle, start a new Netty-based shuffle sender service.
@@ -293,7 +293,7 @@ private[spark] class BlockManager(
    * never deletes (recent) items.
    */
   def getLocalFromDisk(blockId: BlockId, serializer: Serializer): Option[Iterator[Any]] = {
-    memoryStore.getValues(blockId, serializer).orElse(
+    diskStore.getValues(blockId, serializer).orElse(
       sys.error("Block " + blockId + " not found on disk, though it should be"))
   }
 
@@ -313,7 +313,7 @@ private[spark] class BlockManager(
     // As an optimization for map output fetches, if the block is for a shuffle, return it
     // without acquiring a lock; the disk store never deletes (recent) items so this should work
     if (blockId.isShuffle) {
-      memoryStore.getBytes(blockId) match {
+      diskStore.getBytes(blockId) match {
         case Some(bytes) =>
           Some(bytes)
         case None =>
@@ -831,7 +831,7 @@ private[spark] class BlockManager(
     if (info != null) info.synchronized {
       // Removals are idempotent in disk store and memory store. At worst, we get a warning.
       val removedFromMemory = memoryStore.remove(blockId)
-      val removedFromDisk = false //diskStore.remove(blockId)
+      val removedFromDisk = diskStore.remove(blockId)
       if (!removedFromMemory && !removedFromDisk) {
         logWarning("Block " + blockId + " could not be removed as it was not found in either " +
           "the disk or memory store")

core/src/main/scala/org/apache/spark/storage/BlockObjectWriter.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark.storage
 
-import java.io.{ByteArrayOutputStream, FileOutputStream, File, OutputStream}
+import java.io.{FileOutputStream, File, OutputStream}
 import java.nio.channels.FileChannel
 
 import it.unimi.dsi.fastutil.io.FastBufferedOutputStream
@@ -44,7 +44,7 @@ private[spark] abstract class BlockObjectWriter(val blockId: BlockId) {
    * Flush the partial writes and commit them as a single atomic block. Return the
    * number of bytes written for this commit.
    */
-  def commit(): Array[Byte]
+  def commit(): Long
 
   /**
    * Reverts writes that haven't been flushed yet. Callers should invoke this function
@@ -106,7 +106,7 @@ private[spark] class DiskBlockObjectWriter(
   /** The file channel, used for repositioning / truncating the file. */
   private var channel: FileChannel = null
   private var bs: OutputStream = null
-  private var fos: ByteArrayOutputStream = null
+  private var fos: FileOutputStream = null
   private var ts: TimeTrackingOutputStream = null
   private var objOut: SerializationStream = null
   private val initialPosition = file.length()
@@ -115,8 +115,9 @@ private[spark] class DiskBlockObjectWriter(
   private var _timeWriting = 0L
 
   override def open(): BlockObjectWriter = {
-    fos = new ByteArrayOutputStream()
+    fos = new FileOutputStream(file, true)
     ts = new TimeTrackingOutputStream(fos)
+    channel = fos.getChannel()
     lastValidPosition = initialPosition
     bs = compressStream(new FastBufferedOutputStream(ts, bufferSize))
     objOut = serializer.newInstance().serializeStream(bs)
@@ -129,6 +130,9 @@ private[spark] class DiskBlockObjectWriter(
       if (syncWrites) {
         // Force outstanding writes to disk and track how long it takes
         objOut.flush()
+        val start = System.nanoTime()
+        fos.getFD.sync()
+        _timeWriting += System.nanoTime() - start
       }
       objOut.close()
 
@@ -145,18 +149,18 @@ private[spark] class DiskBlockObjectWriter(
 
   override def isOpen: Boolean = objOut != null
 
-  override def commit(): Array[Byte] = {
+  override def commit(): Long = {
     if (initialized) {
       // NOTE: Because Kryo doesn't flush the underlying stream we explicitly flush both the
       //       serializer stream and the lower level stream.
       objOut.flush()
       bs.flush()
       val prevPos = lastValidPosition
-      lastValidPosition = fos.size()
-      fos.toByteArray
+      lastValidPosition = channel.position()
+      lastValidPosition - prevPos
     } else {
       // lastValidPosition is zero if stream is uninitialized
-      null
+      lastValidPosition
     }
   }
 
@@ -166,7 +170,7 @@ private[spark] class DiskBlockObjectWriter(
       // truncate the file to the last valid position.
       objOut.flush()
       bs.flush()
-      throw new UnsupportedOperationException("Revert temporarily broken due to in memory shuffle code changes.")
+      channel.truncate(lastValidPosition)
     }
   }
 
@@ -178,7 +182,7 @@ private[spark] class DiskBlockObjectWriter(
   }
 
   override def fileSegment(): FileSegment = {
-    new FileSegment(null, initialPosition, bytesWritten)
+    new FileSegment(file, initialPosition, bytesWritten)
   }
 
   // Only valid if called after close()

core/src/main/scala/org/apache/spark/storage/MemoryStore.scala

@@ -23,7 +23,6 @@ import java.util.LinkedHashMap
 import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.util.{SizeEstimator, Utils}
-import org.apache.spark.serializer.Serializer
 
 /**
  * Stores blocks in memory, either as ArrayBuffers of deserialized Java objects or as
@@ -120,14 +119,6 @@ private class MemoryStore(blockManager: BlockManager, maxMemory: Long)
     }
   }
 
-  /**
-   * A version of getValues that allows a custom serializer. This is used as part of the
-   * shuffle short-circuit code.
-   */
-  def getValues(blockId: BlockId, serializer: Serializer): Option[Iterator[Any]] = {
-    getBytes(blockId).map(bytes => blockManager.dataDeserialize(blockId, bytes, serializer))
-  }
-
   override def remove(blockId: BlockId): Boolean = {
     entries.synchronized {
       val entry = entries.remove(blockId)

core/src/main/scala/org/apache/spark/storage/ShuffleBlockManager.scala

@@ -187,17 +187,6 @@ class ShuffleBlockManager(blockManager: BlockManager) extends Logging {
       }
     })
   }
-
-  def removeAllShuffleStuff() {
-    for (state <- shuffleStates.values;
-         group <- state.allFileGroups;
-         (mapId, _) <- group.mapIdToIndex.iterator;
-         reducerId <- 0 until group.files.length) {
-      val blockId = new ShuffleBlockId(group.shuffleId, mapId, reducerId)
-      blockManager.removeBlock(blockId, tellMaster = false)
-    }
-    shuffleStates.clear()
-  }
 }
 
 private[spark]
@@ -211,7 +200,7 @@ object ShuffleBlockManager {
      * Stores the absolute index of each mapId in the files of this group. For instance,
      * if mapId 5 is the first block in each file, mapIdToIndex(5) = 0.
      */
-    val mapIdToIndex = new PrimitiveKeyOpenHashMap[Int, Int]()
+    private val mapIdToIndex = new PrimitiveKeyOpenHashMap[Int, Int]()
 
     /**
      * Stores consecutive offsets of blocks into each reducer file, ordered by position in the file.

graphx/src/main/scala/org/apache/spark/graphx/Pregel.scala

@@ -19,7 +19,6 @@ package org.apache.spark.graphx
 
 import scala.reflect.ClassTag
 import org.apache.spark.Logging
-import org.apache.spark.SparkEnv
 
 
 /**
@@ -147,11 +146,6 @@ object Pregel extends Logging {
       // vertices of prevG (depended on by newVerts, oldMessages, and the vertices of g).
       activeMessages = messages.count()
 
-      // Very ugly code to clear the in-memory shuffle data
-      messages.foreachPartition { iter =>
-        SparkEnv.get.blockManager.shuffleBlockManager.removeAllShuffleStuff()
-      }
-
       logWarning("Pregel finished iteration " + i)
 
       // Unpersist the RDDs hidden by newly-materialized RDDs