[SPARK-3453] Refactor Netty module to use BlockTransferService.

Also includes some partial support for uploading blocks.

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

@@ -22,7 +22,8 @@ import java.nio.ByteBuffer
 import java.nio.channels.FileChannel.MapMode
 
 import com.google.common.io.ByteStreams
-import io.netty.buffer.{ByteBufInputStream, ByteBuf}
+import io.netty.buffer.{Unpooled, ByteBufInputStream, ByteBuf}
+import io.netty.channel.DefaultFileRegion
 
 import org.apache.spark.util.ByteBufferInputStream
 
@@ -35,7 +36,7 @@ import org.apache.spark.util.ByteBufferInputStream
  * - NioByteBufferManagedBuffer: data backed by a NIO ByteBuffer
  * - NettyByteBufManagedBuffer: data backed by a Netty ByteBuf
  */
-sealed abstract class ManagedBuffer {
+abstract class ManagedBuffer {
   // Note that all the methods are defined with parenthesis because their implementations can
   // have side effects (io operations).
 
@@ -54,6 +55,11 @@ sealed abstract class ManagedBuffer {
    * it does not go over the limit.
    */
   def inputStream(): InputStream
+
+  /**
+   * Convert the buffer into an Netty object, used to write the data out.
+   */
+  private[network] def convertToNetty(): AnyRef
 }
 
 
@@ -75,6 +81,11 @@ final class FileSegmentManagedBuffer(val file: File, val offset: Long, val lengt
     is.skip(offset)
     ByteStreams.limit(is, length)
   }
+
+  private[network] override def convertToNetty(): AnyRef = {
+    val fileChannel = new FileInputStream(file).getChannel
+    new DefaultFileRegion(fileChannel, offset, length)
+  }
 }
 
 
@@ -88,6 +99,8 @@ final class NioByteBufferManagedBuffer(buf: ByteBuffer) extends ManagedBuffer {
   override def nioByteBuffer() = buf.duplicate()
 
   override def inputStream() = new ByteBufferInputStream(buf)
+
+  private[network] override def convertToNetty(): AnyRef = Unpooled.wrappedBuffer(buf)
 }
 
 
@@ -102,6 +115,8 @@ final class NettyByteBufManagedBuffer(buf: ByteBuf) extends ManagedBuffer {
 
   override def inputStream() = new ByteBufInputStream(buf)
 
+  private[network] override def convertToNetty(): AnyRef = buf
+
   // TODO(rxin): Promote this to top level ManagedBuffer interface and add documentation for it.
   def release(): Unit = buf.release()
 }

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

@@ -15,36 +15,35 @@
  * limitations under the License.
  */
 
-package org.apache.spark.network.netty.client
+package org.apache.spark.network.netty
 
 import java.util.concurrent.TimeoutException
 
 import io.netty.bootstrap.Bootstrap
 import io.netty.buffer.PooledByteBufAllocator
 import io.netty.channel.socket.SocketChannel
-import io.netty.channel.{ChannelFutureListener, ChannelFuture, ChannelInitializer, ChannelOption}
-import io.netty.handler.codec.LengthFieldBasedFrameDecoder
-import io.netty.handler.codec.string.StringEncoder
-import io.netty.util.CharsetUtil
+import io.netty.channel.{ChannelFuture, ChannelFutureListener, ChannelInitializer, ChannelOption}
 
 import org.apache.spark.Logging
+import org.apache.spark.network.BlockFetchingListener
+
 
 /**
- * Client for fetching data blocks from [[org.apache.spark.network.netty.server.BlockServer]].
- * Use [[BlockFetchingClientFactory]] to instantiate this client.
+ * Client for [[NettyBlockTransferService]]. Use [[BlockClientFactory]] to
+ * instantiate this client.
  *
  * The constructor blocks until a connection is successfully established.
  *
- * See [[org.apache.spark.network.netty.server.BlockServer]] for client/server protocol.
- *
  * Concurrency: thread safe and can be called from multiple threads.
  */
 @throws[TimeoutException]
-private[spark]
-class BlockFetchingClient(factory: BlockFetchingClientFactory, hostname: String, port: Int)
+private[netty]
+class BlockClient(factory: BlockClientFactory, hostname: String, port: Int)
   extends Logging {
 
-  private val handler = new BlockFetchingClientHandler
+  private val handler = new BlockClientHandler
+  private val encoder = new ClientRequestEncoder
+  private val decoder = new ServerResponseDecoder
 
   /** Netty Bootstrap for creating the TCP connection. */
   private val bootstrap: Bootstrap = {
@@ -61,9 +60,9 @@ class BlockFetchingClient(factory: BlockFetchingClientFactory, hostname: String,
     b.handler(new ChannelInitializer[SocketChannel] {
       override def initChannel(ch: SocketChannel): Unit = {
         ch.pipeline
-          .addLast("encoder", new StringEncoder(CharsetUtil.UTF_8))
-          // maxFrameLength = 2G, lengthFieldOffset = 0, lengthFieldLength = 4
-          .addLast("framedLengthDecoder", new LengthFieldBasedFrameDecoder(Int.MaxValue, 0, 4))
+          .addLast("clientRequestEncoder", encoder)
+          .addLast("frameDecoder", ProtocolUtils.createFrameDecoder())
+          .addLast("serverResponseDecoder", decoder)
           .addLast("handler", handler)
       }
     })
@@ -86,12 +85,7 @@ class BlockFetchingClient(factory: BlockFetchingClientFactory, hostname: String,
    * @param blockIds sequence of block ids to fetch.
    * @param listener callback to fire on fetch success / failure.
    */
-  def fetchBlocks(blockIds: Seq[String], listener: BlockClientListener): Unit = {
-    // It's best to limit the number of "write" calls since it needs to traverse the whole pipeline.
-    // It's also best to limit the number of "flush" calls since it requires system calls.
-    // Let's concatenate the string and then call writeAndFlush once.
-    // This is also why this implementation might be more efficient than multiple, separate
-    // fetch block calls.
+  def fetchBlocks(blockIds: Seq[String], listener: BlockFetchingListener): Unit = {
     var startTime: Long = 0
     logTrace {
       startTime = System.nanoTime
@@ -102,8 +96,7 @@ class BlockFetchingClient(factory: BlockFetchingClientFactory, hostname: String,
       handler.addRequest(blockId, listener)
     }
 
-    val writeFuture = cf.channel().writeAndFlush(blockIds.mkString("\n") + "\n")
-    writeFuture.addListener(new ChannelFutureListener {
+    cf.channel().writeAndFlush(BlockFetchRequest(blockIds)).addListener(new ChannelFutureListener {
       override def operationComplete(future: ChannelFuture): Unit = {
         if (future.isSuccess) {
           logTrace {
@@ -116,9 +109,9 @@ class BlockFetchingClient(factory: BlockFetchingClientFactory, hostname: String,
             s"Failed to send request $blockIds to $hostname:$port: ${future.cause.getMessage}"
           logError(errorMsg, future.cause)
           blockIds.foreach { blockId =>
-            listener.onFetchFailure(blockId, errorMsg)
             handler.removeRequest(blockId)
           }
+          listener.onBlockFetchFailure(new RuntimeException(errorMsg))
         }
       }
     })

core/src/main/scala/org/apache/spark/network/netty/client/BlockFetchingClientFactory.scala → core/src/main/scala/org/apache/spark/network/netty/BlockClientFactory.scala

@@ -15,36 +15,34 @@
  * limitations under the License.
  */
 
-package org.apache.spark.network.netty.client
+package org.apache.spark.network.netty
 
-import io.netty.channel.epoll.{EpollEventLoopGroup, EpollSocketChannel}
+import io.netty.channel.epoll.{Epoll, EpollEventLoopGroup, EpollSocketChannel}
 import io.netty.channel.nio.NioEventLoopGroup
 import io.netty.channel.oio.OioEventLoopGroup
 import io.netty.channel.socket.nio.NioSocketChannel
 import io.netty.channel.socket.oio.OioSocketChannel
-import io.netty.channel.{EventLoopGroup, Channel}
+import io.netty.channel.{Channel, EventLoopGroup}
 
 import org.apache.spark.SparkConf
-import org.apache.spark.network.netty.NettyConfig
 import org.apache.spark.util.Utils
 
+
 /**
- * Factory for creating [[BlockFetchingClient]] by using createClient. This factory reuses
+ * Factory for creating [[BlockClient]] by using createClient. This factory reuses
  * the worker thread pool for Netty.
- *
- * Concurrency: createClient is safe to be called from multiple threads concurrently.
  */
-private[spark]
-class BlockFetchingClientFactory(val conf: NettyConfig) {
+private[netty]
+class BlockClientFactory(val conf: NettyConfig) {
 
   def this(sparkConf: SparkConf) = this(new NettyConfig(sparkConf))
 
   /** A thread factory so the threads are named (for debugging). */
-  val threadFactory = Utils.namedThreadFactory("spark-shuffle-client")
+  private[netty] val threadFactory = Utils.namedThreadFactory("spark-shuffle-client")
 
   /** The following two are instantiated by the [[init]] method, depending ioMode. */
-  var socketChannelClass: Class[_ <: Channel] = _
-  var workerGroup: EventLoopGroup = _
+  private[netty] var socketChannelClass: Class[_ <: Channel] = _
+  private[netty] var workerGroup: EventLoopGroup = _
 
   init()
 
@@ -63,20 +61,12 @@ class BlockFetchingClientFactory(val conf: NettyConfig) {
       workerGroup = new EpollEventLoopGroup(0, threadFactory)
     }
 
+    // For auto mode, first try epoll (only available on Linux), then nio.
     conf.ioMode match {
       case "nio" => initNio()
       case "oio" => initOio()
       case "epoll" => initEpoll()
-      case "auto" =>
-        // For auto mode, first try epoll (only available on Linux), then nio.
-        try {
-          initEpoll()
-        } catch {
-          // TODO: Should we log the throwable? But that always happen on non-Linux systems.
-          // Perhaps the right thing to do is to check whether the system is Linux, and then only
-          // call initEpoll on Linux.
-          case e: Throwable => initNio()
-        }
+      case "auto" => if (Epoll.isAvailable) initEpoll() else initNio()
     }
   }
 
@@ -87,8 +77,8 @@ class BlockFetchingClientFactory(val conf: NettyConfig) {
    *
    * Concurrency: This method is safe to call from multiple threads.
    */
-  def createClient(remoteHost: String, remotePort: Int): BlockFetchingClient = {
-    new BlockFetchingClient(this, remoteHost, remotePort)
+  def createClient(remoteHost: String, remotePort: Int): BlockClient = {
+    new BlockClient(this, remoteHost, remotePort)
   }
 
   def stop(): Unit = {

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

@@ -0,0 +1,86 @@
+/*
+ * 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.network.netty
+
+import io.netty.channel.{ChannelHandlerContext, SimpleChannelInboundHandler}
+
+import org.apache.spark.Logging
+import org.apache.spark.network.BlockFetchingListener
+
+
+/**
+ * Handler that processes server responses.
+ *
+ * Concurrency: thread safe and can be called from multiple threads.
+ */
+private[netty]
+class BlockClientHandler extends SimpleChannelInboundHandler[ServerResponse] with Logging {
+
+  /** Tracks the list of outstanding requests and their listeners on success/failure. */
+  private val outstandingRequests = java.util.Collections.synchronizedMap {
+    new java.util.HashMap[String, BlockFetchingListener]
+  }
+
+  def addRequest(blockId: String, listener: BlockFetchingListener): Unit = {
+    outstandingRequests.put(blockId, listener)
+  }
+
+  def removeRequest(blockId: String): Unit = {
+    outstandingRequests.remove(blockId)
+  }
+
+  override def exceptionCaught(ctx: ChannelHandlerContext, cause: Throwable): Unit = {
+    val errorMsg = s"Exception in connection from ${ctx.channel.remoteAddress}: ${cause.getMessage}"
+    logError(errorMsg, cause)
+
+    // Fire the failure callback for all outstanding blocks
+    outstandingRequests.synchronized {
+      val iter = outstandingRequests.entrySet().iterator()
+      while (iter.hasNext) {
+        val entry = iter.next()
+        entry.getValue.onBlockFetchFailure(cause)
+      }
+      outstandingRequests.clear()
+    }
+
+    ctx.close()
+  }
+
+  override def channelRead0(ctx: ChannelHandlerContext, response: ServerResponse) {
+    val server = ctx.channel.remoteAddress.toString
+    response match {
+      case BlockFetchSuccess(blockId, buf) =>
+        val listener = outstandingRequests.get(blockId)
+        if (listener == null) {
+          logWarning(s"Got a response for block $blockId from $server but it is not outstanding")
+        } else {
+          outstandingRequests.remove(blockId)
+          listener.onBlockFetchSuccess(blockId, buf)
+        }
+      case BlockFetchFailure(blockId, errorMsg) =>
+        val listener = outstandingRequests.get(blockId)
+        if (listener == null) {
+          logWarning(
+            s"Got a response for block $blockId from $server ($errorMsg) but it is not outstanding")
+        } else {
+          outstandingRequests.remove(blockId)
+          listener.onBlockFetchFailure(new RuntimeException(errorMsg))
+        }
+    }
+  }
+}

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

@@ -15,48 +15,33 @@
  * limitations under the License.
  */
 
-package org.apache.spark.network.netty.server
+package org.apache.spark.network.netty
 
 import java.net.InetSocketAddress
 
 import io.netty.bootstrap.ServerBootstrap
 import io.netty.buffer.PooledByteBufAllocator
-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 io.netty.channel.{ChannelInitializer, ChannelFuture, ChannelOption}
+import io.netty.handler.codec.LengthFieldBasedFrameDecoder
 
 import org.apache.spark.{Logging, SparkConf}
-import org.apache.spark.network.netty.NettyConfig
-import org.apache.spark.storage.BlockDataProvider
+import org.apache.spark.network.BlockDataManager
 import org.apache.spark.util.Utils
 
 
 /**
- * Server for serving Spark data blocks.
- * This should be used together with [[org.apache.spark.network.netty.client.BlockFetchingClient]].
- *
- * Protocol for requesting blocks (client to server):
- *   One block id per line, e.g. to request 3 blocks: "block1\nblock2\nblock3\n"
- *
- * Protocol for sending blocks (server to client):
- *   frame-length (4 bytes), block-id-length (4 bytes), block-id, block-data.
- *
- *   frame-length should not include the length of itself.
- *   If block-id-length is negative, then this is an error message rather than block-data. The real
- *   length is the absolute value of the frame-length.
- *
+ * Server for the [[NettyBlockTransferService]].
  */
-private[spark]
-class BlockServer(conf: NettyConfig, dataProvider: BlockDataProvider) extends Logging {
+private[netty]
+class BlockServer(conf: NettyConfig, dataProvider: BlockDataManager) extends Logging {
 
-  def this(sparkConf: SparkConf, dataProvider: BlockDataProvider) = {
+  def this(sparkConf: SparkConf, dataProvider: BlockDataManager) = {
     this(new NettyConfig(sparkConf), dataProvider)
   }
 
@@ -129,10 +114,10 @@ class BlockServer(conf: NettyConfig, dataProvider: BlockDataProvider) extends Lo
 
     bootstrap.childHandler(new ChannelInitializer[SocketChannel] {
       override def initChannel(ch: SocketChannel): Unit = {
-        ch.pipeline
-          .addLast("frameDecoder", new LineBasedFrameDecoder(1024))  // max block id length 1024
-          .addLast("stringDecoder", new StringDecoder(CharsetUtil.UTF_8))
-          .addLast("blockHeaderEncoder", new BlockHeaderEncoder)
+        val p = ch.pipeline
+          .addLast("frameDecoder", ProtocolUtils.createFrameDecoder())
+          .addLast("clientRequestDecoder", new ClientRequestDecoder)
+          .addLast("serverResponseEncoder", new ServerResponseEncoder)
           .addLast("handler", new BlockServerHandler(dataProvider))
       }
     })