[#2111] improvement(client): cleanup the useless shuffle server clients when unregister shuffle

client/src/main/java/org/apache/uniffle/client/impl/ShuffleWriteClientImpl.java

@@ -1078,6 +1078,8 @@ public void unregisterShuffle(String appId) {
         executorService.shutdownNow();
       }
       shuffleServerInfoMap.remove(appId);
+      ShuffleServerClientFactory.getInstance()
+          .cleanUselessShuffleServerClients(getAllShuffleServers());
     }
   }
 
@@ -1148,6 +1150,14 @@ Set<ShuffleServerInfo> getAllShuffleServers(String appId) {
     return serverInfos;
   }
 
+  Set<ShuffleServerInfo> getAllShuffleServers() {
+    Set<ShuffleServerInfo> serverInfos = Sets.newHashSet();
+    shuffleServerInfoMap
+        .values()
+        .forEach(appServerMap -> appServerMap.values().forEach(serverInfos::addAll));
+    return serverInfos;
+  }
+
   @VisibleForTesting
   public ShuffleServerClient getShuffleServerClient(ShuffleServerInfo shuffleServerInfo) {
     return ShuffleServerClientFactory.getInstance()
@@ -1177,7 +1187,11 @@ void addShuffleServer(String appId, int shuffleId, ShuffleServerInfo serverInfo)
   void removeShuffleServer(String appId, int shuffleId) {
     Map<Integer, Set<ShuffleServerInfo>> appServerMap = shuffleServerInfoMap.get(appId);
     if (appServerMap != null) {
-      appServerMap.remove(shuffleId);
+      Set<ShuffleServerInfo> removed = appServerMap.remove(shuffleId);
+      if (removed != null) {
+        ShuffleServerClientFactory.getInstance()
+            .cleanUselessShuffleServerClients(getAllShuffleServers());
+      }
     }
   }
 }

internal-client/src/main/java/org/apache/uniffle/client/factory/ShuffleServerClientFactory.java

@@ -18,6 +18,10 @@
 package org.apache.uniffle.client.factory;
 
 import java.util.Map;
+import java.util.Set;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import org.apache.uniffle.client.api.ShuffleServerClient;
 import org.apache.uniffle.client.impl.grpc.ShuffleServerGrpcClient;
@@ -28,6 +32,7 @@
 import org.apache.uniffle.common.util.JavaUtils;
 
 public class ShuffleServerClientFactory {
+  private static final Logger LOG = LoggerFactory.getLogger(ShuffleServerClientFactory.class);
 
   private Map<String, Map<ShuffleServerInfo, ShuffleServerClient>> clients;
 
@@ -75,9 +80,30 @@ public synchronized ShuffleServerClient getShuffleServerClient(
     return serverToClients.get(shuffleServerInfo);
   }
 
+  public synchronized void cleanUselessShuffleServerClients(
+      Set<ShuffleServerInfo> usefulShuffleServerInfos) {
+    clients
+        .values()
+        .forEach(
+            (serverToClients -> {
+              serverToClients.forEach(
+                  (shuffleServerInfo, client) -> {
+                    if (!usefulShuffleServerInfos.contains(shuffleServerInfo)) {
+                      LOG.info("Clean useless shuffle server client for " + shuffleServerInfo);
+                      client.close();
+                      serverToClients.remove(shuffleServerInfo);
+                    }
+                  });
+            }));
+  }
+
   // Only for tests
   public synchronized void cleanupCache() {
     clients.values().stream().flatMap(x -> x.values().stream()).forEach(ShuffleServerClient::close);
     this.clients = JavaUtils.newConcurrentMap();
   }
+
+  public Map<String, Map<ShuffleServerInfo, ShuffleServerClient>> getClients() {
+    return clients;
+  }
 }

internal-client/src/test/java/org/apache/uniffle/client/factory/ShuffleServerClientTest.java

@@ -0,0 +1,55 @@
+/*
+ * 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.uniffle.client.factory;
+
+import java.util.HashSet;
+import java.util.Map;
+
+import com.google.common.collect.Sets;
+import org.apache.commons.collections4.CollectionUtils;
+import org.junit.jupiter.api.Test;
+
+import org.apache.uniffle.client.api.ShuffleServerClient;
+import org.apache.uniffle.common.ShuffleServerInfo;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class ShuffleServerClientTest {
+
+  @Test
+  public void testCleanUselessClient() {
+    ShuffleServerInfo shuffleServerInfo1 = new ShuffleServerInfo("127.0.0.1", 19999);
+    ShuffleServerInfo shuffleServerInfo2 = new ShuffleServerInfo("127.0.0.1", 29999);
+    ShuffleServerInfo shuffleServerInfo3 = new ShuffleServerInfo("127.0.0.1", 39999);
+    String clientType = "grpc_netty";
+    ShuffleServerClientFactory shuffleServerClientFactory =
+        ShuffleServerClientFactory.getInstance();
+    shuffleServerClientFactory.getShuffleServerClient(clientType, shuffleServerInfo1);
+    shuffleServerClientFactory.getShuffleServerClient(clientType, shuffleServerInfo2);
+    shuffleServerClientFactory.getShuffleServerClient(clientType, shuffleServerInfo3);
+    Map<ShuffleServerInfo, ShuffleServerClient> clientMap =
+        shuffleServerClientFactory.getClients().get(clientType);
+    assertEquals(3, clientMap.size());
+    HashSet<ShuffleServerInfo> remainShuffleServers =
+        Sets.newHashSet(shuffleServerInfo1, shuffleServerInfo2);
+    shuffleServerClientFactory.cleanUselessShuffleServerClients(remainShuffleServers);
+    assertEquals(2, clientMap.size());
+    assertTrue(CollectionUtils.isEqualCollection(clientMap.keySet(), remainShuffleServers));
+  }
+}