[CARMEL-6222]Improve congestion in shuffle service (#1068)

* [CARMEL-6222]Improve congestion in shuffle service

Author: wakun

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

@@ -333,4 +333,8 @@ public Meter getChunkFetchRequestHandleSuccessRate() {
   public Meter getChunkFetchRequestHandleFailRate() {
     return chunkFetchRequestHandleFailRate;
   }
+
+  public long getChunksBeingTransferred() {
+    return rpcHandler.getStreamManager().chunksBeingTransferred();
+  }
 }

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

@@ -166,7 +166,8 @@ private String findHeaviestShuffle() {
    *    HIGH_CONGESTION --> NO_CONGESTION
    */
   private void checkCongestion() {
-    int requestPendingTasks = transportContext.getPendingTasks();
+    int requestPendingTasks =
+        transportContext.getPendingTasks() + (int)streamManager.chunksBeingTransferred();
     CongestionMode prevMode = congestionMode;
     if (requestPendingTasks <= lowWaterMark) {
       congestionMode = CongestionMode.NO;

common/network-yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java

@@ -207,6 +207,9 @@ protected void serviceInit(Configuration conf) throws Exception {
       blockHandler.getAllMetrics().getMetrics().put("numNettyIOPendingTasks",
         (Gauge) () -> shuffleServer.getPendingTasks());
 
+      blockHandler.getAllMetrics().getMetrics().put("chunksBeingTransferred",
+          (Gauge) () -> transportContext.getChunksBeingTransferred());
+
       blockHandler.getAllMetrics().getMetrics().put("numNettyChunkFetchPendingTasks",
         (Gauge) () -> transportContext.getPendingTasks());