Added some basic test cases.

Author: rxin

core/src/main/scala/org/apache/spark/network/netty/NettyConfig.scala

@@ -22,22 +22,26 @@ import org.apache.spark.SparkConf
 /**
  * A central location that tracks all the settings we exposed to users.
  */
+private[spark]
 class NettyConfig(conf: SparkConf) {
 
+  /** Port the server listens on. Default to a random port. */
+  private[netty] val serverPort = conf.getInt("spark.shuffle.io.port", 0)
+
   /** IO mode: nio, oio, epoll, or auto (try epoll first and then nio). */
-  val ioMode = conf.get("spark.shuffle.io.mode", "auto").toLowerCase
+  private[netty] val ioMode = conf.get("spark.shuffle.io.mode", "auto").toLowerCase
 
   /** Connect timeout in secs. Default 60 secs. */
-  val connectTimeoutMs = conf.getInt("spark.shuffle.io.connectionTimeout", 60) * 1000
+  private[netty] val connectTimeoutMs = conf.getInt("spark.shuffle.io.connectionTimeout", 60) * 1000
 
   /**
    * Percentage of the desired amount of time spent for I/O in the child event loops.
    * Only applicable in nio and epoll.
    */
-  val ioRatio = conf.getInt("spark.shuffle.io.netty.ioRatio", 80)
+  private[netty] val ioRatio = conf.getInt("spark.shuffle.io.netty.ioRatio", 80)
 
   /** Requested maximum length of the queue of incoming connections. */
-  val backLog: Option[Int] = conf.getOption("spark.shuffle.io.backLog").map(_.toInt)
+  private[netty] val backLog: Option[Int] = conf.getOption("spark.shuffle.io.backLog").map(_.toInt)
 
   /**
    * Receive buffer size (SO_RCVBUF).
@@ -46,8 +50,10 @@ class NettyConfig(conf: SparkConf) {
    * Assuming latency = 1ms, network_bandwidth = 10Gbps
    *  buffer size should be ~ 1.25MB
    */
-  val receiveBuf: Option[Int] = conf.getOption("spark.shuffle.io.sendBuffer").map(_.toInt)
+  private[netty] val receiveBuf: Option[Int] =
+    conf.getOption("spark.shuffle.io.sendBuffer").map(_.toInt)
 
   /** Send buffer size (SO_SNDBUF). */
-  val sendBuf: Option[Int] = conf.getOption("spark.shuffle.io.sendBuffer").map(_.toInt)
+  private[netty] val sendBuf: Option[Int] =
+    conf.getOption("spark.shuffle.io.sendBuffer").map(_.toInt)
 }

core/src/main/scala/org/apache/spark/network/netty/client/BlockFetchingClient.scala

@@ -129,7 +129,5 @@ class BlockFetchingClient(factory: BlockFetchingClientFactory, hostname: String,
     cf.channel().closeFuture().sync()
   }
 
-  def close(): Unit = {
-    // TODO: Should we ever close the client? Probably ...
-  }
+  def close(): Unit = cf.channel().close()
 }

core/src/main/scala/org/apache/spark/network/netty/client/BlockFetchingClientHandler.scala

@@ -17,8 +17,6 @@
 
 package org.apache.spark.network.netty.client
 
-import java.net.InetSocketAddress
-
 import io.netty.buffer.ByteBuf
 import io.netty.channel.{ChannelHandlerContext, SimpleChannelInboundHandler}
 
@@ -43,10 +41,7 @@ class BlockFetchingClientHandler extends SimpleChannelInboundHandler[ByteBuf] wi
     val blockId = new String(blockIdBytes)
     val blockLen = math.abs(totalLen) - blockIdLen - 4
 
