[SPARK-38005][core] Support cleaning up merged shuffle files and state from external shuffle service

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

@@ -0,0 +1,33 @@
+/*
+ * 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.util;
+
+public class PushBasedShuffleUtils {
+
+  /**
+   * Push based shuffle requires a comparable Id to distinguish the shuffle data among multiple
+   * application attempts. This variable is derived from the String typed appAttemptId.
+   * If no appAttemptId is set, the default comparableAppAttemptId is -1.
+   */
+  public static final int DEFAUT_APP_ATTEMPT_ID = -1;
+
+  /**
+   * The flag for deleting the current merged shuffle data.
+   */
+  public static final int DELETE_CURRENT_MERGED_SHUFFLE_ID = -1;
+}

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockHandler.java

@@ -198,6 +198,10 @@ protected void handleMessage(
       int numRemovedBlocks = blockManager.removeBlocks(msg.appId, msg.execId, msg.blockIds);
       callback.onSuccess(new BlocksRemoved(numRemovedBlocks).toByteBuffer());
 
+    } else if (msgObj instanceof RemoveShuffleMerge) {
+      RemoveShuffleMerge msg = (RemoveShuffleMerge) msgObj;
+      checkAuth(client, msg.appId);
+      mergeManager.removeShuffleMerge(msg);
     } else if (msgObj instanceof GetLocalDirsForExecutors) {
       GetLocalDirsForExecutors msg = (GetLocalDirsForExecutors) msgObj;
       checkAuth(client, msg.appId);

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockStoreClient.java

@@ -28,6 +28,8 @@
 
 import com.codahale.metrics.MetricSet;
 import com.google.common.collect.Lists;
+import static org.apache.spark.network.util.PushBasedShuffleUtils.DEFAUT_APP_ATTEMPT_ID;
+import static org.apache.spark.network.util.PushBasedShuffleUtils.DELETE_CURRENT_MERGED_SHUFFLE_ID;
 
 import org.apache.spark.network.TransportContext;
 import org.apache.spark.network.buffer.ManagedBuffer;
@@ -38,7 +40,14 @@
 import org.apache.spark.network.crypto.AuthClientBootstrap;
 import org.apache.spark.network.sasl.SecretKeyHolder;
 import org.apache.spark.network.server.NoOpRpcHandler;
-import org.apache.spark.network.shuffle.protocol.*;
+import org.apache.spark.network.shuffle.protocol.BlockTransferMessage;
+import org.apache.spark.network.shuffle.protocol.BlocksRemoved;
+import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo;
+import org.apache.spark.network.shuffle.protocol.FinalizeShuffleMerge;
+import org.apache.spark.network.shuffle.protocol.MergeStatuses;
+import org.apache.spark.network.shuffle.protocol.RegisterExecutor;
+import org.apache.spark.network.shuffle.protocol.RemoveBlocks;
+import org.apache.spark.network.shuffle.protocol.RemoveShuffleMerge;
 import org.apache.spark.network.util.TransportConf;
 
 /**
@@ -52,10 +61,8 @@ public class ExternalBlockStoreClient extends BlockStoreClient {
   private final boolean authEnabled;
   private final SecretKeyHolder secretKeyHolder;
   private final long registrationTimeoutMs;
-  // Push based shuffle requires a comparable Id to distinguish the shuffle data among multiple
-  // application attempts. This variable is derived from the String typed appAttemptId. If no
-  // appAttemptId is set, the default comparableAppAttemptId is -1.
-  private int comparableAppAttemptId = -1;
+
+  private int comparableAppAttemptId = DEFAUT_APP_ATTEMPT_ID;
 
   /**
    * Creates an external shuffle client, with SASL optionally enabled. If SASL is not enabled,
@@ -315,6 +322,23 @@ public void onFailure(Throwable e) {
     return numRemovedBlocksFuture;
   }
 
+  public Boolean removeShuffleMerge(String host, int port, int shuffleId) {
+    try {
+      checkInit();
+      TransportClient client = clientFactory.createClient(host, port);
+      String appAttemptId =
+          getAppAttemptId() == null ? String.valueOf(DEFAUT_APP_ATTEMPT_ID) : getAppAttemptId();
+      ByteBuffer removeMessage = new RemoveShuffleMerge(appId, appAttemptId, shuffleId,
+          DELETE_CURRENT_MERGED_SHUFFLE_ID).toByteBuffer();
+      client.send(removeMessage);
+    } catch (Throwable e) {
+      logger.error("Error while sending RemoveShuffleMerge request to {}:{}",
+          host, port, e);
+      return false;
+    }
+    return true;
+  }
+
   @Override
   public void close() {
     checkInit();

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/MergedShuffleFileManager.java

@@ -26,6 +26,7 @@
 import org.apache.spark.network.shuffle.protocol.FinalizeShuffleMerge;
 import org.apache.spark.network.shuffle.protocol.MergeStatuses;
 import org.apache.spark.network.shuffle.protocol.PushBlockStream;
+import org.apache.spark.network.shuffle.protocol.RemoveShuffleMerge;
 
 /**
  * The MergedShuffleFileManager is used to process push based shuffle when enabled. It works
@@ -121,6 +122,14 @@ MergedBlockMeta getMergedBlockMeta(
    */
   String[] getMergedBlockDirs(String appId);
 
+  /**
+   * Handles the request to remove shuffle merge files.
+   *
+   * @param msg contains shuffle details (appId, shuffleId, etc) to uniquely identify
+   *            a shuffle to be removed
+   */
+  void removeShuffleMerge(RemoveShuffleMerge msg);
+
   /**
    * Optionally close any resources associated the MergedShuffleFileManager, such as the
    * leveldb for state persistence.

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/NoOpMergedShuffleFileManager.java

@@ -26,6 +26,7 @@
 import org.apache.spark.network.shuffle.protocol.FinalizeShuffleMerge;
 import org.apache.spark.network.shuffle.protocol.MergeStatuses;
 import org.apache.spark.network.shuffle.protocol.PushBlockStream;
+import org.apache.spark.network.shuffle.protocol.RemoveShuffleMerge;
 import org.apache.spark.network.util.TransportConf;
 
 /**
@@ -80,6 +81,11 @@ public MergedBlockMeta getMergedBlockMeta(
     throw new UnsupportedOperationException("Cannot handle shuffle block merge");
   }
 
+  @Override
+  public void removeShuffleMerge(RemoveShuffleMerge msg) {
+    throw new UnsupportedOperationException("Cannot handle merged shuffle remove");
+  }
+
   @Override
   public String[] getMergedBlockDirs(String appId) {
     throw new UnsupportedOperationException("Cannot handle shuffle block merge");