[#1476] feat(spark): Provide dedicated unregister app rpc interface

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

@@ -60,6 +60,7 @@
 import org.apache.uniffle.client.request.RssReportShuffleResultRequest;
 import org.apache.uniffle.client.request.RssSendCommitRequest;
 import org.apache.uniffle.client.request.RssSendShuffleDataRequest;
+import org.apache.uniffle.client.request.RssUnregisterAppRequest;
 import org.apache.uniffle.client.request.RssUnregisterShuffleRequest;
 import org.apache.uniffle.client.response.ClientResponse;
 import org.apache.uniffle.client.response.RssAppHeartBeatResponse;
@@ -73,6 +74,7 @@
 import org.apache.uniffle.client.response.RssReportShuffleResultResponse;
 import org.apache.uniffle.client.response.RssSendCommitResponse;
 import org.apache.uniffle.client.response.RssSendShuffleDataResponse;
+import org.apache.uniffle.client.response.RssUnregisterAppResponse;
 import org.apache.uniffle.client.response.RssUnregisterShuffleResponse;
 import org.apache.uniffle.client.response.SendShuffleDataResult;
 import org.apache.uniffle.client.util.ClientUtils;
@@ -984,14 +986,49 @@ public void unregisterShuffle(String appId, int shuffleId) {
 
   @Override
   public void unregisterShuffle(String appId) {
+    RssUnregisterAppRequest request = new RssUnregisterAppRequest(appId);
+
     if (appId == null) {
       return;
     }
     Map<Integer, Set<ShuffleServerInfo>> appServerMap = shuffleServerInfoMap.get(appId);
     if (appServerMap == null) {
       return;
     }
-    appServerMap.keySet().forEach(shuffleId -> unregisterShuffle(appId, shuffleId));
+    Set<ShuffleServerInfo> shuffleServerInfos = getAllShuffleServers(appId);
+
+    ExecutorService executorService = null;
+    try {
+      executorService =
+          ThreadUtils.getDaemonFixedThreadPool(
+              Math.min(unregisterThreadPoolSize, shuffleServerInfos.size()), "unregister-shuffle");
+
+      ThreadUtils.executeTasks(
+          executorService,
+          shuffleServerInfos,
+          shuffleServerInfo -> {
+            try {
+              ShuffleServerClient client =
+                  ShuffleServerClientFactory.getInstance()
+                      .getShuffleServerClient(clientType, shuffleServerInfo, rssConf);
+              RssUnregisterAppResponse response = client.unregisterApp(request);
+              if (response.getStatusCode() != StatusCode.SUCCESS) {
+                LOG.warn("Failed to unregister shuffle to " + shuffleServerInfo);
+              }
+            } catch (Exception e) {
+              LOG.warn("Error happened when unregistering to " + shuffleServerInfo, e);
+            }
+            return null;
+          },
+          unregisterRequestTimeSec,
+          "unregister shuffle server");
+
+    } finally {
+      if (executorService != null) {
+        executorService.shutdownNow();
+      }
+      shuffleServerInfoMap.remove(appId);
+    }
   }
 
   private void throwExceptionIfNecessary(ClientResponse response, String errorMsg) {

internal-client/src/main/java/org/apache/uniffle/client/api/ShuffleServerClient.java

@@ -28,6 +28,7 @@
 import org.apache.uniffle.client.request.RssReportShuffleResultRequest;
 import org.apache.uniffle.client.request.RssSendCommitRequest;
 import org.apache.uniffle.client.request.RssSendShuffleDataRequest;
+import org.apache.uniffle.client.request.RssUnregisterAppRequest;
 import org.apache.uniffle.client.request.RssUnregisterShuffleRequest;
 import org.apache.uniffle.client.response.RssAppHeartBeatResponse;
 import org.apache.uniffle.client.response.RssFinishShuffleResponse;
@@ -39,6 +40,7 @@
 import org.apache.uniffle.client.response.RssReportShuffleResultResponse;
 import org.apache.uniffle.client.response.RssSendCommitResponse;
 import org.apache.uniffle.client.response.RssSendShuffleDataResponse;
+import org.apache.uniffle.client.response.RssUnregisterAppResponse;
 import org.apache.uniffle.client.response.RssUnregisterShuffleResponse;
 
 public interface ShuffleServerClient {
@@ -47,6 +49,8 @@ public interface ShuffleServerClient {
 
   RssRegisterShuffleResponse registerShuffle(RssRegisterShuffleRequest request);
 
+  RssUnregisterAppResponse unregisterApp(RssUnregisterAppRequest request);
+
   RssSendShuffleDataResponse sendShuffleData(RssSendShuffleDataRequest request);
 
   RssSendCommitResponse sendCommit(RssSendCommitRequest request);

internal-client/src/main/java/org/apache/uniffle/client/impl/grpc/ShuffleServerGrpcClient.java

@@ -45,6 +45,7 @@
 import org.apache.uniffle.client.request.RssReportShuffleResultRequest;
 import org.apache.uniffle.client.request.RssSendCommitRequest;
 import org.apache.uniffle.client.request.RssSendShuffleDataRequest;
+import org.apache.uniffle.client.request.RssUnregisterAppRequest;
 import org.apache.uniffle.client.request.RssUnregisterShuffleRequest;
 import org.apache.uniffle.client.response.RssAppHeartBeatResponse;
 import org.apache.uniffle.client.response.RssFinishShuffleResponse;
@@ -56,6 +57,7 @@
 import org.apache.uniffle.client.response.RssReportShuffleResultResponse;
 import org.apache.uniffle.client.response.RssSendCommitResponse;
 import org.apache.uniffle.client.response.RssSendShuffleDataResponse;
+import org.apache.uniffle.client.response.RssUnregisterAppResponse;
 import org.apache.uniffle.client.response.RssUnregisterShuffleResponse;
 import org.apache.uniffle.common.BufferSegment;
 import org.apache.uniffle.common.PartitionRange;
@@ -328,6 +330,35 @@ public long requirePreAllocation(
     return result;
   }
 
+  private RssProtos.AppUnregisterResponse doUnregisterApp(String appId) {
+    RssProtos.AppUnregisterRequest request =
+        RssProtos.AppUnregisterRequest.newBuilder().setAppId(appId).build();
+    return blockingStub.unregisterApp(request);
+  }
+
+  @Override
+  public RssUnregisterAppResponse unregisterApp(RssUnregisterAppRequest request) {
+    RssProtos.AppUnregisterResponse rpcResponse = doUnregisterApp(request.getAppId());
+
+    RssUnregisterAppResponse response;
+    RssProtos.StatusCode statusCode = rpcResponse.getStatus();
+
+    switch (statusCode) {
+      case SUCCESS:
+        response = new RssUnregisterAppResponse(StatusCode.SUCCESS);
+        break;
+      default:
+        String msg =
+            String.format(
+                "Errors on unregister app to %s:%s for appId[%s], error: %s",
+                host, port, request.getAppId(), rpcResponse.getRetMsg());
+        LOG.error(msg);
+        throw new RssException(msg);
+    }
+
+    return response;
+  }
+
   private RssProtos.ShuffleUnregisterResponse doUnregisterShuffle(String appId, int shuffleId) {
     RssProtos.ShuffleUnregisterRequest request =
         RssProtos.ShuffleUnregisterRequest.newBuilder()

internal-client/src/main/java/org/apache/uniffle/client/request/RssUnregisterAppRequest.java

@@ -0,0 +1,30 @@
+/*
+ * 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.request;
+
+public class RssUnregisterAppRequest {
+  private String appId;
+
+  public RssUnregisterAppRequest(String appId) {
+    this.appId = appId;
+  }
+
+  public String getAppId() {
+    return appId;
+  }
+}

internal-client/src/main/java/org/apache/uniffle/client/response/RssUnregisterAppResponse.java

@@ -0,0 +1,27 @@
+/*
+ * 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.response;
+
+import org.apache.uniffle.common.rpc.StatusCode;
+
+public class RssUnregisterAppResponse extends ClientResponse {
+
+  public RssUnregisterAppResponse(StatusCode statusCode) {
+    super(statusCode);
+  }
+}

proto/src/main/proto/Rss.proto

@@ -26,6 +26,7 @@ package rss.common;
 service ShuffleServer {
   rpc registerShuffle (ShuffleRegisterRequest) returns (ShuffleRegisterResponse);
   rpc unregisterShuffle(ShuffleUnregisterRequest) returns (ShuffleUnregisterResponse);
+  rpc unregisterApp(AppUnregisterRequest) returns (AppUnregisterResponse);
   rpc sendShuffleData (SendShuffleDataRequest) returns (SendShuffleDataResponse);
   rpc getLocalShuffleIndex (GetLocalShuffleIndexRequest) returns (GetLocalShuffleIndexResponse);
   rpc getLocalShuffleData (GetLocalShuffleDataRequest) returns (GetLocalShuffleDataResponse);
@@ -197,6 +198,15 @@ message ShuffleRegisterResponse {
   string retMsg = 2;
 }
 
+message AppUnregisterRequest {
+  string appId = 1;
+}
+
+message AppUnregisterResponse {
+  StatusCode status = 1;
+  string retMsg = 2;
+}
+
 message SendShuffleDataRequest {
   string appId = 1;
   int32 shuffleId = 2;

server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java

@@ -18,13 +18,15 @@
 package org.apache.uniffle.server;
 
 import java.nio.ByteBuffer;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
+import com.google.common.collect.RangeMap;
 import com.google.common.collect.Sets;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.UnsafeByteOperations;
@@ -86,6 +88,7 @@
 import org.apache.uniffle.proto.RssProtos.ShuffleRegisterResponse;
 import org.apache.uniffle.proto.ShuffleServerGrpc.ShuffleServerImplBase;
 import org.apache.uniffle.server.buffer.PreAllocatedBufferInfo;
+import org.apache.uniffle.server.buffer.ShuffleBuffer;
 import org.apache.uniffle.storage.common.Storage;
 import org.apache.uniffle.storage.common.StorageReadMetrics;
 import org.apache.uniffle.storage.util.ShuffleStorageUtils;
@@ -99,6 +102,37 @@ public ShuffleServerGrpcService(ShuffleServer shuffleServer) {
     this.shuffleServer = shuffleServer;
   }
 
+  @Override
+  public void unregisterApp(
+      RssProtos.AppUnregisterRequest request,
+      StreamObserver<RssProtos.AppUnregisterResponse> responseStreamObserver) {
+    String appId = request.getAppId();
+
+    StatusCode result = StatusCode.SUCCESS;
+    String responseMessage = "OK";
+    try {
+      Map<Integer, RangeMap<Integer, ShuffleBuffer>> shuffleIdToBuffers =
+          shuffleServer.getShuffleBufferManager().getBufferPool().get(appId);
+      if (shuffleIdToBuffers != null) {
+        HashSet<Integer> shuffleIds = Sets.newHashSet(shuffleIdToBuffers.keySet());
+        shuffleIds.forEach(
+            shuffleId ->
+                shuffleServer.getShuffleTaskManager().removeShuffleDataAsync(appId, shuffleId));
+      }
+
+    } catch (Exception e) {
+      result = StatusCode.INTERNAL_ERROR;
+    }
+
+    RssProtos.AppUnregisterResponse reply =
+        RssProtos.AppUnregisterResponse.newBuilder()
+            .setStatus(result.toProto())
+            .setRetMsg(responseMessage)
+            .build();
+    responseStreamObserver.onNext(reply);
+    responseStreamObserver.onCompleted();
+  }
+
   @Override
   public void unregisterShuffle(
       RssProtos.ShuffleUnregisterRequest request,