-    def server = {
-      val remoteAddr = ctx.channel.remoteAddress.asInstanceOf[InetSocketAddress]
-      remoteAddr.getHostName + ":" + remoteAddr.getPort
-    }
+    def server = ctx.channel.remoteAddress.toString
 
     // totalLen is negative when it is an error message.
     if (totalLen < 0) {

core/src/main/scala/org/apache/spark/network/netty/client/ClientTester.scala

@@ -21,26 +21,27 @@ import java.net.InetSocketAddress
 
 import io.netty.bootstrap.ServerBootstrap
 import io.netty.buffer.PooledByteBufAllocator
-import io.netty.channel.socket.SocketChannel
-import io.netty.channel.{ChannelInitializer, ChannelOption, ChannelFuture}
-import io.netty.channel.epoll.{EpollServerSocketChannel, EpollEventLoopGroup}
+import io.netty.channel.{ChannelFuture, ChannelInitializer, ChannelOption}
+import io.netty.channel.epoll.{EpollEventLoopGroup, EpollServerSocketChannel}
 import io.netty.channel.nio.NioEventLoopGroup
 import io.netty.channel.oio.OioEventLoopGroup
+import io.netty.channel.socket.SocketChannel
 import io.netty.channel.socket.nio.NioServerSocketChannel
 import io.netty.channel.socket.oio.OioServerSocketChannel
 import io.netty.handler.codec.LineBasedFrameDecoder
 import io.netty.handler.codec.string.StringDecoder
 import io.netty.util.CharsetUtil
 
-import org.apache.spark.{SparkConf, Logging}
-import org.apache.spark.network.netty.{PathResolver, NettyConfig}
+import org.apache.spark.{Logging, SparkConf}
+import org.apache.spark.network.netty.NettyConfig
+import org.apache.spark.storage.BlockDataProvider
 import org.apache.spark.util.Utils
 
+
 // TODO: Remove dependency on BlockId. This layer should not be coupled with storage.
 
 // TODO: PathResolver is not general enough. It only works for on-disk blocks.
 
-// TODO: Allow user-configured port
 
 /**
  * Server for serving Spark data blocks.
@@ -58,15 +59,18 @@ import org.apache.spark.util.Utils
  *
  */
 private[spark]
-class BlockServer(conf: NettyConfig, pResolver: PathResolver) extends Logging {
+class BlockServer(conf: NettyConfig, dataProvider: BlockDataProvider) extends Logging {
 
-  def this(sparkConf: SparkConf, pResolver: PathResolver) = {
-    this(new NettyConfig(sparkConf), pResolver)
+  def this(sparkConf: SparkConf, dataProvider: BlockDataProvider) = {
+    this(new NettyConfig(sparkConf), dataProvider)
   }
 
   def port: Int = _port
 
-  private var _port: Int = 0
+  def hostName: String = _hostName
+
+  private var _port: Int = conf.serverPort
+  private var _hostName: String = ""
   private var bootstrap: ServerBootstrap = _
   private var channelFuture: ChannelFuture = _
 
@@ -134,7 +138,7 @@ class BlockServer(conf: NettyConfig, pResolver: PathResolver) extends Logging {
           .addLast("frameDecoder", new LineBasedFrameDecoder(1024))  // max block id length 1024
           .addLast("stringDecoder", new StringDecoder(CharsetUtil.UTF_8))
           .addLast("blockHeaderEncoder", new BlockHeaderEncoder)
-          .addLast("handler", new BlockServerHandler(pResolver))
+          .addLast("handler", new BlockServerHandler(dataProvider))
       }
     })
 
@@ -143,6 +147,7 @@ class BlockServer(conf: NettyConfig, pResolver: PathResolver) extends Logging {
 
     val addr = channelFuture.channel.localAddress.asInstanceOf[InetSocketAddress]
     _port = addr.getPort
+    _hostName = addr.getHostName
   }
 
   /** Shutdown the server. */

core/src/main/scala/org/apache/spark/network/netty/server/BlockServer.scala

@@ -18,14 +18,14 @@
 package org.apache.spark.network.netty.server
 
 import java.io.FileInputStream
-import java.net.InetSocketAddress
+import java.nio.ByteBuffer
 import java.nio.channels.FileChannel
 
+import io.netty.buffer.Unpooled
 import io.netty.channel._
 
 import org.apache.spark.Logging
-import org.apache.spark.network.netty.PathResolver
-import org.apache.spark.storage.BlockId
+import org.apache.spark.storage.{FileSegment, BlockDataProvider}
 
 
 /**
@@ -35,15 +35,11 @@ import org.apache.spark.storage.BlockId
  * so channelRead0 is called once per line (i.e. per block id).
  */
 private[server]
-class BlockServerHandler(p: PathResolver)
+class BlockServerHandler(dataProvider: BlockDataProvider)
   extends SimpleChannelInboundHandler[String] with Logging {
 
   override def channelRead0(ctx: ChannelHandlerContext, blockId: String): Unit = {
-    // client in the form of hostname:port
-    val client = {
-      val remoteAddr = ctx.channel.remoteAddress.asInstanceOf[InetSocketAddress]
-      remoteAddr.getHostName + ":" + remoteAddr.getPort
-    }
+    def client = ctx.channel.remoteAddress.toString
 
     // A helper function to send error message back to the client.
     def respondWithError(error: String): Unit = {
@@ -60,48 +56,80 @@ class BlockServerHandler(p: PathResolver)
       )
     }
 
-    logTrace(s"Received request from $client to fetch block $blockId")
+    def writeFileSegment(segment: FileSegment): Unit = {
+      // Send error message back if the block is too large. Even though we are capable of sending
+      // large (2G+) blocks, the receiving end cannot handle it so let's fail fast.
+      // Once we fixed the receiving end to be able to process large blocks, this should be removed.
+      // Also make sure we update BlockHeaderEncoder to support length > 2G.
 
-    var fileChannel: FileChannel = null
-    var offset: Long = 0
-    var blockSize: Long = 0
+      // See [[BlockHeaderEncoder]] for the way length is encoded.
+      if (segment.length + blockId.length + 4 > Int.MaxValue) {
+        respondWithError(s"Block $blockId size ($segment.length) greater than 2G")
+        return
+      }
 
-    // First make sure we can find the block. If not, send error back to the user.
-    try {
-      val segment = p.getBlockLocation(BlockId(blockId))
-      fileChannel = new FileInputStream(segment.file).getChannel
-      offset = segment.offset
-      blockSize = segment.length
-    } catch {
-      case e: Exception =>
-        logError(s"Error opening block $blockId for request from $client", e)
-        blockSize = -1
-        respondWithError(e.getMessage)
-    }
+      var fileChannel: FileChannel = null
+      try {
+        fileChannel = new FileInputStream(segment.file).getChannel
+      } catch {
+        case e: Exception =>
+          logError(
+            s"Error opening channel for $blockId in ${segment.file} for request from $client", e)
+          respondWithError(e.getMessage)
+      }
+
+      // Found the block. Send it back.
+      if (fileChannel != null) {
+        // Write the header and block data. In the case of failures, the listener on the block data
+        // write should close the connection.
+        ctx.write(new BlockHeader(segment.length.toInt, blockId))
 
-    // Send error message back if the block is too large. Even though we are capable of sending
-    // large (2G+) blocks, the receiving end cannot handle it so let's fail fast.
-    // Once we fixed the receiving end to be able to process large blocks, this should be removed.
-    // Also make sure we update BlockHeaderEncoder to support length > 2G.
-    if (blockSize > Int.MaxValue) {
-      respondWithError(s"Block $blockId size ($blockSize) greater than 2G")
+        val region = new DefaultFileRegion(fileChannel, segment.offset, segment.length)
+        ctx.writeAndFlush(region).addListener(new ChannelFutureListener {
+          override def operationComplete(future: ChannelFuture) {
+            if (future.isSuccess) {
+              logTrace(s"Sent block $blockId (${segment.length} B) back to $client")
+            } else {
+              logError(s"Error sending block $blockId to $client; closing connection", future.cause)
+              ctx.close()
+            }
+          }
+        })
+      }
     }
 
-    // Found the block. Send it back.
-    if (fileChannel != null && blockSize >= 0) {
-      val listener = new ChannelFutureListener {
+    def writeByteBuffer(buf: ByteBuffer): Unit = {
+      ctx.write(new BlockHeader(buf.remaining, blockId))
+      ctx.writeAndFlush(Unpooled.wrappedBuffer(buf)).addListener(new ChannelFutureListener {
         override def operationComplete(future: ChannelFuture) {
           if (future.isSuccess) {
-            logTrace(s"Sent block $blockId ($blockSize B) back to $client")
+            logTrace(s"Sent block $blockId (${buf.remaining} B) back to $client")
           } else {
             logError(s"Error sending block $blockId to $client; closing connection", future.cause)
             ctx.close()
           }
         }
-      }
-      val region = new DefaultFileRegion(fileChannel, offset, blockSize)
-      ctx.writeAndFlush(new BlockHeader(blockSize.toInt, blockId)).addListener(listener)
-      ctx.writeAndFlush(region).addListener(listener)
+      })
+    }
+
+    logTrace(s"Received request from $client to fetch block $blockId")
+
+    var blockData: Either[FileSegment, ByteBuffer] = null
+
+    // First make sure we can find the block. If not, send error back to the user.
+    try {
+      blockData = dataProvider.getBlockData(blockId)
+    } catch {
+      case e: Exception =>
+        logError(s"Error opening block $blockId for request from $client", e)
+        respondWithError(e.getMessage)
+        return
+    }
+
+    blockData match {
+      case Left(segment) => writeFileSegment(segment)
+      case Right(buf) => writeByteBuffer(buf)
     }
+
   }  // end of channelRead0
 }

core/src/main/scala/org/apache/spark/network/netty/server/BlockServerHandler.scala

@@ -15,21 +15,18 @@
  * limitations under the License.
  */
 
-package org.apache.spark.network.netty.server
+package org.apache.spark.storage
 
-import org.apache.spark.SparkConf
-import org.apache.spark.network.netty.PathResolver
-import org.apache.spark.storage.{TestBlockId, FileSegment, BlockId}
+import java.nio.ByteBuffer
 
-/** A simple main function for testing the server. */
-object ServerTester {
-  def main(args: Array[String]): Unit = {
-    new BlockServer(new SparkConf, new PathResolver {
-      override def getBlockLocation(blockId: BlockId): FileSegment = {
-        val file = new java.io.File(blockId.asInstanceOf[TestBlockId].id)
-        new FileSegment(file, 0, file.length())
-      }
-    })
-    Thread.sleep(1000000)
-  }
+
+/**
+ * An interface for providing data for blocks.
+ *
+ * getBlockData returns either a FileSegment (for zero-copy send), or a ByteBuffer.
+ *
+ * Aside from unit tests, [[BlockManager]] is the main class that implements this.
+ */
+private[spark] trait BlockDataProvider {
+  def getBlockData(blockId: String): Either[FileSegment, ByteBuffer]
 }

core/src/main/scala/org/apache/spark/network/netty/server/ServerTester.scala renamed to core/src/main/scala/org/apache/spark/storage/BlockDataProvider.scala

@@ -277,7 +277,7 @@ object BlockFetcherIterator {
 
       bytesInFlight += req.size
       val sizeMap = req.blocks.toMap // so we can look up the size of each blockID
-      val client = blockManager.diskBlockManager.nettyBlockClientFactory.createClient(
+      val client = blockManager.nettyBlockClientFactory.createClient(
         cmId.host, req.address.nettyPort)
       val blocks = req.blocks.map(_._1.toString)