simplify code

Author: xuanyuanking

common/network-common/src/main/java/org/apache/spark/network/TransportContext.java

@@ -123,7 +123,7 @@ public TransportContext(
 
     if (conf.getModuleName() != null &&
         conf.getModuleName().equalsIgnoreCase("shuffle") &&
-        !isClientOnly) {
+        !isClientOnly && conf.separateChunkFetchRequest()) {
       chunkFetchWorkers = NettyUtils.createEventLoop(
           IOMode.valueOf(conf.ioMode()),
           conf.chunkFetchHandlerThreads(),
@@ -187,8 +187,6 @@ public TransportChannelHandler initializePipeline(
       RpcHandler channelRpcHandler) {
     try {
       TransportChannelHandler channelHandler = createChannelHandler(channel, channelRpcHandler);
-      ChunkFetchRequestHandler chunkFetchHandler =
-        createChunkFetchHandler(channelHandler, channelRpcHandler);
       ChannelPipeline pipeline = channel.pipeline()
         .addLast("encoder", ENCODER)
         .addLast(TransportFrameDecoder.HANDLER_NAME, NettyUtils.createFrameDecoder())
@@ -200,6 +198,9 @@ public TransportChannelHandler initializePipeline(
         .addLast("handler", channelHandler);
       // Use a separate EventLoopGroup to handle ChunkFetchRequest messages for shuffle rpcs.
       if (chunkFetchWorkers != null) {
+        ChunkFetchRequestHandler chunkFetchHandler = new ChunkFetchRequestHandler(
+          channelHandler.getClient(), rpcHandler.getStreamManager(),
+          conf.maxChunksBeingTransferred(), true /* syncModeEnabled */);
         pipeline.addLast(chunkFetchWorkers, "chunkFetchHandler", chunkFetchHandler);
       }
       return channelHandler;
@@ -217,19 +218,19 @@ public TransportChannelHandler initializePipeline(
   private TransportChannelHandler createChannelHandler(Channel channel, RpcHandler rpcHandler) {
     TransportResponseHandler responseHandler = new TransportResponseHandler(channel);
     TransportClient client = new TransportClient(channel, responseHandler);
+    boolean separateChunkFetchRequest = conf.separateChunkFetchRequest();
+    ChunkFetchRequestHandler chunkFetchRequestHandler;
+    if (!separateChunkFetchRequest) {
+      chunkFetchRequestHandler = new ChunkFetchRequestHandler(
+        client, rpcHandler.getStreamManager(),
+        conf.maxChunksBeingTransferred(), false /* syncModeEnabled */);
+    } else {
+      chunkFetchRequestHandler = null;
+    }
     TransportRequestHandler requestHandler = new TransportRequestHandler(channel, client,
-      rpcHandler, conf.maxChunksBeingTransferred());
+      rpcHandler, conf.maxChunksBeingTransferred(), chunkFetchRequestHandler);
     return new TransportChannelHandler(client, responseHandler, requestHandler,
-      conf.connectionTimeoutMs(), closeIdleConnections, this);
-  }
-
-  /**
-   * Creates the dedicated ChannelHandler for ChunkFetchRequest messages.
-   */
-  private ChunkFetchRequestHandler createChunkFetchHandler(TransportChannelHandler channelHandler,
-      RpcHandler rpcHandler) {
-    return new ChunkFetchRequestHandler(channelHandler.getClient(),
-      rpcHandler.getStreamManager(), conf.maxChunksBeingTransferred());
+      conf.connectionTimeoutMs(), separateChunkFetchRequest, closeIdleConnections, this);
   }
 
   public TransportConf getConf() { return conf; }

common/network-common/src/main/java/org/apache/spark/network/server/ChunkFetchRequestHandler.java

@@ -55,14 +55,17 @@ public class ChunkFetchRequestHandler extends SimpleChannelInboundHandler<ChunkF
   private final StreamManager streamManager;
   /** The max number of chunks being transferred and not finished yet. */
   private final long maxChunksBeingTransferred;
+  private final boolean syncModeEnabled;
 
   public ChunkFetchRequestHandler(
       TransportClient client,
       StreamManager streamManager,
-      Long maxChunksBeingTransferred) {
+      Long maxChunksBeingTransferred,
+      boolean syncModeEnabled) {
     this.client = client;
     this.streamManager = streamManager;
     this.maxChunksBeingTransferred = maxChunksBeingTransferred;
+    this.syncModeEnabled = syncModeEnabled;
   }
 
   @Override
@@ -76,6 +79,11 @@ protected void channelRead0(
       ChannelHandlerContext ctx,
       final ChunkFetchRequest msg) throws Exception {
     Channel channel = ctx.channel();
+    processFetchRequest(channel, msg);
+  }
+
+  public void processFetchRequest(
+      final Channel channel, final ChunkFetchRequest msg) throws Exception {
     if (logger.isTraceEnabled()) {
       logger.trace("Received req from {} to fetch block {}", getRemoteAddress(channel),
         msg.streamChunkId);
@@ -112,19 +120,26 @@ protected void channelRead0(
    * channel will be handled by the EventLoop the channel is registered to. So even
    * though we are processing the ChunkFetchRequest in a separate thread pool, the actual I/O,
    * which is the potentially blocking call that could deplete server handler threads, is still
-   * being processed by TransportServer's default EventLoopGroup. In order to throttle the max
-   * number of threads that channel I/O for sending response to ChunkFetchRequest, the thread
-   * calling channel.writeAndFlush will wait for the completion of sending response back to
-   * client by invoking await(). This will throttle the rate at which threads from
-   * ChunkFetchRequest dedicated EventLoopGroup submit channel I/O requests to TransportServer's
-   * default EventLoopGroup, thus making sure that we can reserve some threads in
-   * TransportServer's default EventLoopGroup for handling other RPC messages.
+   * being processed by TransportServer's default EventLoopGroup.
+   *
+   * When syncModeEnabled is true, Spark will throttle the max number of threads that channel I/O
+   * for sending response to ChunkFetchRequest, the thread calling channel.writeAndFlush will wait
+   * for the completion of sending response back to client by invoking await(). This will throttle
+   * the rate at which threads from ChunkFetchRequest dedicated EventLoopGroup submit channel I/O
+   * requests to TransportServer's default EventLoopGroup, thus making sure that we can reserve
+   * some threads in TransportServer's default EventLoopGroup for handling other RPC messages.
    */
   private ChannelFuture respond(
       final Channel channel,
       final Encodable result) throws InterruptedException {
     final SocketAddress remoteAddress = channel.remoteAddress();
-    return channel.writeAndFlush(result).await().addListener((ChannelFutureListener) future -> {
+    ChannelFuture channelFuture;
+    if (syncModeEnabled) {
+      channelFuture = channel.writeAndFlush(result).await();
+    } else {
+      channelFuture = channel.writeAndFlush(result);
+    }
+    return channelFuture.addListener((ChannelFutureListener) future -> {
       if (future.isSuccess()) {
         logger.trace("Sent result {} to client {}", result, remoteAddress);
       } else {

common/network-common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java

@@ -58,19 +58,22 @@ public class TransportChannelHandler extends SimpleChannelInboundHandler<Message
   private final TransportRequestHandler requestHandler;
   private final long requestTimeoutNs;
   private final boolean closeIdleConnections;
+  private final boolean skipChunkFetchRequest;
   private final TransportContext transportContext;
 
   public TransportChannelHandler(
       TransportClient client,
       TransportResponseHandler responseHandler,
       TransportRequestHandler requestHandler,
       long requestTimeoutMs,
+      boolean skipChunkFetchRequest,
       boolean closeIdleConnections,
       TransportContext transportContext) {
     this.client = client;
     this.responseHandler = responseHandler;
     this.requestHandler = requestHandler;
     this.requestTimeoutNs = requestTimeoutMs * 1000L * 1000;
+    this.skipChunkFetchRequest = skipChunkFetchRequest;
     this.closeIdleConnections = closeIdleConnections;
     this.transportContext = transportContext;
   }
@@ -124,7 +127,7 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception {
    */
   @Override
   public boolean acceptInboundMessage(Object msg) throws Exception {
-    if (msg instanceof ChunkFetchRequest) {
+    if (skipChunkFetchRequest && msg instanceof ChunkFetchRequest) {
       return false;
     } else {
       return super.acceptInboundMessage(msg);

common/network-common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java

@@ -62,16 +62,21 @@ public class TransportRequestHandler extends MessageHandler<RequestMessage> {
   /** The max number of chunks being transferred and not finished yet. */
   private final long maxChunksBeingTransferred;
 
+  /** The dedicated ChannelHandler for ChunkFetchRequest messages. */
+  private final ChunkFetchRequestHandler chunkFetchRequestHandler;
+
   public TransportRequestHandler(
       Channel channel,
       TransportClient reverseClient,
       RpcHandler rpcHandler,
-      Long maxChunksBeingTransferred) {
+      Long maxChunksBeingTransferred,
+      ChunkFetchRequestHandler chunkFetchRequestHandler) {
     this.channel = channel;
     this.reverseClient = reverseClient;
     this.rpcHandler = rpcHandler;
     this.streamManager = rpcHandler.getStreamManager();
     this.maxChunksBeingTransferred = maxChunksBeingTransferred;
+    this.chunkFetchRequestHandler = chunkFetchRequestHandler;
   }
 
   @Override
@@ -97,8 +102,10 @@ public void channelInactive() {
   }
 
   @Override
-  public void handle(RequestMessage request) {
-    if (request instanceof RpcRequest) {
+  public void handle(RequestMessage request) throws Exception {
+    if (request instanceof ChunkFetchRequest) {
+      chunkFetchRequestHandler.processFetchRequest(channel, (ChunkFetchRequest) request);
+    } else if (request instanceof RpcRequest) {
       processRpcRequest((RpcRequest) request);
     } else if (request instanceof OneWayMessage) {
       processOneWayMessage((OneWayMessage) request);

common/network-common/src/main/java/org/apache/spark/network/util/TransportConf.java

@@ -18,6 +18,7 @@
 package org.apache.spark.network.util;
 
 import java.util.Locale;
+import java.util.NoSuchElementException;
 import java.util.Properties;
 
 import com.google.common.primitives.Ints;
@@ -316,7 +317,8 @@ public long maxChunksBeingTransferred() {
 
   /**
    * Percentage of io.serverThreads used by netty to process ChunkFetchRequest.
-   * Shuffle server will use a separate EventLoopGroup to process ChunkFetchRequest messages.
+   * When the config `spark.shuffle.server.chunkFetchHandlerThreadsPercent` is set,
+   * shuffle server will use a separate EventLoopGroup to process ChunkFetchRequest messages.
    * Although when calling the async writeAndFlush on the underlying channel to send
    * response back to client, the I/O on the channel is still being handled by
    * {@link org.apache.spark.network.server.TransportServer}'s default EventLoopGroup
@@ -339,12 +341,25 @@ public int chunkFetchHandlerThreads() {
       return 0;
     }
     int chunkFetchHandlerThreadsPercent =
-      conf.getInt("spark.shuffle.server.chunkFetchHandlerThreadsPercent", 100);
+      Integer.parseInt(conf.get("spark.shuffle.server.chunkFetchHandlerThreadsPercent"));
     int threads =
       this.serverThreads() > 0 ? this.serverThreads() : 2 * NettyRuntime.availableProcessors();
     return (int) Math.ceil(threads * (chunkFetchHandlerThreadsPercent / 100.0));
   }
 
+  /**
+   * Whether to use a separate EventLoopGroup to process ChunkFetchRequest messages, it is decided
+   * by the config `spark.shuffle.server.chunkFetchHandlerThreadsPercent` is set or not.
+   */
+  public boolean separateChunkFetchRequest() {
+    try {
+      conf.get("spark.shuffle.server.chunkFetchHandlerThreadsPercent");
+      return true;
+    } catch (NoSuchElementException e) {
+      return false;
+    }
+  }
+
   /**
    * Whether to use the old protocol while doing the shuffle block fetching.
    * It is only enabled while we need the compatibility in the scenario of new spark version

common/network-common/src/test/java/org/apache/spark/network/ChunkFetchRequestHandlerSuite.java

@@ -22,7 +22,6 @@
 import java.util.List;
 
 import io.netty.channel.Channel;
-import org.apache.spark.network.server.ChunkFetchRequestHandler;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -33,6 +32,7 @@
 import org.apache.spark.network.buffer.ManagedBuffer;
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.protocol.*;
+import org.apache.spark.network.server.ChunkFetchRequestHandler;
 import org.apache.spark.network.server.NoOpRpcHandler;
 import org.apache.spark.network.server.OneForOneStreamManager;
 import org.apache.spark.network.server.RpcHandler;
@@ -68,7 +68,7 @@ public void handleChunkFetchRequest() throws Exception {
     long streamId = streamManager.registerStream("test-app", managedBuffers.iterator(), channel);
     TransportClient reverseClient = mock(TransportClient.class);
     ChunkFetchRequestHandler requestHandler = new ChunkFetchRequestHandler(reverseClient,
-      rpcHandler.getStreamManager(), 2L);
+      rpcHandler.getStreamManager(), 2L, false);
 
     RequestMessage request0 = new ChunkFetchRequest(new StreamChunkId(streamId, 0));
     requestHandler.channelRead(context, request0);

common/network-common/src/test/java/org/apache/spark/network/TransportRequestHandlerSuite.java

@@ -39,7 +39,7 @@
 public class TransportRequestHandlerSuite {
 
   @Test
-  public void handleStreamRequest() {
+  public void handleStreamRequest() throws Exception {
     RpcHandler rpcHandler = new NoOpRpcHandler();
     OneForOneStreamManager streamManager = (OneForOneStreamManager) (rpcHandler.getStreamManager());
     Channel channel = mock(Channel.class);
@@ -66,7 +66,7 @@ public void handleStreamRequest() {
 
     TransportClient reverseClient = mock(TransportClient.class);
     TransportRequestHandler requestHandler = new TransportRequestHandler(channel, reverseClient,
-      rpcHandler, 2L);
+      rpcHandler, 2L, null);
 
     RequestMessage request0 = new StreamRequest(String.format("%d_%d", streamId, 0));
     requestHandler.handle(request0);