[SPARK-43987][SHUFFLE] Separate finalizeShuffleMerge Processing to Dedicated Thread Pools

### What changes were proposed in this pull request?

In this PR, we proposed to separate finalizeShuffleMerge processing into dedicated thread pools.
1. We introduce `ShuffleTransportContext` extend `TransportContext`.
2. We override the  `initializePipeline` by adding `FinalizeHandler` when the newly added configuration `spark.shuffle.server.finalizeShuffleMergeThreads` is positive.
3. We override the `decode` within `ShuffleMessageDecoder` so that an  `FINALIZE_SHUFFLE_MERGE` type `RpcRequest` will not be processed by current IO threads. We will further encapsulate it as `RpcRequestInternal`.
4. A dedicated `FinalizedHandler` will be attached to the channel pipeline, which only handles `RpcRequestInternal` type.

authors: otterc shuwang21
### Why are the changes needed?

In our production environment, `finalizeShuffleMerge` processing took longer time (p90 is around 20s) than other PRC requests. This is due to `finalizeShuffleMerge` invoking IO operations like truncate and file open/close.

More importantly, processing this `finalizeShuffleMerge` can block other critical lightweight messages like authentications, which can cause authentication timeout as well as fetch failures. Those timeout and fetch failures affect the stability of the Spark job executions.

### Does this PR introduce _any_ user-facing change?

No.
### How was this patch tested?

1. We add a few related UTs.
2. We also tested internally with this patch. Under the same setting, without this patch, the shuffle server can encounter authentication timeouts thus fetch failures.  With this path, `finalizeShuffleMerge` is processed by  `finalizeHandler` we introduced, and authentication RPC requests are not blocked.
3. In terms of performance, we have deployed internally. We saw Shuffle fetch delays have improved. P80 is reduced by 98.1% (120 s to 2.3 s). P90 is reduced by 83% (30mins to ~5mins). P99 is reduced by by 70% (44h to 13h). Furthermore, we saw The SASL timeout exceptions are reduced by ~40%. Spark job runtime is also improved. P50 is reduced by 35.2% (164s to 107s). P80 is reduced by 50.1% (15 mins s to 7.7mins). P90 is reduced by 45.5% (37 mins to 20 mins). P99 is reduced by 10% (4.5h to 4.0h). Finally, less runtime also indicates resources saving, and we are seeing 20% resources saving.

Closes #41489 from shuwang21/SPARK-43987.

Lead-authored-by: Shu Wang <swang7@linkedin.com>
Co-authored-by: Shu Wang <wangshu1990@gmail.com>
Co-authored-by: Chandni Singh <singh.chandni@gmail.com>
Signed-off-by: Mridul Muralidharan <mridul<at>gmail.com>

Author: 2 people

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

@@ -17,6 +17,9 @@
 
 package org.apache.spark.network;
 
+import io.netty.buffer.ByteBuf;
+import io.netty.handler.codec.MessageToMessageDecoder;
+
 import java.io.Closeable;
 import java.util.ArrayList;
 import java.util.List;
@@ -196,7 +199,7 @@ public TransportChannelHandler initializePipeline(
       pipeline
         .addLast("encoder", ENCODER)
         .addLast(TransportFrameDecoder.HANDLER_NAME, NettyUtils.createFrameDecoder())
-        .addLast("decoder", DECODER)
+        .addLast("decoder", getDecoder())
         .addLast("idleStateHandler",
           new IdleStateHandler(0, 0, conf.connectionTimeoutMs() / 1000))
         // NOTE: Chunks are currently guaranteed to be returned in the order of request, but this
@@ -216,6 +219,10 @@ public TransportChannelHandler initializePipeline(
     }
   }
 
+  protected MessageToMessageDecoder<ByteBuf> getDecoder() {
+    return DECODER;
+  }
+
   /**
    * Creates the server- and client-side handler which is used to handle both RequestMessages and
    * ResponseMessages. The channel is expected to have been successfully created, though certain

common/network-common/src/main/java/org/apache/spark/network/protocol/MessageWithHeader.java

@@ -35,7 +35,8 @@
  *
  * The header must be a ByteBuf, while the body can be a ByteBuf or a FileRegion.
  */
