(celeborn) branch main updated: [CELEBORN-1419] Avoid adding shuffle id repeatedly

[ni...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

nicholasjiang pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git


The following commit(s) were added to refs/heads/main by this push:
     new 059195963 [CELEBORN-1419] Avoid adding shuffle id repeatedly
059195963 is described below

commit 059195963e1fb567c2f2d295b5726148419e73cf
Author: huangxiaoping <17...@qq.com>
AuthorDate: Tue May 14 17:15:49 2024 +0800

    [CELEBORN-1419] Avoid adding shuffle id repeatedly
    
    ### What changes were proposed in this pull request?
    Avoid adding shuffle id repeatedly
    
    ### Why are the changes needed?
    `registerShuffle` has added the shuffle id to `sortShuffleIds`, and there is no need to add it repeatedly when calling `getWriter`
    
    ### Does this PR introduce _any_ user-facing change?
    None
    
    ### How was this patch tested?
    
    Closes #2503 from huangxiaopingRD/CELEBORN-1419.
    
    Authored-by: huangxiaoping <17...@qq.com>
    Signed-off-by: SteNicholas <pr...@163.com>
---
 .../java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java    | 3 +--
 .../java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java    | 3 +--
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
index 0344ca6d3..c1a45f5f1 100644
--- a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
+++ b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -137,7 +137,7 @@ public class SparkShuffleManager implements ShuffleManager {
 
   @Override
   public boolean unregisterShuffle(int appShuffleId) {
-    if (sortShuffleIds.contains(appShuffleId)) {
+    if (sortShuffleIds.remove(appShuffleId)) {
       return sortShuffleManager().unregisterShuffle(appShuffleId);
     }
     // For Spark driver side trigger unregister shuffle.
@@ -214,7 +214,6 @@ public class SparkShuffleManager implements ShuffleManager {
               "Unrecognized shuffle write mode!" + celebornConf.shuffleWriterMode());
         }
       } else {
-        sortShuffleIds.add(handle.shuffleId());
         return sortShuffleManager().getWriter(handle, mapId, context);
       }
     } catch (IOException e) {
diff --git a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
index e705e5d01..26fa6fb00 100644
--- a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
+++ b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -190,7 +190,7 @@ public class SparkShuffleManager implements ShuffleManager {
 
   @Override
   public boolean unregisterShuffle(int appShuffleId) {
-    if (sortShuffleIds.contains(appShuffleId)) {
+    if (sortShuffleIds.remove(appShuffleId)) {
       return sortShuffleManager().unregisterShuffle(appShuffleId);
     }
     // For Spark driver side trigger unregister shuffle.
@@ -293,7 +293,6 @@ public class SparkShuffleManager implements ShuffleManager {
               "Unrecognized shuffle write mode!" + celebornConf.shuffleWriterMode());
         }
       } else {
-        sortShuffleIds.add(handle.shuffleId());
         return sortShuffleManager().getWriter(handle, mapId, context, metrics);
       }
     } catch (IOException e) {