From eb6456cd0c4ed98f627c13be49ffee651aa5dcc3 Mon Sep 17 00:00:00 2001
From: mingji <fengmingxiao.fmx@alibaba-inc.com>
Date: Fri, 18 Oct 2024 11:42:43 +0800
Subject: [PATCH] fix incompatible changes.

---
 .../deploy/master/clustermeta/AbstractMetaManager.java   | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
index d2520fbc8fa..031b64b8aac 100644
--- a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
@@ -29,6 +29,7 @@
 import java.util.concurrent.atomic.LongAdder;
 import java.util.stream.Collectors;
 
+import scala.Int;
 import scala.Option;
 import scala.Tuple2;
 
@@ -129,7 +130,13 @@ public void updateUnregisterShuffleMeta(String shuffleKey) {
   }
 
   public void updateBatchUnregisterShuffleMeta(List<String> shuffleKeys) {
-    registeredShuffle.removeAll(shuffleKeys);
+    for (String shuffleKey : shuffleKeys) {
+      Tuple2<String, Object> appIdShuffleId = Utils.splitShuffleKey(shuffleKey);
+      String appId = appIdShuffleId._1;
+      if (registeredAppAndShuffles.containsKey(appId)) {
+        registeredAppAndShuffles.get(appId).remove((Int) appIdShuffleId._2);
+      }
+    }
   }
 
   public void updateAppHeartbeatMeta(String appId, long time, long totalWritten, long fileCount) {