-class MessageWithHeader extends AbstractFileRegion {
+
+public class MessageWithHeader extends AbstractFileRegion {
 
   @Nullable private final ManagedBuffer managedBuffer;
   private final ByteBuf header;

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

@@ -184,6 +184,10 @@ public TransportResponseHandler getResponseHandler() {
     return responseHandler;
   }
 
+  public TransportRequestHandler getRequestHandler() {
+    return requestHandler;
+  }
+
   @Override
   public void channelRegistered(ChannelHandlerContext ctx) throws Exception {
     transportContext.getRegisteredConnections().inc();

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

@@ -20,7 +20,7 @@
 import java.util.Locale;
 import java.util.Properties;
 import java.util.concurrent.TimeUnit;
-
+import com.google.common.base.Preconditions;
 import com.google.common.primitives.Ints;
 import io.netty.util.NettyRuntime;
 
@@ -324,6 +324,35 @@ public boolean separateChunkFetchRequest() {
     return conf.getInt("spark.shuffle.server.chunkFetchHandlerThreadsPercent", 0) > 0;
   }
 
+  /**
+   * Percentage of io.serverThreads used by netty to process FinalizeShuffleMerge. When the config
+   * `spark.shuffle.server.finalizeShuffleMergeThreadsPercent` is set, shuffle server will use a
+   * separate EventLoopGroup to process FinalizeShuffleMerge messages, which are I/O intensive and
+   * could take long time to process due to disk contentions. The number of threads used for
+   * handling finalizeShuffleMerge requests are percentage of io.serverThreads (if defined) else it
+   * is a percentage of 2 * #cores.
+   */
+  public int finalizeShuffleMergeHandlerThreads() {
+    if (!this.getModuleName().equalsIgnoreCase("shuffle")) {
+      return 0;
+    }
+    Preconditions.checkArgument(separateFinalizeShuffleMerge(),
+        "Please set spark.shuffle.server.finalizeShuffleMergeThreadsPercent to a positive value");
+    int finalizeShuffleMergeThreadsPercent =
+        Integer.parseInt(conf.get("spark.shuffle.server.finalizeShuffleMergeThreadsPercent"));
+    int threads =
+        this.serverThreads() > 0 ? this.serverThreads() : 2 * NettyRuntime.availableProcessors();
+    return (int) Math.ceil(threads * (finalizeShuffleMergeThreadsPercent / 100.0));
+  }
+
+  /**
+   * Whether to use a separate EventLoopGroup to process FinalizeShuffleMerge messages, it is
+   * decided by the config `spark.shuffle.server.finalizeShuffleMergeThreadsPercent` is set or not.
+   */
+  public boolean separateFinalizeShuffleMerge() {
+    return conf.getInt("spark.shuffle.server.finalizeShuffleMergeThreadsPercent", 0) > 0;
+  }
+
   /**
    * 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-shuffle/src/main/java/org/apache/spark/network/shuffle/ShuffleTransportContext.java

@@ -0,0 +1,195 @@
+/*
+ * 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.shuffle;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.EventLoopGroup;
+import io.netty.channel.SimpleChannelInboundHandler;
+import io.netty.channel.socket.SocketChannel;
+import io.netty.handler.codec.MessageToMessageDecoder;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.spark.network.TransportContext;
+import org.apache.spark.network.protocol.Message;
+import org.apache.spark.network.protocol.MessageDecoder;
+import org.apache.spark.network.protocol.RpcRequest;
+import org.apache.spark.network.server.RpcHandler;
+import org.apache.spark.network.server.TransportChannelHandler;
+import org.apache.spark.network.server.TransportRequestHandler;
+import org.apache.spark.network.shuffle.protocol.BlockTransferMessage;
+import org.apache.spark.network.util.IOMode;
+import org.apache.spark.network.util.NettyUtils;
+import org.apache.spark.network.util.TransportConf;
+
+import static org.apache.spark.network.util.NettyUtils.getRemoteAddress;
+
+/**
+ * Extends {@link TransportContext} to support customized shuffle service. Specifically, we
+ * modified the Netty Channel Pipeline so that IO expensive messages such as FINALIZE_SHUFFLE_MERGE
+ * are processed in the separate handlers.
+ * */
+public class ShuffleTransportContext extends TransportContext {
+  private static final Logger logger = LoggerFactory.getLogger(ShuffleTransportContext.class);
+  private static final ShuffleMessageDecoder SHUFFLE_DECODER =
+      new ShuffleMessageDecoder(MessageDecoder.INSTANCE);
+  private final EventLoopGroup finalizeWorkers;
+
+  public ShuffleTransportContext(
+    TransportConf conf,
+    ExternalBlockHandler rpcHandler,
+    boolean closeIdleConnections) {
+    this(conf, rpcHandler, closeIdleConnections, false);
+  }
+
+  public ShuffleTransportContext(TransportConf conf,
+      RpcHandler rpcHandler,
+      boolean closeIdleConnections,
+      boolean isClientOnly) {
+    super(conf, rpcHandler, closeIdleConnections, isClientOnly);
+
+    if ("shuffle".equalsIgnoreCase(conf.getModuleName()) && conf.separateFinalizeShuffleMerge()) {
+      finalizeWorkers = NettyUtils.createEventLoop(
+          IOMode.valueOf(conf.ioMode()),
+          conf.finalizeShuffleMergeHandlerThreads(),
+          "shuffle-finalize-merge-handler");
+      logger.info("finalize shuffle merged workers created");
+    } else {
+      finalizeWorkers = null;
+    }
+  }
+
+  @Override
+  public TransportChannelHandler initializePipeline(SocketChannel channel) {
+    TransportChannelHandler ch = super.initializePipeline(channel);
+    addHandlerToPipeline(channel, ch);
+    return ch;
+  }
+
+  @Override
+  public TransportChannelHandler initializePipeline(SocketChannel channel,
+      RpcHandler channelRpcHandler) {
+    TransportChannelHandler ch = super.initializePipeline(channel, channelRpcHandler);
+    addHandlerToPipeline(channel, ch);
+    return ch;
+  }
+
+  /**
+   * Add finalize handler to pipeline if needed. This is needed only when
+   * separateFinalizeShuffleMerge is enabled.
+   */
+  private void addHandlerToPipeline(SocketChannel channel,
+      TransportChannelHandler transportChannelHandler) {
+    if (finalizeWorkers != null) {
+      channel.pipeline().addLast(finalizeWorkers, FinalizedHandler.HANDLER_NAME,
+        new FinalizedHandler(transportChannelHandler.getRequestHandler()));
+    }
+  }
+
+  @Override
+  protected MessageToMessageDecoder<ByteBuf> getDecoder() {
+    return finalizeWorkers == null ? super.getDecoder() : SHUFFLE_DECODER;
+  }
+
+  static class ShuffleMessageDecoder extends MessageToMessageDecoder<ByteBuf> {
+
+    private final MessageDecoder delegate;
+    ShuffleMessageDecoder(MessageDecoder delegate) {
+      super();
+      this.delegate = delegate;
+    }
+
+    /**
+     * Decode the message and check if it is a finalize merge request. If yes, then create an
+     * internal rpc request message and add it to the list of messages to be handled by
+     * {@link TransportChannelHandler}
+    */
+    @Override
+    protected void decode(ChannelHandlerContext channelHandlerContext,
+        ByteBuf byteBuf,
+        List<Object> list) throws Exception {
+      delegate.decode(channelHandlerContext, byteBuf, list);
+      Object msg = list.get(list.size() - 1);
+      if (msg instanceof RpcRequest) {
+        RpcRequest req = (RpcRequest) msg;
+        ByteBuffer buffer = req.body().nioByteBuffer();
+        byte type = Unpooled.wrappedBuffer(buffer).readByte();
+        if (type == BlockTransferMessage.Type.FINALIZE_SHUFFLE_MERGE.id()) {
+          list.remove(list.size() - 1);
+          RpcRequestInternal rpcRequestInternal =
+            new RpcRequestInternal(BlockTransferMessage.Type.FINALIZE_SHUFFLE_MERGE, req);
+          logger.trace("Created internal rpc request msg with rpcId {} for finalize merge req",
+            req.requestId);
+          list.add(rpcRequestInternal);
+        }
+      }
+    }
+  }
+
+  /**
+   * Internal message to handle rpc requests that is not accepted by
+   * {@link TransportChannelHandler} as this message doesn't extend {@link Message}. It will be
+   * accepted by {@link FinalizedHandler} instead, which is configured to execute in a separate
+   * EventLoopGroup.
+   */
+  static class RpcRequestInternal {
+    public final BlockTransferMessage.Type messageType;
+    public final RpcRequest rpcRequest;
+
+    RpcRequestInternal(BlockTransferMessage.Type messageType,
+        RpcRequest rpcRequest) {
+      this.messageType = messageType;
+      this.rpcRequest = rpcRequest;
+    }
+  }
+
+  static class FinalizedHandler extends SimpleChannelInboundHandler<RpcRequestInternal> {
+    private static final Logger logger = LoggerFactory.getLogger(FinalizedHandler.class);
+    public static final String HANDLER_NAME = "finalizeHandler";
+    private final TransportRequestHandler transportRequestHandler;
+
+    @Override
+    public boolean acceptInboundMessage(Object msg) throws Exception {
+      if (msg instanceof RpcRequestInternal) {
+        RpcRequestInternal rpcRequestInternal = (RpcRequestInternal) msg;
+        return rpcRequestInternal.messageType == BlockTransferMessage.Type.FINALIZE_SHUFFLE_MERGE;
+      }
+      return false;
+    }
+
+    FinalizedHandler(TransportRequestHandler transportRequestHandler) {
+      this.transportRequestHandler = transportRequestHandler;
+    }
+
+    @Override
+    protected void channelRead0(ChannelHandlerContext channelHandlerContext,
+        RpcRequestInternal req) throws Exception {
+      if (logger.isTraceEnabled()) {
+        logger.trace("Finalize shuffle req from {} for rpc request {}",
+                getRemoteAddress(channelHandlerContext.channel()), req.rpcRequest.requestId);
+      }
+      this.transportRequestHandler.handle(req.rpcRequest);
+    }
+  }
+}