shuttle: Allow closing flow control requests

Author: bigbigxu

src/main/java/com/oppo/shuttle/rss/clients/NettyClient.java

@@ -55,6 +55,7 @@ public class NettyClient {
     private final long netWorkTimeout;
     private final int ioMaxRetry;
     private final long retryBaseWaitTime;
+    private final boolean flowControlEnable;
 
     private final ResponseCallback callback;
 
@@ -64,6 +65,7 @@ public NettyClient(List<Ors2ServerGroup> groupList, SparkConf conf, AppTaskInfo
         netWorkTimeout = clientFactory.getNetWorkTimeout();
         ioMaxRetry = Math.max((int) conf.get(Ors2Config.sendDataMaxRetries()), 1);
         retryBaseWaitTime = (long) conf.get(Ors2Config.retryBaseWaitTime());
+        flowControlEnable = (boolean) conf.get(Ors2Config.flowControlEnable());
         long networkSlowTime = (long) conf.get(Ors2Config.networkSlowTime());
         mapWriteDispersion = (boolean) conf.get(Ors2Config.mapWriteDispersion());
         int workerRetryNumber = (int) conf.get(Ors2Config.workerRetryNumber());
@@ -146,7 +148,7 @@ public void send(int workerId, ShufflePacket packet) {
             }
 
             Tuple2<ShuffleClient, ShuffleClient> tuple2 = getClient(workerId, 0, Optional.empty());
-            Request request = new Request(packet, workerId, tuple2._1, tuple2._2, callback);
+            Request request = new Request(flowControlEnable, packet, workerId, tuple2._1, tuple2._2, callback);
 
             readySend.incrementAndGet();
             request.writeBuild();
@@ -203,7 +205,8 @@ public Request createRetryRequest(Request request, boolean channelError) {
         }
 
         Tuple2<ShuffleClient, ShuffleClient> tuple2 = getClient(request.getWorkerId(), retry, errorServer);
-        return new Request(request.getPacket(), request.getWorkerId(), tuple2._1, tuple2._2, request.getCallback());
+        return new Request(flowControlEnable, request.getPacket(), request.getWorkerId(),
+                tuple2._1, tuple2._2, request.getCallback());
     }
 
     public Ors2ClientFactory getClientFactory() {

src/main/java/com/oppo/shuttle/rss/clients/handler/Request.java

@@ -17,11 +17,13 @@
 package com.oppo.shuttle.rss.clients.handler;
 
 import com.oppo.shuttle.rss.clients.ShuffleClient;
+import com.oppo.shuttle.rss.common.Constants;
 import com.oppo.shuttle.rss.common.Ors2WorkerDetail;
 import com.oppo.shuttle.rss.messages.ShufflePacket;
 import org.apache.spark.util.Utils;
 
 public class Request {
+    private final boolean flowControlEnable;
     private final ShufflePacket packet;
     private final ShuffleClient buildClient;
     private final ShuffleClient dataClient;
@@ -33,11 +35,13 @@ public class Request {
     private final int workerId;
 
     public Request(
+            boolean flowControlEnable,
             ShufflePacket packet,
             int workerId,
             ShuffleClient buildClient,
             ShuffleClient dataClient,
             ResponseCallback callback) {
+        this.flowControlEnable = flowControlEnable;
         this.packet = packet;
         this.buildClient = buildClient;
         this.dataClient = dataClient;
@@ -66,7 +70,11 @@ public void onError(Throwable e) {
     }
 
     public void writeBuild() {
-        buildClient.writeAndFlush(this, packet.buildBuffer());
+        if (flowControlEnable) {
+            buildClient.writeAndFlush(this, packet.buildBuffer());
+        } else {
+            writeData(Constants.SKIP_CHECK_BUILD_ID, Constants.SKIP_CHECK_BUILD_ID);
+        }
     }
 
     public void writeData(int id, long value) {

src/main/java/com/oppo/shuttle/rss/common/Constants.java

@@ -74,4 +74,5 @@ public class Constants {
 
   public static final String SHUTTLE_HOST_IP_ENV_NAME = "SHUTTLE_HOST_IP";
 
+  public static final int SKIP_CHECK_BUILD_ID = -100;
 }

src/main/java/com/oppo/shuttle/rss/handlers/ShuffleDataIncomingHandler.java

@@ -16,6 +16,7 @@
 
 package com.oppo.shuttle.rss.handlers;
 
+import com.oppo.shuttle.rss.common.Constants;
 import com.oppo.shuttle.rss.common.StageShuffleId;
 import com.oppo.shuttle.rss.exceptions.Ors2InvalidDataException;
 import com.oppo.shuttle.rss.execution.BuildConnectionExecutor;
@@ -58,11 +59,13 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) {
         ShuffleData shuffleData = (ShuffleData) msg;
         ShuffleMessage.UploadPackageRequest up = shuffleData.getUploadPackage();
 
-        if (!buildConnExecutor.checkConnectionIdValue(up.getBuildId(), up.getBuildValue())) {
-            logger.warn("invalid or timeout connectionId&Value: {}, {}, for {}, address {}",
-                    up.getBuildId(), up.getBuildValue(), up.getAppId(), shuffleConnInfo);
-        } else {
-            buildConnExecutor.releaseConnection(up.getBuildId());
+        if (up.getBuildId() != Constants.SKIP_CHECK_BUILD_ID) {
+            if(!buildConnExecutor.checkConnectionIdValue(up.getBuildId(), up.getBuildValue())) {
+                logger.warn("invalid or timeout connectionId&Value: {}, {}, for {}, address {}",
+                        up.getBuildId(), up.getBuildValue(), up.getAppId(), shuffleConnInfo);
+            } else {
+                buildConnExecutor.releaseConnection(up.getBuildId());
+            }
         }
 
         if (appId == null) {

src/main/scala/org/apache/spark/shuffle/Ors2Config.scala

@@ -327,6 +327,12 @@ object Ors2Config {
       .booleanConf
       .createWithDefault(true)
 
+  val flowControlEnable: ConfigEntry[Boolean] =
+    ConfigBuilder("spark.shuffle.rss.flowControlEnable")
+      .doc("Whether to enable flow control. Defaults to true")
+      .booleanConf
+      .createWithDefault(true)
+
   val SHUFFLE_TYPE_ASYNC = "async"
   val SHUFFLE_TYPE_SYNC = "sync"
 

src/test/scala/org/apache/spark/shuffle/Ors2ShuffleManagerTest.scala

@@ -161,4 +161,12 @@ class Ors2ShuffleManagerTest {
     conf.set(Ors2Config.serviceRegistryZKServers, zk)
     runWithConf(conf)
   }
+
+  @Test
+  def skipFlowControl(): Unit = {
+    val conf = new SparkConf()
+      .set(Ors2Config.shuffleWriterType, Ors2Config.SHUFFLE_WRITER_UNSAFE)
+      .set(Ors2Config.flowControlEnable, false)
+    runWithConf(conf)
+  }
 }