Updated tests, although DistributedSuite is hanging.

Author: rxin

core/src/main/scala/org/apache/spark/network/BlockTransferService.scala

@@ -63,8 +63,39 @@ abstract class BlockTransferService {
    * }}}
    */
   def fetchBlock(hostName: String, port: Int, blockId: String): ManagedBuffer = {
-    //fetchBlocks(hostName, port, Seq(blockId)).iterator().next()._2
-    null
+    // TODO(rxin): Add timeout?
+    val lock = new Object
+    @volatile var result: Either[ManagedBuffer, Exception] = null
+    fetchBlocks(hostName, port, Seq(blockId), new BlockFetchingListener {
+      override def onBlockFetchFailure(exception: Exception): Unit = {
+        lock.synchronized {
+          result = Right(exception)
+          lock.notify()
+        }
+      }
+      override def onBlockFetchSuccess(blockId: String, data: ManagedBuffer): Unit = {
+        lock.synchronized {
+          result = Left(data)
+          lock.notify()
+        }
+      }
+    })
+
+    // Sleep until result is no longer null
+    lock.synchronized {
+      while (result == null) {
+        try {
+          lock.wait()
+        } catch {
+          case e: InterruptedException =>
+        }
+      }
+    }
+
+    result match {
+      case Left(data: ManagedBuffer) => data
+      case Right(e: Exception) => throw e
+    }
   }
 
   /**

core/src/main/scala/org/apache/spark/network/ManagedBuffer.scala

@@ -17,8 +17,9 @@
 
 package org.apache.spark.network
 
-import java.io.{File, FileInputStream, InputStream}
+import java.io.{RandomAccessFile, File, FileInputStream, InputStream}
 import java.nio.ByteBuffer
+import java.nio.channels.FileChannel.MapMode
 
 import io.netty.buffer.{ByteBufInputStream, ByteBuf, Unpooled}
 import io.netty.channel.DefaultFileRegion
@@ -34,7 +35,7 @@ abstract class ManagedBuffer {
   // Note that all the methods are defined with parenthesis because their implementations can
   // have side effects (io operations).
 
-  def byteBuffer(): ByteBuffer = throw new UnsupportedOperationException
+  def byteBuffer(): ByteBuffer
 
   def fileSegment(): Option[FileSegment] = None
 
@@ -56,6 +57,11 @@ final class FileSegmentManagedBuffer(file: File, offset: Long, length: Long)
 
   override def size: Long = length
 
+  override def byteBuffer(): ByteBuffer = {
+    val channel = new RandomAccessFile(file, "r").getChannel
+    channel.map(MapMode.READ_ONLY, offset, length)
+  }
+
   override private[network] def toNetty(): AnyRef = {
     val fileChannel = new FileInputStream(file).getChannel
     new DefaultFileRegion(fileChannel, offset, length)

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

@@ -63,9 +63,9 @@ private[spark] class BlockManager(
     mapOutputTracker: MapOutputTracker,
     shuffleManager: ShuffleManager,
     blockTransferService: BlockTransferService)
-  extends BlockDataProvider with Logging {
+  extends BlockDataManager with Logging {
 
-  //blockTransferService.init(this)
+  blockTransferService.init(this)
 
   private val port = conf.getInt("spark.blockManager.port", 0)
   val shuffleBlockManager = new ShuffleBlockManager(this, shuffleManager)
@@ -207,20 +207,34 @@ private[spark] class BlockManager(
     }
   }
 
-  override def getBlockData(blockId: String): Either[FileSegment, ByteBuffer] = {
+  /**
+   * Interface to get local block data.
+   *
+   * @return Some(buffer) if the block exists locally, and None if it doesn't.
+   */
+  override def getBlockData(blockId: String): Option[ManagedBuffer] = {
     val bid = BlockId(blockId)
     if (bid.isShuffle) {
-      Left(diskBlockManager.getBlockLocation(bid))
+      val fileSegment = diskBlockManager.getBlockLocation(bid)
+      Some(new FileSegmentManagedBuffer(fileSegment.file, fileSegment.offset, fileSegment.length))
     } else {
       val blockBytesOpt = doGetLocal(bid, asBlockResult = false).asInstanceOf[Option[ByteBuffer]]
       if (blockBytesOpt.isDefined) {
-        Right(blockBytesOpt.get)
+        val buffer = blockBytesOpt.get
+        Some(new NioByteBufferManagedBuffer(buffer))
       } else {
-        throw new BlockNotFoundException(blockId)
+        None
       }
     }
   }
 
+  /**
+   * Put the block locally, using the given storage level.
+   */
+  override def putBlockData(blockId: String, data: ManagedBuffer, level: StorageLevel): Unit = {
+    putBytes(BlockId(blockId), data.byteBuffer(), level)
+  }
+
   /**
    * Get the BlockStatus for the block identified by the given ID, if it exists.
    * NOTE: This is mainly for testing, and it doesn't fetch information from Tachyon.

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

@@ -231,7 +231,8 @@ final class ShuffleBlockFetcherIterator(
     if (!result.failed) {
       bytesInFlight -= result.size
     }
-    while (!fetchRequests.isEmpty &&
+    // Send fetch requests up to maxBytesInFlight
+    while (fetchRequests.nonEmpty &&
       (bytesInFlight == 0 || bytesInFlight + fetchRequests.front.size <= maxBytesInFlight)) {
       sendRequest(fetchRequests.dequeue())
     }

core/src/test/scala/org/apache/spark/DistributedSuite.scala

@@ -17,7 +17,6 @@
 
 package org.apache.spark
 
-import org.apache.spark.network.cm.{GetBlock, BlockManagerWorker, ConnectionManagerId}
 import org.scalatest.BeforeAndAfter
 import org.scalatest.FunSuite
 import org.scalatest.concurrent.Timeouts._
@@ -202,12 +201,13 @@ class DistributedSuite extends FunSuite with Matchers with BeforeAndAfter
     val blockIds = data.partitions.indices.map(index => RDDBlockId(data.id, index)).toArray
     val blockId = blockIds(0)
     val blockManager = SparkEnv.get.blockManager
-    blockManager.master.getLocations(blockId).foreach(id => {
-      val bytes = BlockManagerWorker.syncGetBlock(
-        GetBlock(blockId), ConnectionManagerId(id.host, id.port))
-      val deserialized = blockManager.dataDeserialize(blockId, bytes).asInstanceOf[Iterator[Int]].toList
+    val blockTransfer = SparkEnv.get.blockTransferService
+    blockManager.master.getLocations(blockId).foreach { cmId =>
+      val bytes = blockTransfer.fetchBlock(cmId.host, cmId.port, blockId.toString)
+      val deserialized = blockManager.dataDeserialize(blockId, bytes.byteBuffer())
+        .asInstanceOf[Iterator[Int]].toList
       assert(deserialized === (1 to 100).toList)
-    })
+    }
   }
 
   test("compute without caching when no partitions fit in memory") {
@@ -268,6 +268,8 @@ class DistributedSuite extends FunSuite with Matchers with BeforeAndAfter
     DistributedSuite.amMaster = true
     sc = new SparkContext(clusterUrl, "test")
     for (i <- 1 to 3) {
+      println("i = " + i)
+      Console.out.flush()
       val data = sc.parallelize(Seq(true, true), 2)
       assert(data.count === 2)
       assert(data.map(markNodeIfIdentity).collect.size === 2)
@@ -339,6 +341,7 @@ object DistributedSuite {
   // Act like an identity function, but if the argument is true, set mark to true.
   def markNodeIfIdentity(item: Boolean): Boolean = {
     if (item) {
+      println("marking node!!!!!!!!!!!!!!!")
       assert(!amMaster)
       mark = true
     }
@@ -349,6 +352,7 @@ object DistributedSuite {
   // crashing the entire JVM.
   def failOnMarkedIdentity(item: Boolean): Boolean = {
     if (mark) {
+      println("failing node !!!!!!!!!!!!!!!")
       System.exit(42)
     }
     item